Spark SQL如何按列值分布占比随机抽取指定总量的样本行？

核心逻辑是分层抽样，按指定列（本示例为prod_name）分组后，匹配原数据的分组占比抽取样本，最终抽样结果的分布和原数据集一致。

方法1：DataFrame API内置分层抽样（大数据量场景推荐） #

Spark的stat.sampleBy是官方提供的分层抽样接口，底层做了性能优化，适合处理大规模数据。
实现步骤：

1. 计算各分组的总条数，结合预期总抽样行数算出每个分组的抽样比例
2. 调用sampleBy按比例抽样，可指定随机种子保证结果可复现

示例代码（PySpark）：

from pyspark.sql import functions as F

# 假设你的源DataFrame名为df
# 配置总抽样行数
total_sample_rows = 6
total_original_rows = df.count()

# 计算每个分组的抽样比例
group_frac_df = df.groupBy("prod_name") \
                  .agg(F.count("*").alias("group_count")) \
                  .withColumn("sample_frac", 
                      # 分组占比 * 总抽样行数 / 分组原有行数 = 该分组抽样比例
                      (F.col("group_count") / total_original_rows * total_sample_rows) / F.col("group_count")
                  )
# 转成分组->抽样比例的字典
sample_frac_map = dict(group_frac_df.rdd.map(lambda x: (x["prod_name"], x["sample_frac"])).collect())

# 无放回分层抽样，seed为随机种子可自行修改
sampled_df = df.stat.sampleBy("prod_name", fractions=sample_frac_map, seed=42)

本示例中A、B、C三个分组各占1/3，总抽6行对应每个分组抽2行，抽样比例为2/4=0.5，运行后即可得到每个分组约2行的抽样结果。

方法2：纯Spark SQL实现（适合SQL开发场景） #

如果偏好SQL开发，可以通过窗口函数实现分层抽样：

WITH group_meta AS (
    -- 统计各分组的基础指标
    SELECT
        prod_name,
        COUNT(*) AS group_cnt,
        SUM(COUNT(*)) OVER() AS total_cnt,
        6 AS total_sample_rows -- 配置总抽样行数
    FROM your_table_name -- 替换为你的表名
    GROUP BY prod_name
),
group_ranked AS (
    -- 给每个分组内的行随机排序
    SELECT
        t.*,
        ROW_NUMBER() OVER(PARTITION BY t.prod_name ORDER BY RAND()) AS rn,
        -- 计算每个分组需要抽取的行数
        ROUND(gm.group_cnt / gm.total_cnt * gm.total_sample_rows) AS group_sample_cnt
    FROM your_table_name t
    JOIN group_meta gm ON t.prod_name = gm.prod_name
)
-- 过滤取每个分组前N行
SELECT prod_name, value
FROM group_ranked
WHERE rn <= group_sample_cnt

注意事项 #

• 如需固定抽样结果方便复现，可以替换RAND()为带种子的RAND(42)，或给sampleBy指定固定seed参数
• 如果某分组需要抽取的行数超过该分组原有数据行数，需要开启有放回抽样：sampleBy的第一个参数设为True；SQL写法可以将ROW_NUMBER替换为多次随机生成逻辑
• 随机抽样存在小幅波动，如对比例精度要求极高，可以多次抽样取最符合预期的结果

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