如何在Presto SQL中基于offset列的递增子序列对表分区？

问题描述 #

存在表T，结构及数据如下：

Table T:
row_num, offset ,col_1, col_2, col_3 ....
     1       0
     2      10
     3      22
     4      32
     5       0
     6       0
     7      12
     8      32 
     9      44

需要基于offset列将表分区，每个严格递增子序列的行归为一个分区，预期分区结果为[1,2,3,4]、[5]、[6,7,8,9]。此前尝试过LEAD和LAG窗口函数，但误以为需要循环功能，想了解SQL中具体实现方式。

解决方案 #

无需循环，仅通过窗口函数即可实现核心逻辑。核心思路是：先标记每个分区的起始行，再通过累加起始标记生成唯一的分区编号，最终按分区编号聚合。

完整SQL代码 #

WITH partition_markers AS (
    SELECT 
        row_num,
        offset,
        -- 标记分区起始行：offset为0，或当前offset不大于前一行offset（严格递增中断）
        CASE 
            WHEN offset = 0 THEN 1
            WHEN offset <= LAG(offset) OVER (ORDER BY row_num) THEN 1
            ELSE 0
        END AS is_start
    FROM T
),
partitioned_data AS (
    SELECT 
        row_num,
        offset,
        -- 累加起始标记，生成连续的分区编号
        SUM(is_start) OVER (ORDER BY row_num) AS partition_id
    FROM partition_markers
)
-- 按分区编号聚合，得到每个分区的行号集合
SELECT 
    partition_id,
    ARRAY_AGG(row_num ORDER BY row_num) AS partition_rows
FROM partitioned_data
GROUP BY partition_id
ORDER BY partition_id;

逻辑说明 #

1. 标记分区起始行：

   • 使用LAG(offset) OVER (ORDER BY row_num)获取当前行的前一行offset值
   • 当offset为0时，必然是新分区的起点；当当前offset不大于前一行offset时，说明严格递增序列中断，同样是新分区的起点，此时标记is_start=1，否则为0

2. 生成分区编号：

   • 通过SUM(is_start) OVER (ORDER BY row_num)按行号顺序累加起始标记，每遇到一个起始行，分区编号自动加1，同一分区内的行将拥有相同的partition_id

3. 聚合分区结果：

   • 使用ARRAY_AGG函数将同一partition_id下的row_num聚合为数组，得到预期的分区行号集合

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