Repository: dathere/qsv (language: rust)

Issue:
feat: `luau` add accumulate helper function
It is a generalized loop where the value of the prior iteration is fed back into the current iteration.  That is, it is basically a loop iterating over the elements of a vector x such that the output vector y is recursively given by `y[i] = f(y[i-1], x[i])`. To start off `y[1]` is `x[1]` if `init` is missing and `f(init, x[1])` if `init` is present.  

For example, if the function passed to it is addition and `init` is missing or equal to 0 then then `y[i] = x[1] + x[2] + ... + x[i]` . 

The lua Functional module may have some relevant code as well.

Here is some lua code and some tests.

```
-- x is input vector, f is function of two args, init is optional
-- output is vector same length as x
function accumulate(x, f, init) do
  y = {}
  for i = 1, #x do
    if (i == 1) then
      if (init) then y[1] = f(init, x[1]) else y[1] = x[1] end
    else
      y[i] = f(y[i-1], x[i])
    end
  end
  return y end
end

-- tests

print("cumulative sum")
x = {3, 5, 4, 1}
add = function(x, y) do return x + y end end
y = accumulate(x, add)
for i = 1, #y do print(x[i], y[i]) end

print("\ncumulative sum with reset once sum exceeds 7")
add2 = function(x, y) do 
  if x > 7 then z = y else z = x+y end
  return z end
end
y = accumulate(x, add2)
for i = 1, #y do print(x[i], y[i]) end
```

giving

```
cumulative sum
3       3
5       8
4       12
1       13

cumulative sum with reset once sum exceeds 7
3       3
5       8
4       4  <-- note how it reset here since 8 > 7
1       5	
```

_Originally posted by @ggrothendieck in https://github.com/dathere/qsv/discussions/2526#discussioncomment-12229670_


Discussion / hints:

Fix patch:
diff --git a/src/cmd/luau.rs b/src/cmd/luau.rs
index 5bba954e0..cd1ae7f68 100644
--- a/src/cmd/luau.rs
+++ b/src/cmd/luau.rs
@@ -2439,6 +2439,85 @@ fn setup_helpers(
     })?;
     luau.globals().set("qsv_cumall", qsv_cumall)?;
 
+    // this is a helper function that implements accumulate functionality.
+    // It takes a column, a function of two arguments, and an optional initial value.
+    // If y = accumulate(x, f, init) then y[1] = init and for i > 1 we have y[i] = f(y[i-1], x[i])
+    // The default for init is x[1].
+    //
+    //   qsv_accumulate(column: string, func: function, init?: any)
+    //          column: the name of the column to accumulate over
+    //            func: function that takes two arguments (prev accumulated value, current val)
+    //            init: optional initial value. If not provided, first row of column is used
+    //         returns: table containing accumulated values
+    //                  or Luau runtime error if invalid arguments
+    //
+    let qsv_accumulate = luau.create_function(|luau, args: mlua::MultiValue| {
+        // Convert args to Vec for easier handling
+        let args: Vec<mlua::Value> = args.into_iter().collect();
+
+        if args.len() < 2 {
+            return helper_err!(
+                "qsv_accumulate",
+                "requires at least 2 arguments: column name and function"
+            );
+        }
+
+        // Get the column name
+        let column_name = args[0].to_string()?;
+
+        // Get the current value from the column
+        let curr_value = luau
+            .globals()
+            .get::<String>(&*column_name)?
+            .parse::<f64>()
+            .unwrap_or(0.0);
+
+        // Get the function
+        let func = match &args[1] {
+            mlua::Value::Function(f) => f.clone(),
+            _ => {
+                return helper_err!("qsv_accumulate", "second argument must be a function");
+            },
+        };
+
+        // Generate unique name for the accumulator state
+        let state_name = format!("_qsv_accumulate_state_{column_name}");
+
+        // Get existing accumulator value or use initial value
+        let prev_acc = if let Ok(prev) = luau.globals().get::<f64>(&*state_name) {
+            prev
+        } else {
+            // Get initial value from args or default to 0.0
+            let init_value = if args.len() > 2 {
+                match &args[2] {
+                    mlua::Value::Number(n) => *n,
+                    mlua::Value::Integer(i) => *i as f64,
+                    mlua::Value::String(s) => s.to_string_lossy().parse::<f64>().unwrap_or(0.0),
+                    _ => 0.0,
+                }
+            } else {
+                0.0 // Always start with 0.0 as default initial value
+            };
+            luau.globals().set(&*state_name, init_value)?;
+            init_value
+        };
+
+        // Call the accumulator function
+        let result = match func.call::<mlua::Value>((prev_acc, curr_value)) {
+            Ok(mlua::Value::Number(n)) => n,
+            Ok(mlua::Value::Integer(i)) => i as f64,
+            Ok(mlua::Value::String(s)) => s.to_string_lossy().parse::<f64>().unwrap_or(prev_acc),
+            Ok(_) => prev_acc,
+            Err(e) => return Err(e),
+        };
+
+        // Store the new accumulated value
+        luau.globals().set(&*state_name, result)?;
+
+        Ok(result)
+    })?;
+    luau.globals().set("qsv_accumulate", qsv_accumulate)?;
+
     // qsv_diff - returns difference between current and previous value
     //
     //   qsv_diff(name, value[, periods])


