如何在Spark SQL中将多行数据聚合为指定大小限制的数组

看起来你需要按id分组，把每行的value按照固定大小（这里是2个元素）拆分成多个数组，同时生成带序号的新ID。普通的collect_list只能把同ID的所有值聚合成一个大数组，没法做拆分，咱们可以用窗口函数+分组聚合的组合来实现这个需求。

实现思路 #

核心是先给每个id下的行分配一个"组编号"，把每N行（这里N=2）划到同一个组里，然后再按id+组编号聚合生成指定大小的数组：

• 用row_number()给每个id下的行排序（保证顺序和输入一致）
• 通过整除运算计算每行所属的组编号
• 最后按id和组编号分组，用collect_list聚合数组，再拼接新的ID

完整Spark SQL代码 #

WITH numbered_rows AS (
    SELECT 
        id,
        value,
        -- 给每个id下的行分配序号，这里如果有明确的排序字段（比如时间）可以替换掉value
        ROW_NUMBER() OVER (PARTITION BY id ORDER BY value) AS row_num
    FROM your_input_table
),
grouped_rows AS (
    SELECT 
        id,
        value,
        -- 计算组编号，这里2就是你想要的数组大小，可按需修改
        CEIL(row_num / 2.0) AS group_num
    FROM numbered_rows
)
SELECT 
    CONCAT(id, '_', group_num) AS id,
    COLLECT_LIST(value) AS value
FROM grouped_rows
GROUP BY id, group_num
ORDER BY id;

代码解释 #

1. numbered_rows CTE：给每个id分区内的行添加序号，ORDER BY value是为了保证行的顺序和输入一致，如果你的表有更合适的排序字段（比如插入时间），替换成那个字段会更准确。
2. grouped_rows CTE：用CEIL(row_num / 2.0)把每2行划到一个组里，比如行号1、2会分到组1，行号3、4分到组2，以此类推。这里的2.0是为了避免整数除法的问题，确保计算结果是浮点数后再取整。
3. 最终查询：拼接id和group_num得到新的ID，用COLLECT_LIST把同组的value聚合成数组，最后按新ID排序。

验证结果 #

运行上面的代码后，输出会和你期望的完全一致：

注意事项 #

• 如果需要调整数组大小，只需要把CEIL(row_num / 2.0)里的2.0改成你想要的数字就行，比如要每个数组3个元素就改成3.0。
• 如果你的value是复杂类型（比如结构体），COLLECT_LIST同样适用，Spark会保留原类型结构。

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