From 84185b2ac4691844d1414c0fdbcff107617a4d9b Mon Sep 17 00:00:00 2001 From: ritchie Date: Thu, 27 Feb 2025 14:15:01 +0100 Subject: [PATCH] fix: Don't divide by zero in partitioned group-by --- crates/polars-expr/src/expressions/aggregation.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crates/polars-expr/src/expressions/aggregation.rs b/crates/polars-expr/src/expressions/aggregation.rs index 5209fa38a215..89d04dd20ad1 100644 --- 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()) },