fix: Propagate null in min_by / max_by for all-null by groups (#26919)

Author: gautamvarmadatla

crates/polars-expr/src/expressions/aggregation.rs

@@ -9,7 +9,6 @@ use polars_core::utils::{_split_offsets, NoNull};
 use polars_ops::prelude::ArgAgg;
 #[cfg(feature = "propagate_nans")]
 use polars_ops::prelude::nan_propagating_aggregate;
-use polars_utils::itertools::Itertools;
 use rayon::prelude::*;
 
 use super::*;
@@ -712,21 +711,23 @@ impl PhysicalExpr for AggMinMaxByExpr {
             unsafe { by_col.agg_arg_min(&by_groups) }
         };
         let idxs_in_groups: &IdxCa = idxs_in_groups.as_materialized_series().as_ref().as_ref();
-        let flat_gather_idxs = match input_groups.as_ref().as_ref() {
+        let gather_idxs: IdxCa = match input_groups.as_ref().as_ref() {
             GroupsType::Idx(g) => idxs_in_groups
-                .into_no_null_iter()
+                .iter()
                 .enumerate()
-                .map(|(group_idx, idx_in_group)| g.all()[group_idx][idx_in_group as usize])
-                .collect_vec(),
+                .map(|(group_idx, idx_in_group)| {
+                    idx_in_group.map(|i| g.all()[group_idx][i as usize])
+                })
+                .collect(),
             GroupsType::Slice { groups, .. } => idxs_in_groups
-                .into_no_null_iter()
+                .iter()
                 .enumerate()
-                .map(|(group_idx, idx_in_group)| groups[group_idx][0] + idx_in_group)
-                .collect_vec(),
+                .map(|(group_idx, idx_in_group)| idx_in_group.map(|i| groups[group_idx][0] + i))
+                .collect(),
         };
 
-        // SAFETY: All indices are within input_col's groups.
-        let gathered = unsafe { input_col.take_slice_unchecked(&flat_gather_idxs) };
+        // SAFETY: All non-null indices are within input_col's groups.
+        let gathered = unsafe { input_col.take_unchecked(&gather_idxs) };
         let agg_state = AggregatedScalar(gathered.with_name(keep_name));
         Ok(AggregationContext::from_agg_state(
             agg_state,

py-polars/tests/unit/operations/aggregation/test_aggregations.py

@@ -1512,3 +1512,53 @@ def test_min_max_by_on_boolean_26847(
     df = pl.DataFrame({"a": [1] * 10, "b": [True] * 10})
     result = df.select(agg(pl.col("a"), pl.col("b")))
     assert result.item() == expected
+
+
+@pytest.mark.parametrize("agg", [pl.Expr.min_by, pl.Expr.max_by])
+def test_min_max_by_all_null_by_group(agg: Callable[..., pl.Expr]) -> None:
+    df = pl.DataFrame(
+        {
+            "g": ["a", "a", "b"],
+            "val": [1, 2, 3],
+            "by": pl.Series([None, None, 5], dtype=pl.Int64),
+        }
+    )
+    expected = pl.DataFrame(
+        {"g": ["a", "b"], "val": pl.Series([None, 3], dtype=pl.Int64)}
+    )
+
+    eager = df.group_by("g", maintain_order=True).agg(agg(pl.col("val"), pl.col("by")))
+    assert_frame_equal(eager, expected)
+
+    streaming = (
+        df.lazy()
+        .group_by("g", maintain_order=True)
+        .agg(agg(pl.col("val"), pl.col("by")))
+        .collect(engine="streaming")
+    )
+    assert_frame_equal(streaming, expected)
+
+
+@pytest.mark.parametrize("agg", [pl.Expr.min_by, pl.Expr.max_by])
+def test_min_max_by_all_null_by_group_slice(agg: Callable[..., pl.Expr]) -> None:
+    df = pl.DataFrame(
+        {
+            "dt": [date(2020, 1, 1), date(2020, 1, 1), date(2020, 2, 1)],
+            "val": [1, 2, 3],
+            "by": pl.Series([None, None, 5], dtype=pl.Int64),
+        }
+    )
+    expected = pl.DataFrame(
+        {
+            "dt": [date(2020, 1, 1), date(2020, 2, 1)],
+            "val": pl.Series([None, 3], dtype=pl.Int64),
+        }
+    )
+
+    result = (
+        df.lazy()
+        .group_by_dynamic("dt", every="1mo")
+        .agg(agg(pl.col("val"), pl.col("by")))
+        .collect()
+    )
+    assert_frame_equal(result, expected)