BigQuery UDF使用ARRAY_AGG报错的解决方案咨询

问题描述 #

我尝试在BigQuery中创建UDF将多行数据压缩为单行，编写的函数代码如下：

CREATE OR REPLACE FUNCTION function_name(to_compress_column INT64, order_by_column INT64) AS (
    TO_JSON_STRING(
           ARRAY_AGG(
                IFNULL(to_compress_column,-1) RESPECT NULLS
                ORDER BY order_by_column
                    )
                  )
    );

执行时触发错误：ARRAY_AGG not allowed in SQL function body。我试过用SELECT和UNNEST，但输入不是数组，没能解决问题。

输入表结构及数据 #

期望输出 #

我已经知道用以下SQL可以实现需求：

SELECT TO_JSON_STRING(ARRAY_AGG(IFNULL(visits,-1) RESPECT NULLS
  ORDER BY hour)) visits, 
  place_id, 
  date 
  from input_table
  group by place_id, date

但因为要处理大量表，不想重复写这段聚合逻辑，希望封装成UDF后用简化的查询：

SELECT function_name(visits,hour) visits, 
      place_id, 
      date 
      from input_table
      group by place_id, date

请问怎么解决UDF里用ARRAY_AGG报错的问题？

原因分析 #

BigQuery的标量SQL UDF（即你编写的这种单行函数）不支持聚合函数（如ARRAY_AGG），因为标量UDF的设计目标是处理单条行的输入，而聚合函数是对分组后的多行数据进行计算，两者执行逻辑不兼容。

解决方法 #

方法1：使用表值函数（TVF） #

表值函数可以处理分组后的数据集，返回结构化结果。你可以创建一个接受分组行数组作为输入的TVF：

CREATE OR REPLACE FUNCTION my_dataset.aggregate_and_jsonify(
  input_data ANY TYPE
)
RETURNS STRING
AS ((
  SELECT TO_JSON_STRING(ARRAY_AGG(IFNULL(visits_col, -1) RESPECT NULLS ORDER BY hour_col))
  FROM UNNEST(input_data) AS row
));

使用时需要先将分组内的行打包成数组传入：

SELECT 
  my_dataset.aggregate_and_jsonify(ARRAY_AGG(STRUCT(visits AS visits_col, hour AS hour_col))) AS visits,
  place_id,
  date
FROM input_table
GROUP BY place_id, date;

方法2：使用BigQuery宏（推荐） #

宏是更适配该场景的方案，它会在查询解析阶段直接替换为对应SQL代码，完全复用聚合逻辑，写法和你期望的简化查询几乎一致：

CREATE OR REPLACE MACRO my_dataset.aggregate_visits(col_to_compress, order_col) AS (
  TO_JSON_STRING(ARRAY_AGG(IFNULL(col_to_compress, -1) RESPECT NULLS ORDER BY order_col))
);

使用宏的查询语句完全符合你的需求：

SELECT 
  my_dataset.aggregate_visits(visits, hour) AS visits,
  place_id,
  date
FROM input_table
GROUP BY place_id, date;

宏的优势在于仅做代码替换，无额外执行开销，且保留原SQL的执行计划，适合批量处理大量表的场景。

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