PySpark DataFrame新增随机抽样分组列及多分组实现方法

原sample+leftanti join的实现方式会触发额外shuffle开销，大数量级下性能较差，完全可以通过内置随机函数直接新增分组列实现，无需拆分DataFrame。

两组拆分（对照组+单测试组，各50%占比） #

使用rand()函数为每行生成0~1区间的均匀分布随机值，固定随机种子保证结果可复现，直接通过条件判断赋值分组即可：

from pyspark.sql import functions as F

df_with_group = df_segment.withColumn(
    "group",
    # 固定seed=0保证每次重跑分组结果一致
    F.when(F.rand(seed=0) < 0.5, "control").otherwise("treatment")
)

该实现无任何join、shuffle操作，仅做列级转换，性能远高于原实现，大样本下两组占比会稳定趋近50%。

多组拆分（支持任意分组数+自定义占比） #

如果需要拆分为1个对照组+N个测试组，只需要按占比设置分段阈值，链式调用when判断即可，以三等分（每组占1/3）为例：

df_with_multi_group = df_segment.withColumn(
    "rand_val", F.rand(seed=0)
).withColumn(
    "group",
    F.when(F.col("rand_val") < 1/3, "control")
     .when(F.col("rand_val") < 2/3, "treat_1")
     .otherwise("treat_2")
).drop("rand_val")

如果需要自定义各组占比，直接修改阈值即可：比如对照组占50%、treat1占30%、treat2占20%，只需要把阈值依次改为0.5、0.8。

如果分组数量较多，可以通过配置化方式动态生成分组逻辑，避免硬编码：

# 分组配置：格式为(分组名称, 累计占比阈值)，所有阈值最后一位必须为1，占比总和为1
group_conf = [
    ("control", 0.5),
    ("treat_1", 0.8),
    ("treat_2", 1.0)
]

# 动态拼接判断逻辑
group_expr = F.when(F.lit(False), F.lit(None))
for name, threshold in group_conf[:-1]:
    group_expr = group_expr.when(F.rand(seed=0) < threshold, name)
group_expr = group_expr.otherwise(group_conf[-1][0])

df_config_group = df_segment.withColumn("group", group_expr)

注意事项：分组时必须固定随机种子，否则每次任务重跑分组结果都会随机变化，无法满足AB实验的分组稳定性要求。大样本场景下rand()生成的均匀分布值可以保证各组占比和设定值偏差极小，不需要额外做占比校准。

方案优势 #

• 不需要拆分多个独立DataFrame对象，所有逻辑在原DataFrame上通过新增列完成
• 无join、shuffle等重开销操作，执行效率远高于原sample+join的实现
• 逻辑灵活可扩展，支持任意数量分组、任意占比配置，代码可复用性高

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