fix: Don't divide by zero in partitioned group-by (#21498)

pola-rs · Feb 27, 2025 · fec5425 · fec5425
1 parent 5e280d9
commit fec5425
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/crates/polars-expr/src/expressions/aggregation.rs b/crates/polars-expr/src/expressions/aggregation.rs
@@ -566,6 +566,12 @@ impl PartitionedAggregation for AggregationExpr {
                         let count = &fields[1];
                         let (agg_count, agg_s) =
                             unsafe { POOL.join(|| count.agg_sum(groups), || sum.agg_sum(groups)) };
+
+                        let agg_count = agg_count.idx().unwrap();
+                        // Ensure that we don't divide by zero.
+                        let agg_count = agg_count
+                            .apply_values(|v| if v == 0 { 1 } else { v })
+                            .into_series();
                         let agg_s = &agg_s / &agg_count;
                         Ok(agg_s?.with_name(new_name).into_column())
                     },