PySpark按固定行数分箱聚合低内存实现及剩余行处理咨询

实现方案 #

因为你已经提前按TIME_STAMP完成了数据的全局排序，以下两种方案都无需将全量数据拉取到驱动节点，完全分布式执行不会触发OOM。

方案1：RDD zipWithIndex方案（超大数据集优先推荐） #

该方案利用RDD的zipWithIndex接口分布式生成连续递增行号，比全局窗口排序性能更高，适合超大规模数据集。

from pyspark.sql import functions as F

# df为你已按TIME_STAMP排序完成的原始DataFrame
# 生成带连续行号的RDD，行号从1开始计数
rdd_with_idx = df.rdd.zipWithIndex().map(lambda x: (x[0][0], x[0][1], x[1]+1))
df_with_row = rdd_with_idx.toDF(["TIME_STAMP", "RESULT", "row_idx"])

# 每5条为一组聚合
grouped_df = df_with_row.groupBy(F.floor((F.col("row_idx") -1)/5).alias("group_id"))\
    .agg(
        F.max("row_idx").alias("bin"),
        F.avg("RESULT").alias("mean")
    )\
    .select("bin", "mean")\
    .orderBy("bin")

grouped_df.show()

方案2：DataFrame窗口函数方案 #

如果不想转换RDD，可直接用窗口函数实现，代码更简洁：

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

# 按TIME_STAMP全局排序的窗口
w = Window.orderBy("TIME_STAMP")
df_with_row = df.withColumn("row_idx", F.row_number().over(w))

# 分组聚合逻辑和方案1一致
grouped_df = df_with_row.groupBy(F.floor((F.col("row_idx") -1)/5).alias("group_id"))\
    .agg(
        F.max("row_idx").alias("bin"),
        F.avg("RESULT").alias("mean")
    )\
    .select("bin", "mean")\
    .orderBy("bin")

grouped_df.show()

补充问题解答 #

两种方案天然支持总记录数无法被5整除的场景，比如19条记录的情况，前3组每组5条、bin值对应为5/10/15，最后一组自动保留剩余4条、bin值为19，完全符合你的需求。

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