如何在PySpark中实现扩展窗口结合百分位数计算？

可以实现，无需将整个DataFrame转换为Pandas。PySpark没有直接提供expanding().quantile()的API，但可以通过Pandas分组映射UDAF或窗口函数+自定义UDF两种方式实现，以下是具体方案：

方案一：Pandas GROUPED_MAP UDAF（推荐，性能更优） #

该方式通过按分组将数据转换为Pandas小DataFrame计算，再合并结果，不会将整个大数据集转为Pandas，适合大数据量场景。

完整代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, PandasUDFType

# 初始化Spark会话
spark = SparkSession.builder.appName("ExpandingQuantile").getOrCreate()

# 创建测试数据
data = [
    (1, 90), (1, 75), (1, 60),
    (2, 65), (2, 80), (2, 95),
    (3, 100), (3, 110), (3, 115)
]
df = spark.createDataFrame(data, ["TRIP_ID", "SPEED"])

# 定义分组映射UDAF
@pandas_udf(df.schema.add("95_PERC", "double"), PandasUDFType.GROUPED_MAP)
def compute_expanding_quantile(pdf):
    # 确保分组内的行顺序与原始数据一致（若有排序字段可替换为该字段排序）
    pdf = pdf.reset_index(drop=True)
    # 计算扩展窗口的0.95分位数
    pdf["95_PERC"] = pdf["SPEED"].expanding().quantile(0.95)
    return pdf

# 执行计算并展示结果
result_df = df.groupBy("TRIP_ID").apply(compute_expanding_quantile)
result_df.show()

输出结果 #

+-------+-----+-------+
|TRIP_ID|SPEED|95_PERC|
+-------+-----+-------+
|      1|   90|   90.0|
|      1|   75|  89.25|
|      1|   60|   88.5|
|      2|   65|   65.0|
|      2|   80|  79.25|
|      2|   95|   93.5|
|      3|  100|  100.0|
|      3|  110|  109.5|
|      3|  115|  114.5|
+-------+-----+-------+

方案二：窗口函数+collect_list+自定义UDF #

该方式通过窗口收集当前行及之前的所有SPEED值，再用UDF计算分位数，适合小数据量场景，实现更简单。

完整代码示例 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window
import numpy as np

# 添加行号保证分组内的顺序
df = df.withColumn(
    "row_num",
    F.row_number().over(Window.partitionBy("TRIP_ID").orderBy(F.monotonically_increasing_id()))
)

# 定义扩展窗口，收集当前行及之前的SPEED值
expanding_window = Window.partitionBy("TRIP_ID").orderBy("row_num").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df_with_list = df.withColumn("speed_list", F.collect_list("SPEED").over(expanding_window))

# 定义计算分位数的UDF
quantile_udf = F.udf(lambda x: np.quantile(x, 0.95), F.DoubleType())

# 生成结果并清理临时列
result_df = df_with_list.withColumn("95_PERC", quantile_udf(F.col("speed_list"))).drop("row_num", "speed_list")
result_df.show()

关键注意事项 #

1. 行顺序保证：PySpark DataFrame本身是无序的，必须显式指定分组内的排序规则（如行号、时间戳），否则分位数计算结果会与预期不符。
2. 性能选择：大数据量下优先使用GROUPED_MAP UDAF，避免collect_list带来的内存压力；小数据量可使用窗口+UDF的方式快速实现。

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