如何按标签占比抽取含全部标签的随机行子集（BigQuery/PostgreSQL）

我有一张包含tag列的大表，结构如下：

各标签在表中的占比不同。需要抽取一个随机行子集，要求：

1. 每个标签至少包含一行
2. 子集中各标签的占比与原表中标签的占比一致

当前我用按标签排序后选取每n行的方法，但部分占比较低的标签会被意外遗漏，查询语句如下：

SELECT t.id, t.tag
FROM
(
    SELECT id, tag, ROW_NUMBER() OVER (ORDER BY tag) AS rownum
    FROM datatable
) AS t
WHERE t.rownum % 100 = 0    -- can be any number

当前使用BigQuery，也可切换为PostgreSQL执行查询。

方法1：分层抽样 + 小标签兜底 #

核心思路是按标签分组进行随机抽样，对每个标签按原表占比计算应抽取的行数，同时确保小标签（原表行数少的）至少抽1行，从根源避免低占比标签被遗漏。

BigQuery 实现 #

WITH tag_stats AS (
    -- 计算每个标签的总行数、原表总数据量，以及每个标签的占比
    SELECT 
        tag,
        COUNT(*) AS tag_count,
        SUM(COUNT(*)) OVER () AS total_rows
    FROM datatable
    GROUP BY tag
),
sample_quota AS (
    -- 计算每个标签的抽样配额：按占比计算的行数，至少1行
    SELECT 
        tag,
        GREATEST(
            ROUND((tag_count / total_rows) * @target_sample_size),  -- 替换@target_sample_size为你期望的子集总大小
            1  -- 保证每个标签至少1行
        ) AS sample_size
    FROM tag_stats
)
-- 按每个标签的配额进行随机抽样
SELECT d.id, d.tag, d.data
FROM sample_quota q
JOIN datatable d ON q.tag = d.tag
QUALIFY ROW_NUMBER() OVER (PARTITION BY d.tag ORDER BY RAND()) <= q.sample_size;

PostgreSQL 实现 #

WITH tag_stats AS (
    SELECT 
        tag,
        COUNT(*) AS tag_count,
        SUM(COUNT(*)) OVER () AS total_rows
    FROM datatable
    GROUP BY tag
),
sample_quota AS (
    SELECT 
        tag,
        GREATEST(
            ROUND((tag_count::numeric / total_rows) * :target_sample_size),  -- 替换:target_sample_size为目标子集大小
            1
        )::int AS sample_size
    FROM tag_stats
)
SELECT d.id, d.tag, d.data
FROM (
    SELECT 
        d.*,
        ROW_NUMBER() OVER (PARTITION BY d.tag ORDER BY RANDOM()) AS rn
    FROM datatable d
) d
JOIN sample_quota q ON d.tag = q.tag
WHERE d.rn <= q.sample_size;

方法2：先取所有标签的至少一行，再按比例补充剩余行 #

如果需要更精准控制占比，可分两步操作：

1. 先为每个标签随机选1行，保证无标签遗漏
2. 计算剩余需要抽取的行数，按原表标签占比分配配额，再从剩余行中随机抽取（排除已选的行）

BigQuery 示例 #

DECLARE target_sample_size INT64 DEFAULT 10000;  -- 设置目标子集总大小

WITH tag_min AS (
    -- 每个标签先取1行
    SELECT id, tag, data
    FROM datatable
    QUALIFY ROW_NUMBER() OVER (PARTITION BY tag ORDER BY RAND()) = 1
),
remaining_quota AS (
    -- 计算剩余需要抽取的行数，以及各标签的剩余配额
    SELECT 
        ts.tag,
        GREATEST(
            ROUND((ts.tag_count / ts.total_rows) * (target_sample_size - (SELECT COUNT(*) FROM tag_min))),
            0
        ) AS remaining_sample_size
    FROM (
        SELECT 
            tag,
            COUNT(*) AS tag_count,
            SUM(COUNT(*)) OVER () AS total_rows
        FROM datatable
        GROUP BY tag
    ) ts
)
-- 合并基础行和补充行
SELECT * FROM tag_min
UNION ALL
SELECT d.id, d.tag, d.data
FROM remaining_quota rq
JOIN datatable d ON rq.tag = d.tag
-- 排除已经在基础行中的数据
WHERE NOT EXISTS (SELECT 1 FROM tag_min tm WHERE tm.id = d.id)
QUALIFY ROW_NUMBER() OVER (PARTITION BY d.tag ORDER BY RAND()) <= rq.remaining_sample_size;

当前方法遗漏低占比标签的原因 #

你当前用全局ROW_NUMBER()按tag排序后取模，相当于把所有数据按tag排好后间隔选取。如果某个标签的总行数小于取模步长（比如100），就可能刚好没有行的rownum能被步长整除，导致这个标签完全被遗漏。而分层抽样是针对每个标签单独处理，从根源避免了这个问题。

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