fix incorrect match in agg_mean (#3030)

pola-rs · Apr 1, 2022 · fb2c57e · fb2c57e
1 parent 524f8a5
commit fb2c57e
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git a/polars/polars-core/src/frame/groupby/aggregations.rs b/polars/polars-core/src/frame/groupby/aggregations.rs
@@ -598,7 +598,7 @@ where
             GroupsProxy::Slice(groups) => {
                 agg_helper_slice::<Float64Type, _>(groups, |[first, len]| {
                     debug_assert!(len < self.len() as IdxSize);
-                    match first - len {
+                    match len {
                         0 => None,
                         1 => self.get(first as usize).map(|v| NumCast::from(v).unwrap()),
                         _ => {

diff --git a/py-polars/tests/test_datelike.py b/py-polars/tests/test_datelike.py
@@ -666,3 +666,40 @@ def test_rolling_groupby_by_argument() -> None:
     )
 
     assert out.frame_equal(expected)
+
+
+def test_groupby_rolling_mean_3020() -> None:
+    df = pl.DataFrame(
+        {
+            "Date": [
+                "1998-04-12",
+                "1998-04-19",
+                "1998-04-26",
+                "1998-05-03",
+                "1998-05-10",
+                "1998-05-17",
+                "1998-05-24",
+            ],
+            "val": range(7),
+        }
+    ).with_column(pl.col("Date").str.strptime(pl.Date))
+    assert (
+        df.groupby_rolling(index_column="Date", period="1w")
+        .agg(pl.col("val").mean().alias("val_mean"))
+        .frame_equal(
+            pl.DataFrame(
+                {
+                    "Date": [
+                        date(1998, 4, 12),
+                        date(1998, 4, 19),
+                        date(1998, 4, 26),
+                        date(1998, 5, 3),
+                        date(1998, 5, 10),
+                        date(1998, 5, 17),
+                        date(1998, 5, 24),
+                    ],
+                    "val_mean": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
+                }
+            )
+        )
+    )