如何在Snowflake中对数组执行聚合，按数字长度拆分输出多列数组？

我有如下Snowflake数组格式的输入数据，需要遍历数组中的每个值，根据值的数字长度进行聚合，将5位数字的值放入一列数组、6位数字的值放入另一列数组，输出多列数组结果。

输入数据 #

ID_COL,ARRAY_COL_VALUE
1,[22,333,666666]
2,[1,55555,999999999]
3,[22,444]

期望输出表 #

ID_COL,FIVE_DIGIT_COL,SIX_DIGIT_COL
1,[],[666666]
2,[55555],[]
3,[],[]

请问能否通过遍历数组元素执行SQL聚合来实现该需求？若SQL无法实现，使用JavaScript或Python编写UDF的方案也可。

纯SQL实现方案 #

完全可以用Snowflake原生SQL实现，无需依赖UDF。核心逻辑是先拆分数组为单行元素，按数字长度筛选后再聚合回数组：

WITH split_data AS (
    SELECT 
        ID_COL,
        value AS num,
        LENGTH(TO_VARCHAR(num)) AS digit_length
    FROM your_table,
    LATERAL FLATTEN(input => ARRAY_COL_VALUE)
)
SELECT 
    ID_COL,
    ARRAY_AGG(CASE WHEN digit_length = 5 THEN num END IGNORE NULLS) AS FIVE_DIGIT_COL,
    ARRAY_AGG(CASE WHEN digit_length = 6 THEN num END IGNORE NULLS) AS SIX_DIGIT_COL
FROM split_data
GROUP BY ID_COL
ORDER BY ID_COL;

关键步骤说明 #

• LATERAL FLATTEN：将数组拆分为每行一个独立元素，实现逐个遍历处理
• LENGTH(TO_VARCHAR(num))：把数字转为字符串后计算长度，精准判断位数
• ARRAY_AGG(...) IGNORE NULLS：仅聚合符合条件的元素，自动忽略不符合条件产生的NULL值，最终生成目标数组
• 按ID_COL分组聚合，得到每个ID对应的两个目标数组

JavaScript UDF方案 #

如果偏好UDF实现，可编写JavaScript函数直接处理数组：

CREATE OR REPLACE FUNCTION CLASSIFY_DIGITS(arr ARRAY)
RETURNS OBJECT
LANGUAGE JAVASCRIPT
AS $$
    const fiveDigit = [];
    const sixDigit = [];
    arr.forEach(num => {
        const len = num.toString().length;
        if (len === 5) fiveDigit.push(num);
        else if (len === 6) sixDigit.push(num);
    });
    return { FIVE_DIGIT_COL: fiveDigit, SIX_DIGIT_COL: sixDigit };
$$;

调用方式：

SELECT 
    ID_COL,
    CLASSIFY_DIGITS(ARRAY_COL_VALUE):FIVE_DIGIT_COL AS FIVE_DIGIT_COL,
    CLASSIFY_DIGITS(ARRAY_COL_VALUE):SIX_DIGIT_COL AS SIX_DIGIT_COL
FROM your_table;

Python UDF方案 #

也可以用Python编写UDF，逻辑与JavaScript版本一致：

CREATE OR REPLACE FUNCTION CLASSIFY_DIGITS_PY(arr ARRAY)
RETURNS OBJECT
LANGUAGE PYTHON
RUNTIME_VERSION = '3.8'
HANDLER = 'classify_digits'
AS $$
def classify_digits(arr):
    five_digit = []
    six_digit = []
    for num in arr:
        digit_len = len(str(num))
        if digit_len == 5:
            five_digit.append(num)
        elif digit_len == 6:
            six_digit.append(num)
    return {"FIVE_DIGIT_COL": five_digit, "SIX_DIGIT_COL": six_digit}
$$;

调用方式：

SELECT 
    ID_COL,
    CLASSIFY_DIGITS_PY(ARRAY_COL_VALUE):FIVE_DIGIT_COL AS FIVE_DIGIT_COL,
    CLASSIFY_DIGITS_PY(ARRAY_COL_VALUE):SIX_DIGIT_COL AS SIX_DIGIT_COL
FROM your_table;

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