simplify cum_sum closes #15

docs/cum_sum.md

@@ -1,4 +1,4 @@
-# 4. How to not FOLD under pressure
+# 4. Yes we SCAN
 
 The operations we've seen so far have all been elementwise, e.g.:
 
@@ -40,19 +40,14 @@ fn cum_sum(inputs: &[Series]) -> PolarsResult<Series> {
     let ca: &Int64Chunked = s.i64()?;
     let out: Int64Chunked = ca
         .into_iter()
-        .scan(None, |state: &mut Option<i64>, x: Option<i64>| {
-            let sum = match (*state, x) {
-                (Some(state_inner), Some(x)) => {
-                    *state = Some(state_inner + x);
-                    *state
-                }
-                (None, Some(x)) => {
-                    *state = Some(x);
-                    *state
-                }
-                (_, None) => None,
-            };
-            Some(sum)
+        .scan(0_i64, |state: &mut i64, x: Option<i64>| {
+            match x {
+                Some(x) => {
+                    *state += x;
+                    Some(Some(*state))
+                },
+                None => Some(None),
+            }
         })
         .collect_trusted();
     let out: Int64Chunked = out.with_name(ca.name());

docs/noop.md

@@ -42,12 +42,12 @@ Let's go through this line-by-line:
 
 - when we compile Rust, it generates a Shared Object file.
   The `lib` variable holds its filepath;
-- We'll cover `is_elementwise` in [How to not FOLD under pressure], for now don't pay attention to it;
+- We'll cover `is_elementwise` in [Yes we SCAN], for now don't pay attention to it;
 - We use the utility function `parse_into_expr` to make sure that
   literals will be parsed as column names - this will ensure we'll be
   able to call either `noop('a')` or `noop(pl.col('a'))`.
 
-  [How to not FOLD under pressure]: ../cum_sum/
+  [Yes we SCAN]: ../cum_sum/
 
 ## Let's get Rusty
 

run.py

@@ -7,9 +7,6 @@
         "b": [2.0, 3.1, 2.5],
         "c": ["3", "7", "3"],
     }
-).select(abc=pl.struct("a", "b", "c"))
-print(df.with_columns(abc_shifted=mp.shift_struct("abc")))
-import pprint
-pprint.pprint(df.with_columns(abc_shifted=mp.shift_struct("abc")).schema)
-# print(df.lazy().with_columns(swapped= mp.shift_struct('a')).schema)
-# print(df.lazy().with_columns(swapped= mp.shift_struct('a')).collect().schema)
+)
+
+print(df.with_columns(mp.cum_sum('a')))

src/expressions.rs

@@ -76,19 +76,14 @@ fn cum_sum(inputs: &[Series]) -> PolarsResult<Series> {
     let ca: &Int64Chunked = s.i64()?;
     let out: Int64Chunked = ca
         .into_iter()
-        .scan(None, |state: &mut Option<i64>, x: Option<i64>| {
-            let sum = match (*state, x) {
-                (Some(state_inner), Some(x)) => {
-                    *state = Some(state_inner + x);
-                    *state
-                }
-                (None, Some(x)) => {
-                    *state = Some(x);
-                    *state
-                }
-                (_, None) => None,
-            };
-            Some(sum)
+        .scan(0_i64, |state: &mut i64, x: Option<i64>| {
+            match x {
+                Some(x) => {
+                    *state += x;
+                    Some(Some(*state))
+                },
+                None => Some(None),
+            }
         })
         .collect_trusted();
     let out: Int64Chunked = out.with_name(ca.name());