如何使用 PySpark 的 Window function 到 model 指數衰減？

Question

我正在嘗試應用 PySpark Window function 來執行“指數衰減”。 公式是

todays_score = yesterdays_score * (weight) + todays_raw_score

例如，假設我們有一個 dataframe 以天為單位訂購，並且每天得分為 1：

+---+----+---------+
|day|user|raw_score|
+---+----+---------+
|  0|   a|        1|
|  1|   a|        1|
|  2|   a|        1|
|  3|   a|        1|
+---+----+---------+

如果我要計算 todays_score，它看起來像這樣：

+---+----+---------+------------+
|day|user|raw_score|todays_score| # Here's the math:
+---+----+---------+------------+
|  0|   a|        1|         1.0| (0 * .90) + 1
|  1|   a|        1|         1.9| (1.0 * .90) + 1
|  2|   a|        1|        2.71| (1.9 * .90) + 1
|  3|   a|        1|       3.439| (2.71 * .90) + 1
+---+----+---------+------------+

我試過使用 window 函數； 但是根據我所見，他們只能使用原始 dataframe 中的“靜態值”，而不是我們剛剛計算的值。 我什至嘗試創建一個“虛擬列”來啟動該過程； 但是那也不起作用。

我嘗試的代碼：

df = sqlContext.createDataFrame([
                                 (0, 'a', 1),
                                 (1, 'a', 1),
                                 (2, 'a', 1),
                                 (3, 'a', 1)],
    ['day', 'user', 'raw_score']
)
df.show()

# Create a "dummy column" (weighted score) so we can use it.
df2 = df.select('*', col('raw_score').alias('todays_score'))
df2.show()

w = Window.partitionBy('user') 

df2.withColumn('todays_score', 
              F.lag(F.col('todays_score'), count=1, default=0).over(w.orderBy('day'))* 0.9 + F.col('raw_score')) \
  .show()

這個（不需要的）output 是：

+---+----+---------+------------+
|day|user|raw_score|todays_score|
+---+----+---------+------------+
|  0|   a|        1|         1.0|
|  1|   a|        1|         1.9|
|  2|   a|        1|         1.9|
|  3|   a|        1|         1.9|
+---+----+---------+------------+

它只取前一個值 * (.90)，而不是剛剛計算的值。

如何訪問剛剛由 window function 計算的值？

Answer 1

對於Spark2.4+ ，您可以像這樣使用高階函數transform 、 aggregate 、 filter和arrays_zip 。 它適用於 raw_score 的任何組合，並且比 pandas_udaf 更快。 （假設數據已按每個用戶按天排序，如示例所示）

df.show() #sample dataframe
#+---+----+---------+
#|day|user|raw_score|
#+---+----+---------+
#|  0|   a|        1|
#|  1|   a|        1|
#|  2|   a|        1|
#|  3|   a|        1|
#+---+----+---------+


from pyspark.sql import functions as F

df\
  .groupBy("user").agg(F.collect_list("raw_score").alias("raw_score"),F.collect_list("day").alias("day"))\
   .withColumn("raw_score1", F.expr("""transform(raw_score,(x,i)-> struct(x as raw,i as index))"""))\
   .withColumn("todays_score", F.expr("""transform(raw_score1, x-> aggregate(filter(raw_score1,z-> z.index<=x.index)\
                                             ,cast(0 as double),(acc,y)->(acc*0.9)+y.raw))"""))\
   .withColumn("zip", F.explode(F.arrays_zip("day","raw_score","todays_score")))\
   .select("user", "zip.*")\
   .show(truncate=False)


#+----+---+---------+------------+
#|user|day|raw_score|todays_score|
#+----+---+---------+------------+
#|a   |0  |1        |1.0         |
#|a   |1  |1        |1.9         |
#|a   |2  |1        |2.71        |
#|a   |3  |1        |3.439       |
#+----+---+---------+------------+

UPDATE:

假設數據按天排序，如 sample 所示，您可以像這樣使用Pandas Grouped Map UDAF ：

import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql.functions import pandas_udf, PandasUDFType


@pandas_udf(df.withColumn("raw_score", F.lit(1.2456)).schema, PandasUDFType.GROUPED_MAP)
def grouped_map(df):
     for i in range(1,len(df)):
          df.loc[i,'raw_score']=(df.loc[i-1,'raw_score'] * 0.9)+1   

     return df
df\
  .groupby("user").apply(grouped_map).show()

#+---+----+---------+
#|day|user|raw_score|
#+---+----+---------+
#|  0|   a|      1.0|
#|  1|   a|      1.9|
#|  2|   a|     2.71|
#|  3|   a|    3.439|
#+---+----+---------+