如何用SQL按8:1:1比例基于Column Y拆分数据集

实现思路 #

基于每个Column X对应的总计数Column Y，按80%/10%/10%的比例拆分Train、Test、Val数据集，核心是先计算每个分组的各数据集分配数量，再给每条数据随机标记所属类别。

具体代码（MySQL兼容） #

第一步：计算每个分组的分配数量 #

先算出每个Column X对应的Train、Test、Val具体条数，严格匹配你给出的示例逻辑（余数优先分配给Test）：

WITH group_split AS (
    SELECT
        Column_X,
        Column_Y,
        -- 训练集取80%向下取整
        FLOOR(Column_Y * 0.8) AS train_num,
        -- 测试集取10%向下取整，若有剩余条数优先分配给测试集
        FLOOR(Column_Y * 0.1) + IF(Column_Y - FLOOR(Column_Y * 0.8) - FLOOR(Column_Y * 0.1)*2 > 0, 1, 0) AS test_num,
        -- 验证集用总数减去训练和测试的数量，确保总数匹配
        Column_Y - FLOOR(Column_Y * 0.8) - (FLOOR(Column_Y * 0.1) + IF(Column_Y - FLOOR(Column_Y * 0.8) - FLOOR(Column_Y * 0.1)*2 > 0, 1, 0)) AS val_num
    FROM your_table
)

第二步：给每条数据标记数据集类别 #

通过行号随机分配每条数据到对应集合：

SELECT
    raw.Column_X,
    raw.Column_Y,
    CASE
        WHEN row_id <= gs.train_num THEN 'Train'
        WHEN row_id <= gs.train_num + gs.test_num THEN 'Test'
        ELSE 'Val'
    END AS dataset_type
FROM (
    SELECT
        Column_X,
        Column_Y,
        -- 按Column_X分组，随机生成行号保证拆分的随机性
        ROW_NUMBER() OVER (PARTITION BY Column_X ORDER BY RAND()) AS row_id
    FROM your_table
) raw
JOIN group_split gs ON raw.Column_X = gs.Column_X
ORDER BY raw.Column_X, dataset_type;

关键细节说明 #

• ROW_NUMBER() OVER (PARTITION BY Column_X ORDER BY RAND())：确保每个Column X分组内的数据被随机分配，避免拆分结果有偏差
• 余数分配逻辑：当Column Y无法被10整除时，多出来的条数优先分给Test，和你给出的示例（Y=23拆18/3/2、Y=64拆51/7/6）完全匹配
• 若使用其他SQL方言（如PostgreSQL），只需调整随机排序的语法（比如PostgreSQL用ORDER BY random()）即可

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