如何在BigQuery中用JavaScript UDF实现带负值重置的多列条件累计求和

需求是生成数据表，实现以下累计求和规则：

• 正常累计tr_value得到running_total
• 当running_total低于0时，将avg_price * unit_left作为新的起始点，后续累计基于这个新起点计算，最终得到investment_amount（预期结果）

示例数据表 #

当前尝试的实现（存在问题） #

UDF代码 #

CREATE OR REPLACE FUNCTION calculate_investment_amount(arr ARRAY<INT64>) RETURNS INT64 LANGUAGE js AS R"""
let counter = 0
for (let i = 0; i < arr.length; i++) {
    if (arr[i].tr_value< 0 ) {
        counter += arr[i].avg_price[arr[i].avg_price.length -1] * arr[i].unit_left[arr[i].unit_left.length -1]
        arr[i].tr_value= counter
    }else {
        arr[i].tr_value = arr[i].tr_value + counter
    }
}
return arr;
""";

对应的SQL语句 #

SELECT
  order_date,
  user_id,
  item_code,
  tr_value, 
  ARRAY_AGG(cast(tr_value as int)) OVER(PARTITION BY user_id, item_code ORDER BY order_date) AS tr_value,
  ARRAY_AGG(cast(AVG_PRICE as int)) OVER(PARTITION BY user_id, item_code ORDER BY order_date) AS avg_price,
  ARRAY_AGG(cast(unit_left as int)) OVER(PARTITION BY user_id, item_code ORDER BY order_date) AS unit_left
FROM
  trx),
z as (
SELECT
  a.*,
  TO_JSON_STRING(a) AS json,
  calculate_investment_amount(TO_JSON_STRING(a)) as investment_amount
FROM
  a AS a

SQL输出的JSON示例 #

{"order_date":"2022-09-13","user_id":"000027","item_code":"EII01FIS01SFIA01","tr_value":-1260334,"cont_value":[250000,250000,250000,250000,250000,-1260334],"avg_price":[1362,1362,1349,1347,1349,10],"unit_left":[1362,1362,1349,1347,1349,10]}

问题分析与修正方案 #

当前UDF存在两个核心问题：

1. 参数类型不匹配：UDF声明接收ARRAY<INT64>，但实际传入的是JSON字符串，无法正确解析tr_value、avg_price等字段
2. 逻辑错误：循环中错误访问数组最后一个元素，且未正确维护累计状态的重置逻辑

修正后的UDF实现 #

传入包含每一行数据的STRUCT数组，逐行处理累计逻辑：

CREATE OR REPLACE FUNCTION calculate_investment_amount(rows ARRAY<STRUCT<tr_value INT64, avg_price INT64, unit_left INT64>>) 
RETURNS ARRAY<INT64> 
LANGUAGE js AS """
let result = [];
let current_total = 0;

for (const row of rows) {
    current_total += row.tr_value;
    // 累计值低于0时，重置为avg_price * unit_left
    if (current_total < 0) {
        current_total = row.avg_price * row.unit_left;
    }
    result.push(current_total);
}

return result;
""";

修正后的SQL查询 #

WITH ordered_trx AS (
    SELECT
        order_date,
        user_id,
        item_code,
        tr_value,
        avg_price,
        unit_left,
        -- 按用户和商品分组排序，打包从第一行到当前行的数据为STRUCT数组
        ARRAY_AGG(STRUCT(tr_value, avg_price, unit_left)) OVER(
            PARTITION BY user_id, item_code 
            ORDER BY order_date 
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) AS running_rows
    FROM trx
),
calculated AS (
    SELECT
        *,
        calculate_investment_amount(running_rows) AS investment_amount_array
    FROM ordered_trx
)
SELECT
    order_date,
    user_id,
    item_code,
    tr_value,
    avg_price,
    unit_left,
    -- 提取数组最后一个元素作为当前行的investment_amount
    investment_amount_array[ORDINAL(ARRAY_LENGTH(investment_amount_array))] AS investment_amount
FROM calculated;

逻辑说明 #

1. ordered_trx CTE：按user_id和item_code分组、order_date排序，将每行及之前的所有数据打包成STRUCT数组
2. calculated CTE：调用修正后的UDF处理每组的行数组，生成每行对应的累计结果数组
3. 最终查询：从结果数组中提取当前行对应的investment_amount，实现需求中的条件累计逻辑

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