如何將所有窗口值傳遞給 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 端收到窗口，並能夠在轉換后返回它們的值。