PySpark如何按用户分组后按1:5比例实现正负样本随机采样

• 第一步：统计每个用户的正样本（label=1）总数量，由于原始数据正负样本比例极低，默认保留全部正样本，仅对负样本做随机下采样
• 第二步：给每个用户下的正负样本分别做随机排序，打上分组内的排名序号
• 第三步：筛选样本：正样本全部保留，负样本仅保留排名在前「该用户正样本数*5」的部分，若用户负样本总数不足5倍正样本数，则保留该用户所有负样本
• 第四步：合并筛选后的正负样本，得到最终采样结果

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 1. 统计每个用户的正样本数量
user_pos_count = df.filter(F.col("label") == 1) \
                   .groupBy("userid") \
                   .agg(F.count("*").alias("pos_total"))

# 2. 将正样本数量关联回原数据集
df = df.join(user_pos_count, on="userid", how="left")

# 3. 定义窗口：按用户+标签分组，随机排序后打排名
rank_window = Window.partitionBy("userid", "label").orderBy(F.rand())
df = df.withColumn("group_rn", F.row_number().over(rank_window))

# 4. 按1:5比例筛选样本
sample_df = df.filter(
    # 正样本全部保留
    (F.col("label") == 1) |
    # 负样本保留前 正样本数*5 条
    (F.col("label") == 0) & (F.col("group_rn") <= F.col("pos_total") * 5)
).drop("pos_total", "group_rn")

# 查看结果
sample_df.orderBy("userid", "label").show()

• 若存在没有正样本的用户，上述逻辑会自动过滤掉该用户的所有数据。如果需要保留这类用户，可以在关联正样本数量后用F.fillna(0, subset=["pos_total"])填充空值，再自定义这类用户的采样规则即可
• 若需要固定每次采样的结果，可以将F.rand()替换为F.rand(seed=自定义整数)，比如F.rand(42)即可固定随机种子
• 如果你的正样本量足够大不需要全留，只需要修改正样本的筛选条件，比如(F.col("label") == 1) & (F.col("group_rn") <= 你需要的正样本数)即可

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