BigQuery数组列加权后按列分组的聚合实现方法

问题描述 #

数据表table1包含column1、weight、vector三列，需完成：

1. 每行vector数组与该行weight逐元素相乘，生成加权向量
2. 按column1分组，对组内加权向量的对应位置元素做聚合（示例为取平均值）

实现代码 #

WITH weighted_elements AS (
    -- 生成带位置索引的加权元素
    SELECT
        column1,
        offset,
        dim * weight AS weighted_val
    FROM table1,
         UNNEST(vector) dim WITH OFFSET AS offset
)
-- 分组聚合后重组数组
SELECT
    column1,
    ARRAY_AGG(AVG(weighted_val) ORDER BY offset) AS weighted_vector
FROM weighted_elements
GROUP BY column1

示例匹配说明 #

你的示例中，分组A的两个加权向量(4,6,8,2,1)和(2,4,6,10,7)，通过对应位置元素取平均得到(3,5,7,6,4)，上述代码中的AVG(weighted_val)正好实现该逻辑。若需改为求和，替换为SUM(weighted_val)即可。

核心逻辑 #

• UNNEST(...) WITH OFFSET拆分数组时保留元素的位置索引，确保聚合时不会打乱元素顺序
• ARRAY_AGG(聚合值 ORDER BY offset)将分组聚合后的元素按原位置重新组合成数组，保证输出数组结构与原数组一致

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