How to make running total using multiple columns condition when value below zero in BigQuery using Javascript UDF?

Question

i am trying to make a table where the running total have a condition whenever the value went below 0 it will use the value of other columns (avg_price * unit_left) and work as a new starting point like the table below:

| type | unit | price | avg_price | tr_value | unit_left | running_total | investment_amount (expected)
| Buy  | 10   | 10    | 10        | 100      | 10        | 100           | 100
| Buy  | 15   | 14    | 12,4      | 210      | 25        | 310           | 310
| Sell | 20   | 18    | 12,4      | -360     | 5         | 20            | 20
| Buy  | 15   | 16    | 13        | 240      | 20        | 260           | 260
| Sell | 10   | 20    | 13        | -200     | 10        | 60            | 60
| Sell | 8    | 25    | 13        | -200     | 2         | -140          | 26
| Buy  | 10   | 20    | 18,3      | 200      | 12        | 60            | 226

I've tried to use UDF but havent figured out how to done it with multiple columns.

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;
""";

And this is the 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

JSON from SQL:

{"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]}

Answer 1

From what it looks like all you want to know which column value you should used for calculation. And you have a condition for this usage - where if some column value is negative you don't want to use that value and move on to whatever useful column next.

If above assumption is correctly understood, then see if something like below helps you understand:

select a,b,c, (case when c < 0 then (case when b < 0 then a else b end) else c end) as value_to_be_used_for_calculation, 
from 
(select 1 as a, -1 as b, 2 as c 
union all 
select 10 as a, 20 as b, -2 as c 
union all 
select 15 as a, -4 as b, -3 as c 
)

Here the case conditions are checking if C is positive then use C, if its negative then check on B if positive then use B else use A...

Once you know what final value to use from various available columns from the given row then you can leverage it into any custom function you may have.