PySpark sampleBy函数抽样结果异常的技术咨询

问题背景 #

使用PySpark SQL的sampleBy函数时，发现抽样行为存在明显不确定性：对包含100个值（key=0、1、2各占约1/3）的数据集分层抽样，设置key=0抽样比例10%、key=1抽样比例20%、key=2默认0%，预期得到约3个key=0、7个key=1的样本，但实际结果与预期不符，且修改seed值时抽样分布会变化。

测试代码 #

from pyspark.sql.functions import col
from pyspark.sql import SQLContext
sqlC = SQLContext(sc)

dataset = sqlC.range(0, 100).select((col("id") % 3).alias("key"))
dataset.groupBy("key").count().orderBy("key").show()
#+---+-----+
#|key|count|
#+---+-----+
#|  0|   34|
#|  1|   33|
#|  2|   33|
#+---+-----+

# With seed = 0
sampled = dataset.sampleBy("key", fractions={0: 0.1, 1: 0.2}, seed=0)
sampled.groupBy("key").count().orderBy("key").show()
#+---+-----+
#|key|count|
#+---+-----+
#|  0|    6|
#|  1|   11|
#+---+-----+

# With seed = 123
sampled = dataset.sampleBy("key", fractions={0: 0.1, 1: 0.2}, seed=123)
sampled.groupBy("key").count().orderBy("key").show()
#+---+-----+
#|key|count|
#+---+-----+
#|  0|    1|
#|  1|   11|
#+---+-----+

解答 #

这种现象是完全正常的，核心原因和你可能遗漏的细节如下：

1. sampleBy采用伯努利独立抽样逻辑

   • 该函数对每个数据行独立判断是否被选中，并非严格按比例抽取固定数量。比如key=0的34条数据，每条有10%概率被抽中，实际抽中数量服从二项分布，理论均值是3.4，但实际结果会在合理范围内波动（比如0到7之间都是概率允许的结果）。你预期的“约3个0、7个1”只是均值附近的参考值，抽样本身的随机性必然导致结果偏离均值。

2. seed值直接控制随机序列生成

   • seed是随机数生成器的初始参数，不同seed会生成不同的随机序列，进而导致不同的抽样结果。这是随机抽样的标准行为——只要固定seed，每次抽样结果就会完全一致，这也是seed参数的核心作用：保证抽样的可复现性。

3. 若需精确比例抽样的替代方案

   • 如果需要严格按分层比例抽取固定数量的样本，Spark 3.3及以上版本可使用stratifiedSample函数，它支持指定各分层的精确抽样数量；若使用低版本Spark，可通过先对分层数据进行排序，再按比例截取的方式实现，但这种方式效率较低且无法保证随机性。

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