Snowflake中限制ARRAY_AGG结果条数的实现方案（无需重构主查询）

需求说明：使用Snowflake的ARRAY_AGG聚合函数生成数组时，需要在不重构现有主查询、无法使用用户自定义聚合函数（UDAF）的前提下，实现类似ARRAY_AGG(...) WITHIN GROUP(... LIMIT 3)的条数限制效果，同时避免因聚合生成的数组超过16MB导致ARRAY_SLICE失效的问题。

解决方案思路 #

先通过窗口函数在分组内筛选出指定条数的记录，再对筛选后的结果执行ARRAY_AGG聚合，从根源上避免生成超大数组，同时满足无需重构主查询的要求。

示例实现 #

1. 示例数据准备

CREATE OR REPLACE TABLE tab(grp TEXT, col TEXT) AS
SELECT * FROM VALUES 
  ('Grp1', 'A'),('Grp1', 'B'),('Grp1', 'C'),('Grp1', 'D'), ('Grp1', 'E'),
  ('Grp2', 'X'),('Grp2', 'Y'),('Grp2', 'Z'),('Grp2', 'V'),
  ('Grp3', 'M'),('Grp3', 'N'),('Grp3', 'M');

2. 带条数限制的聚合查询

WITH ranked_data AS (
  SELECT 
    grp, 
    col,
    -- 按分组内的排序规则（可根据需求调整）生成行号
    ROW_NUMBER() OVER (PARTITION BY grp ORDER BY col) AS rn
  FROM tab
  -- 此处可直接替换为你的主查询，无需修改原有逻辑
)
SELECT 
  grp,
  ARRAY_AGG(col) AS arr_limit_3
FROM ranked_data
WHERE rn <= 3  -- 限制每组只取前3条
GROUP BY grp
ORDER BY grp DESC;

3. 预期输出

GRP    ARR_LIMIT_3
Grp3   [   "M",   "M",   "N" ]
Grp2   [   "V",   "X",   "Y" ]
Grp1   [   "A",   "B",   "C" ]

为什么不使用ARRAY_SLICE？ #

如果直接对全量数据执行ARRAY_AGG再用ARRAY_SLICE截取，当分组内数据量极大时，聚合生成的数组可能超过16MB的限制，触发Result array of ARRAYAGG is too large错误。而先筛选再聚合的方式只会处理每组内的指定条数记录，不会生成超大数组，从根本上避免了这个问题。

适配复杂主查询的场景 #

如果你的主查询包含JOIN、过滤等复杂逻辑，只需将主查询替换到ranked_data CTE的子查询中即可，无需修改原有主查询的逻辑：

WITH ranked_data AS (
  SELECT 
    grp, 
    col,
    ROW_NUMBER() OVER (PARTITION BY grp ORDER BY col) AS rn
  FROM (
    -- 放入你的原有主查询
    SELECT t.grp, t.col
    FROM big_table t
    JOIN other_table ot ON t.id = ot.t_id
    WHERE t.status = 'active'
  ) main_query
)
SELECT 
  grp,
  ARRAY_AGG(col) AS arr_limit_3
FROM ranked_data
WHERE rn <= 3
GROUP BY grp;

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