Snowflake中基于前后行条件为数据行分配分组标签

需求说明 #

按id分区、start_time排序，遵循以下规则分配分组标签：

• 当某行grp_start = 1时，开启一个新分组并分配唯一标签
• 后续连续行使用同一标签，直到遇到grp_end = 1的行（该行也属于当前分组）
• 单独一行同时满足grp_start = 1和grp_end = 1时，分配独立的分组标签
• 标签从0开始，每个新分组递增1

示例数据 #

当前数据 #

create or replace temporary table grp_test (id char(4), start_time date, grp_start int, grp_end int)
    as select * from values 
    ('0001','2021-01-10',1,0),
    ('0001','2021-01-11',0,0),
    ('0001','2021-01-14',0,1),
    ('0001','2021-07-01',1,1),
    ('0001','2021-09-25',1,0),
    ('0001','2021-09-29',0,1),
    ('0002','2022-11-04',1,0),
    ('0002','2022-11-25',0,1);
    
select * from grp_test;

期望结果 #

create or replace temporary table desired_result (id char(4), start_time date, grp_start int, grp_end int, label int)
    as select * from values 
    ('0001','2021-01-10',1,0,0),
    ('0001','2021-01-11',0,0,0),
    ('0001','2021-01-14',0,1,0),
    ('0001','2021-07-01',1,1,1),
    ('0001','2021-09-25',1,0,2),
    ('0001','2021-09-29',0,1,2),
    ('0002','2022-11-04',1,0,0),
    ('0002','2022-11-25',0,1,0);
    
select * from desired_result;

解决方案 #

使用Snowflake窗口函数实现高效分组标签分配，代码如下：

SELECT 
    id,
    start_time,
    grp_start,
    grp_end,
    -- 按id分区排序，累计统计grp_start=1的次数，减1得到从0开始的连续标签
    COUNT(CASE WHEN grp_start = 1 THEN 1 END) OVER (
        PARTITION BY id 
        ORDER BY start_time 
        ROWS UNBOUNDED PRECEDING
    ) - 1 AS label
FROM grp_test
ORDER BY id, start_time;

逻辑说明 #

1. PARTITION BY id：确保每个id独立计算分组
2. ORDER BY start_time：保证按时间顺序处理行
3. COUNT(CASE WHEN grp_start = 1 THEN 1 END)：统计从分区开头到当前行的所有分组起始次数
4. 减1操作将标签从0开始编号，与期望结果完全匹配

性能说明 #

该方案基于Snowflake原生优化的窗口函数实现，属于流式逐行计算，无需全表聚合或复杂关联操作，能够高效支撑大数据量场景的处理需求。

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