PySpark 2.4按id分组计算滚动窗口内各值出现占比

实现思路 #

针对PySpark 2.4版本的特性，使用内置窗口函数即可实现需求，核心逻辑如下：

• 按id分区、time升序排序，定义行级滑动窗口，窗口范围为当前行及前序2行，固定窗口大小为3
• 利用avg+when组合计算各value的占比：对匹配目标value的记录标记1，其余标记0，窗口内求平均值即为该value的出现占比，省去额外计数、除法运算步骤
• 过滤窗口长度不足3的前序记录，保留符合要求的窗口结果

完整实现代码 #

from pyspark.sql import Window
import pyspark.sql.functions as F

# 构造示例测试数据
import pandas as pd
foo = pd.DataFrame({'id': [1,1,1,1,1, 2,2,2,2,2],
                    'time': [1,2,3,4,5, 1,2,3,4,5],
                    'value': ['a','a','a','b','b', 'b','b','c','c','c']})
foo_df = spark.createDataFrame(foo)

# 定义大小为3的滚动窗口
roll_window = Window.partitionBy("id").orderBy("time").rowsBetween(-2, 0)

# 计算各value占比
result_df = foo_df.withColumn("window_cnt", F.count("*").over(roll_window)) \
    .withColumn("perc_a", F.avg(F.when(F.col("value") == "a", 1).otherwise(0)).over(roll_window)) \
    .withColumn("perc_b", F.avg(F.when(F.col("value") == "b", 1).otherwise(0)).over(roll_window)) \
    .withColumn("perc_c", F.avg(F.when(F.col("value") == "c", 1).otherwise(0)).over(roll_window)) \
    .filter(F.col("window_cnt") == 3) \
    .select("id", "perc_a", "perc_b", "perc_c")

# 输出查看结果
result_df.show()

运行结果 #

执行代码后输出完全匹配预期：

+---+------------------+------------------+------------------+
| id|            perc_a|            perc_b|            perc_c|
+---+------------------+------------------+------------------+
|  1|               1.0|               0.0|               0.0|
|  1|0.6666666666666666|0.3333333333333333|               0.0|
|  1|0.3333333333333333|0.6666666666666666|               0.0|
|  2|               0.0|0.6666666666666666|0.3333333333333333|
|  2|               0.0|0.3333333333333333|0.6666666666666666|
|  2|               0.0|               0.0|               1.0|
+---+------------------+------------------+------------------+

注意事项 #

• 选择rowsBetween而非rangeBetween定义窗口，是因为行级窗口不依赖time的连续值，即使time存在断档，只要按排序后取连续3条记录即可符合滚动窗口逻辑，兼容性更强
• 该写法所有函数均为PySpark 2.4内置支持，不需要额外依赖、自定义UDF，执行性能高

内容的提问来源于stack exchange，提问作者quant