如何将所有窗口值传递给 pyspark UDF

Question

我想对数据帧执行以下操作：

1. 分组列
2. 窗口数据
3. 对窗口数据执行（udf）自定义操作

这是我尝试过的示例代码：

from pyspark.sql import SparkSession
from pyspark.sql.types import *
ss = SparkSession.builder
from pyspark.sql import functions as F
from pyspark.sql.functions import udf, col

sparkSession = ss.getOrCreate()

sc = sparkSession.sparkContext
sc.setLogLevel("FATAL")

df = sparkSession.createDataFrame([(17.00, "2018-03-10"),
                                   (13.00, "2018-03-11"),
                                   (25.00, "2018-03-12"),
                                   (20.00, "2018-03-13"),
                                   (17.00, "2018-03-14"),
                                   (99.00, "2018-03-15"),
                                   (156.00, "2018-03-22"),
                                   (17.00, "2018-03-31"),
                                   (25.00, "2018-03-15"),
                                   (25.00, "2018-03-16")
                                   ],
                                  ["id", "ts"])

w = F.window(col("ts").cast("timestamp"), "10 days")
windo = w.alias("window")

@udf(ArrayType(FloatType()))
def new_tuple(x):
    #print(type(x))
    return x

df.groupBy("id", windo).agg(new_tuple(F.collect_list("id"))).show(truncate=False)

上面的代码给了我我想要的。 但是，我不确定“collect_list”方法。

我也试过熊猫 UDF。 我使用熊猫获得了预期的输出（见下文）。 但是，“应用”方法不返回窗口列。

问题：

1. collect_list 是在工作节点还是驱动节点上运行？ 如果 collect_list 将所有结果收集到主节点，则此代码可能不可扩展。

2. 有没有什么有效的方法可以在没有 collect_list 的情况下获得以下输出？

3. 我读过熊猫 UDF 是高效的。 但是，我不知道如何传递/返回窗口列。

预期输出：

+-----+------------------------------------------+---------------------------------+
|id   |window                                    |new_tuple(collect_list(id, 0, 0))|
+-----+------------------------------------------+---------------------------------+
|17.0 |[2018-03-29 19:00:00, 2018-04-08 19:00:00]|[17.0]                           |
|25.0 |[2018-03-09 18:00:00, 2018-03-19 19:00:00]|[25.0, 25.0, 25.0]               |
|13.0 |[2018-03-09 18:00:00, 2018-03-19 19:00:00]|[13.0]                           |
|99.0 |[2018-03-09 18:00:00, 2018-03-19 19:00:00]|[99.0]                           |
|156.0|[2018-03-19 19:00:00, 2018-03-29 19:00:00]|[156.0]                          |
|20.0 |[2018-03-09 18:00:00, 2018-03-19 19:00:00]|[20.0]                           |
|17.0 |[2018-03-09 18:00:00, 2018-03-19 19:00:00]|[17.0, 17.0]                     |
+-----+------------------------------------------+---------------------------------+

此处的问题不提供我的问题的答案。 我正在对分组数据应用窗口操作。

Answer 1

要回答您的第三个问题，您只需要显式创建一个用于存储窗口的列，例如：

df = df.withColumn('window', F.window(col("ts").cast("timestamp"), "10 days"))
df.groupby('id', 'window').apply(pandas_udf)

这里新创建的window列将是一列字典，键为start ， end表示窗口的开始时间和结束时间。 您可以通过访问各个元素将其进一步展平为两列开始和结束时间：

df = df.withColumn('start', F.col('window')['start'])
df = df.withColumn('end', F.col('window')['end'])

然后，无论 Spark 数据帧在应用于 Pandas UDF 之前的状态是什么，都将是 UDF 接收到的 Pandas 数据帧的状态。 因此，您将在 UDF 端收到窗口，并能够在转换后返回它们的值。