test: Add test for 16642 (#16646)

pola-rs · Jun 1, 2024 · 4c80271 · 4c80271
1 parent ae70fd4
commit 4c80271
Showing 1 changed file with 61 additions and 0 deletions.
diff --git a/py-polars/tests/unit/operations/aggregation/test_aggregations.py b/py-polars/tests/unit/operations/aggregation/test_aggregations.py
@@ -594,3 +594,64 @@ def test_multi_arg_structify_15834() -> None:
  {"a": 1, "value": 0.5896839894245691},
  ],
  }
+
+
+def test_filter_aggregation_16642() -> None:
+ df = pl.DataFrame(
+ {
+ "datetime": [
+ datetime(2022, 1, 1, 11, 0),
+ datetime(2022, 1, 1, 11, 1),
+ datetime(2022, 1, 1, 11, 2),
+ datetime(2022, 1, 1, 11, 3),
+ datetime(2022, 1, 1, 11, 4),
+ datetime(2022, 1, 1, 11, 5),
+ datetime(2022, 1, 1, 11, 6),
+ datetime(2022, 1, 1, 11, 7),
+ datetime(2022, 1, 1, 11, 8),
+ datetime(2022, 1, 1, 11, 9, 1),
+ datetime(2022, 1, 2, 11, 0),
+ datetime(2022, 1, 2, 11, 1),
+ datetime(2022, 1, 2, 11, 2),
+ datetime(2022, 1, 2, 11, 3),
+ datetime(2022, 1, 2, 11, 4),
+ datetime(2022, 1, 2, 11, 5),
+ datetime(2022, 1, 2, 11, 6),
+ datetime(2022, 1, 2, 11, 7),
+ datetime(2022, 1, 2, 11, 8),
+ datetime(2022, 1, 2, 11, 9, 1),
+ ],
+ "alpha": [
+ "A",
+ "B",
+ "C",
+ "D",
+ "E",
+ "F",
+ "G",
+ "H",
+ "I",
+ "J",
+ "A",
+ "B",
+ "C",
+ "D",
+ "E",
+ "F",
+ "G",
+ "H",
+ "I",
+ "J",
+ ],
+ "num": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+ }
+ )
+ grouped = df.group_by(pl.col("datetime").dt.date())
+
+ ts_filter = pl.col("datetime").dt.time() <= pl.time(11, 3)
+
+ report = grouped.agg(pl.col("num").filter(ts_filter).max()).sort("datetime")
+ assert report.to_dict(as_series=False) == {
+ "datetime": [date(2022, 1, 1), date(2022, 1, 2)],
+ "num": [3, 3],
+ }