如何对PySpark DataFrame采样指定精确行数而非按比例采样

PySpark 内置的sample()函数属于概率性采样，仅能保证返回行数的期望值接近你按比例计算的理论值，无法直接返回固定精确行数的样本。以下是两种适用不同场景的实现方案：

方案1：随机排序+limit（小数据量/低采样占比场景） #

这是最简洁的实现方式，代码如下：

from pyspark.sql.functions import rand

# 定义你需要采样的精确行数
target_count = 200000
sampled_df = df.orderBy(rand()).limit(target_count)

相关说明：

• 无需提前统计DataFrame总行数，实现成本极低
• 缺点是需要对全量数据做排序，当你的DataFrame数据量达到十亿级以上、或者集群资源不足时，排序操作的开销会很高，性能较差
• 如果需要可复现的采样结果，给rand()传入固定随机种子即可，例如rand(seed=42)

方案2：比例采样兜底+二次取数（大数据量场景，性能优先） #

如果数据量很大，全量排序性价比太低，可以先用超量比例采样缩小数据范围，再做二次取数保证精确性：

from pyspark.sql.functions import rand

target_count = 200000
total_count = df.count()
# 采样比例比目标多10%做冗余，避免概率采样行数不足的问题
sample_fraction = (target_count / total_count) * 1.1
# 第一步先按冗余比例采样，大幅缩小数据规模
temp_sample = df.sample(fraction=sample_fraction)
# 第二步对小数据量排序取前N行，得到精确行数的结果
sampled_df = temp_sample.orderBy(rand()).limit(target_count)

相关说明：

• 性能远高于全量排序方案，绝大多数生产场景都适用
• 冗余比例可根据实际场景调整，采样占比越低可以适当把冗余比例调得越高（极端场景可以提到20%~50%），保证第一次采样的行数一定大于目标行数
• 同样可以传入随机种子保证采样结果可复现，sample()也支持传seed参数，例如df.sample(fraction=sample_fraction, seed=42)

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