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

时间:2019-02-15 17:31:44

标签: python pandas apache-spark pyspark user-defined-functions

我要对数据框执行以下操作:

  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]                     |
+-----+------------------------------------------+---------------------------------+

Question here未提供我问题的答案。我正在对分组数据应用开窗操作。

1 个答案:

答案 0 :(得分:0)

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

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

此处新创建的window列将是一列字典,其中包含键startend,分别表示窗口的开始时间和结束时间。您可以通过访问各个元素将其进一步平整为两列开始时间和结束时间:

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

然后,在应用于熊猫UDF之前,Spark数据帧的状态将是UDF接收到的熊猫数据帧的状态。这样,您将收到UDF侧的窗口,并能够在转换后返回它们的值。