如何高效从按date分区的PySpark DataFrame中按id整组采样？

你的场景是：处理一个按date分区的超大PySpark DataFrame，需要实现ID级全量采样——只要某个id被选中，该id的所有行都要纳入样本。你当前的实现逻辑可行，但存在可以优化的空间，以下是更高效的几种方案：

问题分析：当前写法的不足 #

你当前的代码中，df[['id']].sample(fraction=0.001)会保留大量重复的id，后续join时会导致同一id的行被多次匹配，最终结果出现重复数据，既浪费计算资源，还需要额外去重步骤。

优化方案1：先去重ID再采样（最稳妥的通用方案） #

先提取所有唯一id，再对这个无重复的集合采样，最后关联原表。这种方式彻底避免了重复id带来的冗余计算：

# 提取唯一ID集合并采样
unique_ids = df.select('id').distinct()
sampled_ids = unique_ids.sample(fraction=0.001)

# 关联原表得到全量采样数据
sampled_df = sampled_ids.join(df, on='id')

优化方案2：用广播+isin过滤（轻量场景首选） #

如果采样后的id数量较少（比如几千以内），可以直接把采样后的id集合广播到所有Executor，用isin过滤原表，省去join操作的开销：

# 提取唯一ID、采样后转为本地列表
sampled_id_list = df.select('id').distinct().sample(fraction=0.001).rdd.map(lambda x: x.id).collect()

# 广播ID列表并过滤原表
from pyspark.sql.functions import broadcast
sampled_df = df.filter(df.id.isin(broadcast(sampled_id_list)))

⚠️ 注意：如果采样后的ID数量过大（比如超过10万），collect()会占用Driver节点过多内存，此时优先用方案1。

优化方案3：利用分区特性减少全局Shuffle #

因为你的表按date分区，若id在各分区的分布相对均匀，可以先在每个分区内提取唯一ID并采样，再全局去重后关联原表，把部分计算分散到各个Executor，降低全局数据传输压力：

# 分区内提取ID→全局去重→采样→转DataFrame
sampled_ids = df.rdd.mapPartitions(lambda part: [(row.id,) for row in part]) \
                   .distinct() \
                   .sample(False, 0.001) \
                   .toDF(['id'])

# 关联原表
sampled_df = sampled_ids.join(df, on='id')

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