Test patch:
diff --git a/tests/test_luau.rs b/tests/test_luau.rs
index 5a1c30c10..ee9d4fb84 100644
--- a/tests/test_luau.rs
+++ b/tests/test_luau.rs
@@ -3057,3 +3057,156 @@ return "ok"
     assert_eq!(json["empty_table"], json!({}));
     assert_eq!(json["nested_empty"]["empty"], json!({}));
 }
+
+#[test]
+fn luau_accumulate_custom_sum_function() {
+    let wrk = Workdir::new("luau_accumulate_custom_sum_function");
+    wrk.create(
+        "data.csv",
+        vec![
+            svec!["value"],
+            svec!["10"],
+            svec!["20"],
+            svec!["30"],
+            svec!["40"],
+            svec!["50"],
+        ],
+    );
+
+    let mut cmd = wrk.command("luau");
+    cmd.arg("map")
+        .arg("accumulated")
+        .arg(
+            r#"
+BEGIN {
+
+    -- Define a custom accumulator function that keeps a weighted sum
+    -- where each new value is weighted by its position
+    function udf_sum(acc, x)
+         return acc + tonumber(x)
+    end
+}!
+
+-- This is the MAIN LOOP
+accumulated = qsv_accumulate("value", udf_sum, 0)
+
+-- return the accumulated value for the current row
+return accumulated
+"#,
+        )
+        .arg("data.csv");
+
+    let got: Vec<Vec<String>> = wrk.read_stdout(&mut cmd);
+    let expected = vec![
+        svec!["value", "accumulated"],
+        svec!["10", "10"],  // initial value
+        svec!["20", "30"],  // 10 + 20
+        svec!["30", "60"],  // 30 + 30
+        svec!["40", "100"], // 60 + 40
+        svec!["50", "150"], // 100 + 50
+    ];
+    assert_eq!(got, expected);
+}
+
+#[test]
+fn luau_accumulate_custom_sum_function_with_initial_value() {
+    let wrk = Workdir::new("luau_accumulate_custom_sum_function_with_initial_value");
+    wrk.create(
+        "data.csv",
+        vec![
+            svec!["value"],
+            svec!["10"],
+            svec!["20"],
+            svec!["30"],
+            svec!["40"],
+            svec!["50"],
+        ],
+    );
+
+    let mut cmd = wrk.command("luau");
+    cmd.arg("map")
+        .arg("accumulated")
+        .arg(
+            r#"
+BEGIN {
+
+    -- Define a custom accumulator function that keeps a weighted sum
+    -- where each new value is weighted by its position
+    function udf_sum(acc, x)
+         return acc + x
+    end
+}!
+
+-- This is the MAIN LOOP
+accumulated = qsv_accumulate("value", udf_sum, 100)
+
+-- return the accumulated value for the current row
+return accumulated
+"#,
+        )
+        .arg("data.csv");
+
+    let got: Vec<Vec<String>> = wrk.read_stdout(&mut cmd);
+    let expected = vec![
+        svec!["value", "accumulated"],
+        svec!["10", "110"], // initial value
+        svec!["20", "130"], // 110 + 20
+        svec!["30", "160"], // 130 + 30
+        svec!["40", "200"], // 160 + 40
+        svec!["50", "250"], // 200 + 50
+    ];
+    assert_eq!(got, expected);
+}
+
+#[test]
+fn luau_accumulate_custom_function_with_reset() {
+    let wrk = Workdir::new("luau_accumulate_custom_with_reset");
+    wrk.create(
+        "data.csv",
+        vec![
+            svec!["value"],
+            svec!["3"],
+            svec!["5"],
+            svec!["4"],
+            svec!["1"],
+        ],
+    );
+
+    let mut cmd = wrk.command("luau");
+    cmd.arg("map")
+        .arg("accumulated")
+        .arg(
+            r#"
+BEGIN {
+
+    -- Define a custom accumulator function that keeps a weighted sum
+    -- where each new value is weighted by its position
+    function func_with_reset(acc, x)
+         if tonumber(acc) > 7 then
+            z = x
+         else
+            z = acc + x
+         end
+         return z
+    end
+}!
+
+-- This is the MAIN LOOP
+accumulated = qsv_accumulate("value", func_with_reset, 0)
+
+-- return the accumulated value for the current row
+return accumulated
+"#,
+        )
+        .arg("data.csv");
+
+    let got: Vec<Vec<String>> = wrk.read_stdout(&mut cmd);
+    let expected = vec![
+        svec!["value", "accumulated"],
+        svec!["3", "3"], // initial value
+        svec!["5", "8"], // 3 + 5
+        svec!["4", "4"], // reset since 8 > 7
+        svec!["1", "5"], // 4 + 1
+    ];
+    assert_eq!(got, expected);
+}