高效计算PySpark数据帧统计量并转换为Pandas数据帧

核心原则是仅对原始PySpark DataFrame做1次全表扫描，所有统计计算全部在Spark侧分布式完成，最后仅拉取聚合后的小结果集转换为Pandas对象，完全避免重复遍历、全量数据拉取到本地这类性能损耗。

具体实现步骤 #

1. 先做时间范围过滤：利用Spark的谓词下推能力，提前裁剪掉不在指定start_ts到end_ts区间的无效数据，从源头减少后续计算的数据量。
2. 按ts分组做聚合：
   • 平均值、最小值、最大值直接用Spark原生内置的avg/min/max聚合函数，这几个函数经过Catalyst优化器深度优化，性能极高
   • 中位数计算不要自定义Python UDF，也不要把原始数据拉到Driver端用Pandas计算，直接用Spark内置的percentile_approx函数计算0.5分位数即可，对于固定2小时间隔的分组场景，设置足够的精度参数就能得到几乎无误差的中位数结果，且全程分布式执行。如果要求100%精确中位数，替换为内置的percentile函数即可，性能依然远高于自定义实现。
3. 聚合完成后按ts排序，直接调用toPandas()转换为Pandas DataFrame即可。聚合后的结果为2小时间隔、跨度12年的统计值，总数据量仅5万条左右，完全不存在内存压力。

可直接运行的代码 #

from pyspark.sql.functions import avg, min, max, percentile_approx, col, lit
import datetime
import pandas as pd

# 定义统计时间范围
start_ts = datetime.datetime(year=2010, month=2, day=1, hour=0)
end_ts = datetime.datetime(year=2022, month=6, day=1, hour=22)

# 单次扫描完成所有过滤、分组、聚合计算
agg_result = df.filter(
    (col("ts") >= lit(start_ts)) & (col("ts") <= lit(end_ts))
).groupBy("ts").agg(
    avg("value").alias("average"),
    min("value").alias("min"),
    max("value").alias("max"),
    # 第三个参数为精度控制，数值越大精度越高，10000足够覆盖绝大多数业务场景
    percentile_approx(col("value"), 0.5, 10000).alias("median")
)

# 转换为Pandas DataFrame并按时间排序
result_pdf = agg_result.orderBy("ts").toPandas()

性能优势说明 #

• 原始大表仅遍历1次：过滤、分组、所有聚合逻辑在同一个Spark计算流程中完成，无额外重复扫描、多余shuffle开销
• 计算全分布式执行：所有统计逻辑在集群节点并行计算，不需要把海量原始数据拉取到Driver端，从根本上避免Driver端OOM问题
• 全原生函数实现：所有聚合逻辑都用Spark内置优化过的函数，比Python UDF实现性能高1~2个数量级
• 无效数据提前裁剪：时间过滤逻辑会下推到数据源层（如果是表存储的话），不会加载不需要的数据到计算链路

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