Skip to content

Commit

Permalink
fix: ewm_mean_by was skipping initial nulls when it was already sorte…
Browse files Browse the repository at this point in the history
…d by "by" column (#15812)
  • Loading branch information
MarcoGorelli authored Apr 21, 2024
1 parent d58b912 commit efa92f9
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
4 changes: 3 additions & 1 deletion crates/polars-ops/src/series/ops/ewm_by.rs
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,9 @@ where
out.push(Some(prev_result));
skip_rows = idx + 1;
break;
};
} else {
out.push(None)
}
}
values
.iter()
Expand Down
5 changes: 4 additions & 1 deletion py-polars/tests/unit/operations/test_ewm_by.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@
from polars.type_aliases import PolarsIntegerType, TimeUnit


def test_ewma_by_date() -> None:
@pytest.mark.parametrize("sort", [True, False])
def test_ewma_by_date(sort: bool) -> None:
df = pl.LazyFrame(
{
"values": [3.0, 1.0, 2.0, None, 4.0],
Expand All @@ -25,6 +26,8 @@ def test_ewma_by_date() -> None:
],
}
)
if sort:
df = df.sort("times")
result = df.select(
pl.col("values").ewm_mean_by("times", half_life=timedelta(days=2)),
)
Expand Down

0 comments on commit efa92f9

Please sign in to comment.