如何用PySpark实现满足多列比例要求的10万行分层抽样？

要实现同时满足OCCUPA_ID和DIG_LABEL双重比例的10万行抽样，核心思路是先计算两个字段交叉组合的目标抽样量，再针对每个组合单独抽样后合并，具体步骤如下：

1. 定义目标参数 #

先明确总样本量、各维度的比例要求：

total_samples = 100000
# OCCUPA_ID 各取值的占比与对应样本量
occupa_proportions = {2705: 0.7, 2708: 0.2, 2712: 0.1}
occupa_targets = {k: int(total_samples * v) for k, v in occupa_proportions.items()}
# DIG_LABEL 各取值的占比（需应用到每个OCCUPA_ID分组内）
dig_label_proportions = {1: 0.5, 2: 0.2, 7: 0.3}

2. 计算交叉组合的目标抽样量 #

对每个OCCUPA_ID分组，再按DIG_LABEL的比例拆分目标数量：

# 生成所有交叉组合的目标抽样字典
target_counts = {}
for occupa_id, occ_total in occupa_targets.items():
    for dig_label, dig_ratio in dig_label_proportions.items():
        target_counts[(occupa_id, dig_label)] = int(occ_total * dig_ratio)

比如(2705, 1)的目标量是70000 * 0.5 = 35000，(2708,7)的目标量是20000 * 0.3 = 6000，以此类推。

3. 按交叉组合分组抽样 #

使用PySpark的groupBy结合applyInPandas实现分组抽样：

from pyspark.sql.functions import struct

def sample_group(df):
    # 获取当前组的OCCUPA_ID和DIG_LABEL
    group_key = (df['OCCUPA_ID'].iloc[0], df['DIG_LABEL'].iloc[0])
    target = target_counts.get(group_key, 0)
    # 如果组内数据量不足目标，直接取全部；否则随机抽样
    if len(df) <= target:
        return df
    else:
        return df.sample(n=target, random_state=42)

# 按OCCUPA_ID和DIG_LABEL分组，应用抽样函数
sampled_df = df.groupBy(struct("OCCUPA_ID", "DIG_LABEL")).applyInPandas(
    sample_group,
    schema=df.schema
)

4. 验证抽样结果 #

可以通过以下代码检查抽样后的比例是否符合要求：

# 检查OCCUPA_ID比例
sampled_df.groupBy("OCCUPA_ID").count().withColumn(
    "proportion", (sampled_df.groupBy("OCCUPA_ID").count()["count"] / total_samples)
).show()

# 检查DIG_LABEL比例
sampled_df.groupBy("DIG_LABEL").count().withColumn(
    "proportion", (sampled_df.groupBy("DIG_LABEL").count()["count"] / total_samples)
).show()

注意事项 #

• 如果某个交叉组合的原始数据量小于目标抽样量，会直接取该组合的全部数据，最终总样本量可能略小于10万，可根据实际需求调整（比如从其他组合补充，或放宽比例要求）
• 设置random_state是为了保证抽样结果可复现，可根据需要调整
• 对于百万级数据，applyInPandas的性能完全够用，若数据量极大可考虑RDD层面的操作

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