如何在Athena/Presto中用SQL将扁平行转为嵌套映射/结构体数组？

需求说明 #

基于给定的用户和待办测试数据，需要用Athena/Presto的SQL生成指定格式的嵌套JSON：每个用户对象包含user_id、name字段，以及对应的todos数组（每个待办项含todo_id和title）。

测试数据 #

with users (user_id, name) as (
    values (1, 'Alice'),
           (2, 'Bob'),
           (3, 'Charlie')
), todos (todo_id, user_id, title) as (
    values (1, 1, 'todo 1'),
           (2, 1, 'todo 2'),
           (3, 2, 'todo 3'),
           (4, 3, 'todo 4')
)

期望输出JSON #

[
    {
        "user_id": 1,
        "name": "Alice",
        "todos": [
            {
                "todo_id": 1,
                "title": "todo 1"
            },
            {
                "todo_id": 2,
                "title": "todo 2"
            }
        ]
    },
    {
        "user_id": 2,
        "name": "Bob",
        "todos": [
            {
                "todo_id": 3,
                "title": "todo 3"
            }
        ]
    },
    {
        "user_id": 3,
        "name": "Charlie",
        "todos": [
            {
                "todo_id": 4,
                "title": "todo 4"
            }
        ]
    }
]

之前尝试的问题 #

之前用多层map_agg实现的SQL输出不符合预期：把user_id当成了JSON对象的键，缺失name字段，且待办项的字段名错误（显示为id而非todo_id）：

select
    cast(array_agg(res) as JSON) as result
    from
    (
        select map_agg(user_id, m) as res
        from
            (
                select
                    user_id,
                    map_agg('todos', todos) as m
                from
                    (
                        select
                            user_id,
                            array_agg(todo) as todos
                        from
                            (
                                select
                                    user_id,
                                    map_agg('id', todo_id) as todo
                                from
                                    todos
                                group by
                                    user_id,
                                    todo_id
                            ) t
                        group by
                            user_id
                    ) t
                group by
                    user_id
            ) t
        group by
            user_id
    ) t
group by
    true;

输出结果：

[{"2":{"todos":[{"id":3}]}},{"1":{"todos":[{"id":1},{"id":2}]}},{"3":{"todos":[{"id":4}]}}]

正确实现方案 #

Athena/Presto可以原生实现该需求，核心是用row()构造支持多类型字段的结构体，结合array_agg聚合和JSON类型转换即可：

with users (user_id, name) as (
    values (1, 'Alice'),
           (2, 'Bob'),
           (3, 'Charlie')
), todos (todo_id, user_id, title) as (
    values (1, 1, 'todo 1'),
           (2, 1, 'todo 2'),
           (3, 2, 'todo 3'),
           (4, 3, 'todo 4')
)
select cast(array_agg(user_obj) as JSON) as result
from (
    select
        row(
            u.user_id,
            u.name,
            array_agg(row(t.todo_id, t.title)) as todos
        ) as user_obj
    from users u
    left join todos t on u.user_id = t.user_id
    group by u.user_id, u.name
) t;

方案说明 #

1. 关联表数据：用left join关联用户和待办表，确保无待办的用户也能被保留（不需要的话可换成inner join）
2. 聚合待办数组：对每个用户，用array_agg(row(t.todo_id, t.title))生成包含待办ID和标题的结构体数组
3. 构造用户对象：用row()把用户ID、名称、待办数组组合成完整的用户结构体
4. 生成最终JSON：用array_agg聚合所有用户结构体为数组，再转成JSON类型，得到符合需求的嵌套结构

执行上述SQL后，输出结果与期望完全一致。

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