Repository: pola-rs/polars-xdt (language: python)

Issue:
Add `ignore_nulls` (and `ignore_nan`) option to `xdt.ewma_by_time`
Currently, if there are any NaN values in the value column passed to `xdt.ewma_by_time`, then all following values in the output are NaN (see snippet). It would be great if there was n `ignore_nulls` flag, similar to in the builtin `ewma`, to allow for NaN or null values to be ignored during calculation, to prevent this. In this case, the presence or absence of a row containing Null or NaN should have no effect on subsequent rows; i.e. the `ewma`-ed output of the final row of the two following tables should be identical.

```
shape: (2, 2)
┌───────────┬────────────────────────────┐
│ values    ┆ time                       │
│ ---       ┆ ---                        │
│ f64       ┆ datetime[ns]               │
╞═══════════╪════════════════════════════╡
│ -0.042898 ┆ 2000-01-01 00:00:00        │
│ 0.186466  ┆ 2000-01-01 00:00:00.000002 │
└───────────┴────────────────────────────┘
```

```
shape: (3, 2)
┌───────────┬────────────────────────────┐
│ values    ┆ time                       │
│ ---       ┆ ---                        │
│ f64       ┆ datetime[ns]               │
╞═══════════╪════════════════════════════╡
│ -0.042898 ┆ 2000-01-01 00:00:00        │
│ NaN       ┆ 2000-01-01 00:00:00.000001 │
│ 0.186466  ┆ 2000-01-01 00:00:00.000002 │
└───────────┴────────────────────────────┘
```

## Reproducible snippet
```
from datetime import timedelta

import numpy as np
import polars as pl
import polars_xdt as xdt


n = 100

df = pl.DataFrame({
    "values": np.linspace(0, 10, n) + 0.1 * np.random.normal(size=n),
    "time": np.datetime64("2000-01-01 00:00:00") + np.asarray([i*np.timedelta64(1000, "ns") for i in range(n)])
})


new = df.with_columns(xdt.ewma_by_time("values", times="time", half_life=timedelta(microseconds=1)).alias("ewma"))

# True
print(new["ewma"].is_finite().all())

new_with_nan = df.with_columns(xdt.ewma_by_time(
    pl.when(pl.col("values") > 5).then(np.nan).otherwise(pl.col("values")), times="time", half_life=timedelta(microseconds=1)
).alias("ewma"))

# False
print(new_with_nan["ewma"].is_finite().all())

```


Discussion / hints:
thanks @wbeardall for the request! seems reasonable, will take a look

Fix patch:
diff --git a/src/ewma_by_time.rs b/src/ewma_by_time.rs
index 1f50fcb..fc242fd 100644
--- a/src/ewma_by_time.rs
+++ b/src/ewma_by_time.rs
@@ -5,7 +5,7 @@ pub(crate) fn impl_ewma_by_time_float(
     times: &Int64Chunked,
     values: &Float64Chunked,
     half_life: i64,
-    time_unit: TimeUnit,
+    time_unit: TimeUnit
 ) -> Float64Chunked {
     let mut out = Vec::with_capacity(times.len());
     if values.is_empty() {
@@ -18,13 +18,27 @@ pub(crate) fn impl_ewma_by_time_float(
         TimeUnit::Nanoseconds => half_life * 1_000,
     };
 
-    let mut prev_time: i64 = times.get(0).unwrap();
-    let mut prev_result = values.get(0).unwrap();
-    out.push(Some(prev_result));
+    let mut skip_rows: usize = 0;
+    let mut prev_time: i64 = 0;
+    let mut prev_result: f64 = 0.;
+    for (idx, (value, time)) in values.iter().zip(times.iter()).enumerate() {
+        match (time, value) {
+            (Some(time), Some(value)) => {
+                prev_time = time;
+                prev_result = value;
+                out.push(Some(prev_result));
+                skip_rows = idx + 1;
+                break;
+            },
+            _ => {
+                out.push(None);
+            }
+        };
+    }
     values
         .iter()
         .zip(times.iter())
-        .skip(1)
+        .skip(skip_rows)
         .for_each(|(value, time)| {
             match (time, value) {
                 (Some(time), Some(value)) => {
@@ -36,7 +50,7 @@ pub(crate) fn impl_ewma_by_time_float(
                     prev_time = time;
                     prev_result = result;
                     out.push(Some(result));
-                }
+                },
                 _ => out.push(None),
             }
         });
diff --git a/src/expressions.rs b/src/expressions.rs
index ab71e33..fb096e0 100644
--- a/src/expressions.rs
+++ b/src/expressions.rs
@@ -178,7 +178,7 @@ fn arg_previous_greater(inputs: &[Series]) -> PolarsResult<Series> {
 
 #[derive(Deserialize)]
 struct EwmTimeKwargs {
-    half_life: i64,
+    half_life: i64
 }
 
 #[polars_expr(output_type=Float64)]


Test patch:
diff --git a/tests/test_ewma_by_time.py b/tests/test_ewma_by_time.py
index 08a65a7..55a36cd 100644
--- a/tests/test_ewma_by_time.py
+++ b/tests/test_ewma_by_time.py
@@ -1,34 +1,48 @@
 import polars as pl
 from polars.testing import assert_frame_equal
 import polars_xdt as xdt
+import pytest
 from datetime import datetime, timedelta
 
+import os
 
-def test_ewma_by_time():
+os.environ["POLARS_VERBOSE"] = "1"
+
+@pytest.mark.parametrize("start_null", [True, False])
+def test_ewma_by_time(start_null):
+    if start_null:
+        values = [None]
+        times = [datetime(2020, 1, 1)]
+    else:
+        values = []
+        times = []
+        
     df = pl.DataFrame(
         {
-            "values": [0.0, 1, 2, None, 4],
-            "times": [
-                datetime(2020, 1, 1),
-                datetime(2020, 1, 3),
-                datetime(2020, 1, 10),
-                datetime(2020, 1, 15),
-                datetime(2020, 1, 17),
+            "values": values + [0.0, 1., 2., None, 4.],
+            "times": times + [
+                datetime(2020, 1, 2),
+                datetime(2020, 1, 4),
+                datetime(2020, 1, 11),
+                datetime(2020, 1, 16),
+                datetime(2020, 1, 18),
             ],
         }
     )
     result = df.select(
-        xdt.ewma_by_time("values", times="times", half_life=timedelta(days=4)),
+        xdt.ewma_by_time("values", times="times", half_life=timedelta(days=2)),
     )
+        
     expected = pl.DataFrame(
         {
-            "values": [
+            "values": values + [
                 0.0,
-                0.2928932188134524,
-                1.4924741174358913,
+                0.5,
+                1.8674174785275222,
                 None,
-                3.2545080948503213,
+                3.811504554703363,
             ]
         }
     )
+
     assert_frame_equal(result, expected)