PySpark如何实现按列分组后每组随机抽取最多n行

核心性能结论 #

所有按指定字段分组的操作本身就需要按分组键做shuffle，不存在完全规避shuffle的分组取数方案；无排序需求的场景下，使用常量排序的窗口函数性能最优，不存在比该方案性能显著更高的原生实现。

最优实现代码 #

不需要指定业务字段排序，仅需在窗口定义时用常量作为排序键，Spark会在每个分组内按数据天然读取顺序取前N行，不会触发业务字段的全量排序开销，shuffle量和纯groupBy操作完全一致：

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

# 测试数据构造
data = [('A',1), ('B',1), ('C',2)]
columns = ["field_1","field_2"]
df = spark.createDataFrame(data=data, schema = columns)

# 定义仅按分组键分区、无业务字段排序的窗口
w = Window.partitionBy("field_2").orderBy(F.lit(1))
# 取每组最多1行，可修改<=后面的数值调整抽取行数
res = df.withColumn("rn", F.row_number().over(w)) \
        .filter(F.col("rn") <= 1) \
        .drop("rn")

res.show()

执行效果和pandas的df.groupby('field_2').head(1)完全一致，不保证组内返回行的顺序，仅保证每个分组最多返回指定数量的行。

等价SQL写法 #

-- 构造测试数据
CREATE TEMP VIEW test_df AS
SELECT * FROM VALUES ('A',1), ('B',1), ('C',2) AS t(field_1, field_2);

-- 分组取数
SELECT field_1, field_2 FROM (
    SELECT *,
           ROW_NUMBER() OVER(PARTITION BY field_2 ORDER BY 1) AS rn
    FROM test_df
) t
WHERE rn <= 1;

其他方案的局限性 #

• sampleBy()：仅支持按比例抽样，无法实现每个分组精确抽取固定绝对行数的需求，抽样结果存在随机误差，数据倾斜场景下误差会被放大
• 带业务字段排序的窗口函数：会额外触发分组内的业务字段排序，在数据量较大时排序开销占比高，无排序需求时没必要使用
• 原生groupBy：内置聚合函数仅支持返回聚合计算结果，无法直接保留原始行的多字段明细；自定义UDAF实现取数逻辑的开发成本高，性能远低于窗口函数方案
• 自定义RDDmapPartitions实现：需要手动处理shuffle分区逻辑，代码复杂度高，实际性能和常量排序窗口方案基本持平，没有额外收益

低版本Spark（2.4及以下）如果出现常量排序被执行计划优化器忽略的问题，可以将排序键替换为F.monotonically_increasing_id()，该函数生成全局递增的唯一ID，不需要基于业务数据做排序计算，额外开销可以忽略。

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