pyspark - 应用窗口函数计算 pySpark 中的差异

Question

我正在使用 pySpark ，并设置了我的数据框，其中包含代表每日 Assets 价格的两列，如下所示:

ind = sc.parallelize(range(1,5))
prices = sc.parallelize([33.3,31.1,51.2,21.3])
data = ind.zip(prices)
df = sqlCtx.createDataFrame(data,["day","price"])

我在申请时得到 df.show() :

+---+-----+
|day|price|
+---+-----+
|  1| 33.3|
|  2| 31.1|
|  3| 51.2|
|  4| 21.3|
+---+-----+

这很好。我想要另一列包含价格列的日常返回，即类似
(price(day2)-price(day1))/(price(day1))
经过大量研究，我得知这是通过应用 pyspark.sql.window 最有效地实现的。功能，但我无法看到如何。

Answer 1

您可以使用 lag 带上前一天的列函数，并添加额外的列来从两列进行实际的日常返回，但您可能必须告诉 spark 如何对数据进行分区和/或命令它执行滞后，如下所示:

from pyspark.sql.window import Window
import pyspark.sql.functions as func
from pyspark.sql.functions import lit

dfu = df.withColumn('user', lit('tmoore'))

df_lag = dfu.withColumn('prev_day_price',
                        func.lag(dfu['price'])
                                 .over(Window.partitionBy("user")))

result = df_lag.withColumn('daily_return', 
          (df_lag['price'] - df_lag['prev_day_price']) / df_lag['price'] )

>>> result.show()
+---+-----+-------+--------------+--------------------+
|day|price|   user|prev_day_price|        daily_return|
+---+-----+-------+--------------+--------------------+
|  1| 33.3| tmoore|          null|                null|
|  2| 31.1| tmoore|          33.3|-0.07073954983922816|
|  3| 51.2| tmoore|          31.1|         0.392578125|
|  4| 21.3| tmoore|          51.2|  -1.403755868544601|
+---+-----+-------+--------------+--------------------+

这里是对 Window functions in Spark 的详细介绍.