apache-spark - Spark 窗口函数 - rangeBetween 日期

Question

我有一个带有日期列的 Spark SQL DataFrame，我想要获取的是给定日期范围内当前行之前的所有行。例如，我想要获取给定行之前 7 天前的所有行。我发现，我需要使用窗口函数，例如:

Window \
    .partitionBy('id') \
    .orderBy('start')

我想要 rangeBetween 7 天，但 Spark 文档中我找不到任何相关内容。 Spark 甚至提供这样的选项吗？现在我只是获取前面的所有行:

.rowsBetween(-sys.maxsize, 0)

但想要实现类似的目标:

.rangeBetween("7 days", 0)

Answer 1

Spark >= 2.3

自 Spark 2.3 起，可以通过 SQL API 使用间隔对象，但 DataFrame API 支持 still work in progress .

df.createOrReplaceTempView("df")

spark.sql(
    """SELECT *, mean(some_value) OVER (
        PARTITION BY id 
        ORDER BY CAST(start AS timestamp) 
        RANGE BETWEEN INTERVAL 7 DAYS PRECEDING AND CURRENT ROW
     ) AS mean FROM df""").show()

## +---+----------+----------+------------------+       
## | id|     start|some_value|              mean|
## +---+----------+----------+------------------+
## |  1|2015-01-01|      20.0|              20.0|
## |  1|2015-01-06|      10.0|              15.0|
## |  1|2015-01-07|      25.0|18.333333333333332|
## |  1|2015-01-12|      30.0|21.666666666666668|
## |  2|2015-01-01|       5.0|               5.0|
## |  2|2015-01-03|      30.0|              17.5|
## |  2|2015-02-01|      20.0|              20.0|
## +---+----------+----------+------------------+

Spark <2.3

据我所知，无论是在 Spark 还是 Hive 中都不可能直接实现。两者都要求与 RANGE 一起使用的 ORDER BY 子句为数字。我发现的最接近的是转换为时间戳并以秒为单位进行操作。假设 start 列包含 date 类型:

from pyspark.sql import Row

row = Row("id", "start", "some_value")
df = sc.parallelize([
    row(1, "2015-01-01", 20.0),
    row(1, "2015-01-06", 10.0),
    row(1, "2015-01-07", 25.0),
    row(1, "2015-01-12", 30.0),
    row(2, "2015-01-01", 5.0),
    row(2, "2015-01-03", 30.0),
    row(2, "2015-02-01", 20.0)
]).toDF().withColumn("start", col("start").cast("date"))

一个小助手和窗口定义:

from pyspark.sql.window import Window
from pyspark.sql.functions import mean, col


# Hive timestamp is interpreted as UNIX timestamp in seconds*
days = lambda i: i * 86400 

最后查询:

w = (Window()
   .partitionBy(col("id"))
   .orderBy(col("start").cast("timestamp").cast("long"))
   .rangeBetween(-days(7), 0))

df.select(col("*"), mean("some_value").over(w).alias("mean")).show()

## +---+----------+----------+------------------+
## | id|     start|some_value|              mean|
## +---+----------+----------+------------------+
## |  1|2015-01-01|      20.0|              20.0|
## |  1|2015-01-06|      10.0|              15.0|
## |  1|2015-01-07|      25.0|18.333333333333332|
## |  1|2015-01-12|      30.0|21.666666666666668|
## |  2|2015-01-01|       5.0|               5.0|
## |  2|2015-01-03|      30.0|              17.5|
## |  2|2015-02-01|      20.0|              20.0|
## +---+----------+----------+------------------+

远谈不上漂亮，但很有效。

<小时/>

* Hive Language Manual, Types