Spark抽取指定分布数据样本：替代多查询Union的更优方案？

以下是我的数据集 schema：

request_type | request_body
1              body A
2              body B
3              ...
4
5
6
..
32            body XXX

• 总共需要获取500条记录。
• 数据集包含32种request_type。
• request_type需按以下方式分桶抽取：
  • request_type 1、2、3、4各占20%，5至32整体占20%
  • request_type 1、2、3、4各抽取100条，总计400条
  • 最后一个分桶需抽取100条request_type为5至32的记录

我目前采用的暴力解法是运行5条Spark SQL查询后合并结果，请问是否有更优的实现方案？

最优方案：单查询实现分桶抽样 #

可以借助Spark窗口函数+条件过滤的方式，仅扫描一次数据集完成所有抽样逻辑，避免多次查询带来的重复IO开销，同时代码更简洁易维护。

Spark SQL 实现（精确条数控制） #

WITH ranked_records AS (
    SELECT 
        request_type,
        request_body,
        -- 对1-4单独分组，5-32归为同一分组
        CASE 
            WHEN request_type IN (1,2,3,4) THEN request_type
            ELSE 0 
        END AS bucket_key,
        -- 按分组随机排序后编号
        ROW_NUMBER() OVER (
            PARTITION BY CASE 
                            WHEN request_type IN (1,2,3,4) THEN request_type
                            ELSE 0 
                        END 
            ORDER BY RAND()
        ) AS row_num
    FROM your_table_name
)
SELECT request_type, request_body
FROM ranked_records
WHERE 
    -- 1-4类型各取前100条
    (bucket_key IN (1,2,3,4) AND row_num <= 100)
    -- 5-32类型取前100条
    OR (bucket_key = 0 AND row_num <= 100)

方案优势 #

1. 性能更优：仅扫描一次数据集，避免多次查询的重复IO操作。
2. 逻辑统一：所有抽样规则在单个查询中实现，代码易读易维护。
3. 条数精确：通过ROW_NUMBER()配合随机排序，确保每个分组精确抽取指定条数（前提是对应分组的数据量满足抽取需求）。

备选优化方案（分批次抽样） #

如果5-32类型的数据量极大，也可以拆分两个子查询合并，仅扫描两次数据集，逻辑同样清晰：

-- 抽取1-4类型各100条
WITH type_1_4_sample AS (
    SELECT request_type, request_body
    FROM your_table_name
    WHERE request_type IN (1,2,3,4)
    QUALIFY ROW_NUMBER() OVER (PARTITION BY request_type ORDER BY RAND()) <= 100
),
-- 抽取5-32类型100条
type_5_32_sample AS (
    SELECT request_type, request_body
    FROM your_table_name
    WHERE request_type BETWEEN 5 AND 32
    ORDER BY RAND()
    LIMIT 100
)
SELECT * FROM type_1_4_sample
UNION ALL
SELECT * FROM type_5_32_sample

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