求PySpark中等价于pandas groupby('col1').col2.head()的实现方法

你说的这个需求我之前也碰到过，确实Spark没法像pandas那样直接用head()实现，但完全不用循环遍历分组，有两种高效的方案可以搞定：

方案一：使用窗口函数（推荐，支持保留全量列） #

这是最通用的方法，借助row_number()窗口函数给每个分组内的行编号，再筛选出前N条（这里N=10）。如果需要随机采样，还能结合随机排序：

from pyspark.sql import Window
from pyspark.sql.functions import row_number, rand

# 定义窗口规则：按col1分区，随机排序（如果要固定顺序，换成你需要的字段比如col2即可）
window_spec = Window.partitionBy("col1").orderBy(rand())

# 添加行号列，筛选前10条后删除行号列
result_df = df.withColumn("row_idx", row_number().over(window_spec)) \
              .filter("row_idx <= 10") \
              .drop("row_idx")

补充说明： #

• 如果不需要随机样本，把orderBy(rand())改成orderBy("col2")或者其他字段，就能按指定顺序取前10条；
• 这个方法可以保留原DataFrame的所有列，要是你除了col1和col2还有其他字段需要保留，选这个方案最合适；
• 性能表现很友好，Spark会并行处理每个分组，完全没有循环带来的性能瓶颈。

方案二：分组聚合+切片（适合仅需col1和col2的场景） #

如果你的需求只需要保留col1和采样后的col2，可以用collect_list收集每个分组的col2值，再用slice截取前10个，最后通过explode把列表拆成行：

from pyspark.sql.functions import collect_list, slice, explode

# 分组收集col2，截取前10个，再展开成行
result_df = df.groupBy("col1") \
              .agg(slice(collect_list("col2"), 1, 10).alias("sampled_col2")) \
              .selectExpr("col1", "explode(sampled_col2) as col2")

补充说明： #

• collect_list的顺序依赖于数据的原始顺序，如果需要随机采样，建议先对整个DataFrame做随机排序：df.orderBy(rand())再进行后续操作；
• 这个方案代码更简洁，但如果分组的行数特别多，collect_list可能会占用较多内存（要把整个分组的col2都加载到内存里再切片），所以大数据量下优先选方案一。

内容的提问来源于stack exchange，提问作者Renée