如何在Hive SQL中实现基于薪资列的加权数据抽样？

当然可以实现按薪资（或其他数值列）的加权抽样，核心思路是让样本被选中的概率与薪资值成正比，以下是几种可行的实现方式：

方法一：基于累计薪资区间的无放回抽样（固定样本量） #

这种方法通过计算累计薪资区间，结合随机数匹配区间来抽取样本，能保证薪资越高的行被选中的概率越高，且不会重复抽取同一行。

单样本抽取 #

WITH total_salary AS (
    -- 计算全局总薪资
    SELECT SUM(salary) AS total FROM employee
),
ranked_employees AS (
    SELECT 
        empid, deptid, salary,
        -- 计算到当前行的累计薪资
        SUM(salary) OVER (ORDER BY empid) AS cumulative_salary,
        -- 计算到上一行的累计薪资（当前区间的起始值）
        (SUM(salary) OVER (ORDER BY empid) - salary) AS prev_cumulative_salary,
        (SELECT total FROM total_salary) AS total_salary
    FROM employee
)
-- 生成0到总薪资的随机数，匹配区间选中对应行
SELECT empid, deptid, salary
FROM ranked_employees
WHERE rand() * total_salary BETWEEN prev_cumulative_salary AND cumulative_salary
LIMIT 1;

多样本抽取 #

如果需要抽取多个样本，可以生成多个随机数进行匹配：

WITH total_salary AS (
    SELECT SUM(salary) AS total FROM employee
),
ranked_employees AS (
    SELECT 
        empid, deptid, salary,
        SUM(salary) OVER (ORDER BY empid) AS cumulative_salary,
        (SUM(salary) OVER (ORDER BY empid) - salary) AS prev_cumulative_salary,
        (SELECT total FROM total_salary) AS total_salary
    FROM employee
),
random_numbers AS (
    -- 生成3个随机数，对应抽取3个样本
    SELECT rand() * (SELECT total FROM total_salary) AS rnd UNION ALL
    SELECT rand() * (SELECT total FROM total_salary) AS rnd UNION ALL
    SELECT rand() * (SELECT total FROM total_salary) AS rnd
)
-- 匹配随机数所在的薪资区间
SELECT re.empid, re.deptid, re.salary
FROM ranked_employees re
JOIN random_numbers rn 
ON rn.rnd BETWEEN re.prev_cumulative_salary AND re.cumulative_salary;

方法二：基于权重概率的独立抽样（有放回） #

这种方式是让每行按薪资/总薪资的概率被选中，属于有放回抽样，可能会出现重复样本，适合不需要严格固定样本量的场景：

WITH total_salary AS (
    SELECT SUM(salary) AS total FROM employee
)
SELECT empid, deptid, salary
FROM employee, total_salary
-- 乘以10是调整抽样比例，提升命中率，可根据需求调整
WHERE rand() < (salary / total_salary) * 10;

按分区（部门）加权抽样 #

如果需要在每个部门内单独按薪资加权抽样，只需在窗口函数中添加PARTITION BY deptid即可：

WITH dept_total AS (
    -- 计算每个部门的总薪资
    SELECT deptid, SUM(salary) AS dept_total FROM employee GROUP BY deptid
),
ranked_dept_employees AS (
    SELECT 
        empid, e.deptid, salary,
        SUM(salary) OVER (PARTITION BY e.deptid ORDER BY empid) AS cumulative_salary,
        (SUM(salary) OVER (PARTITION BY e.deptid ORDER BY empid) - salary) AS prev_cumulative_salary,
        dt.dept_total
    FROM employee e
    JOIN dept_total dt ON e.deptid = dt.deptid
),
dept_random AS (
    -- 每个部门生成一个随机数
    SELECT deptid, rand() * dept_total AS rnd FROM dept_total
)
-- 每个部门匹配对应的随机数区间
SELECT r.empid, r.deptid, r.salary
FROM ranked_dept_employees r
JOIN dept_random dr ON r.deptid = dr.deptid
WHERE dr.rnd BETWEEN r.prev_cumulative_salary AND r.cumulative_salary;

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