Skip to content

Commit

Permalink
fix: Average expression in Comet Final should handle all null inputs …
Browse files Browse the repository at this point in the history
…from partial Spark aggregation
  • Loading branch information
viirya committed Apr 11, 2024
1 parent b7d2c63 commit 1b5cece
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions core/src/execution/datafusion/expressions/avg.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,9 +176,15 @@ impl Accumulator for AvgAccumulator {
}

fn evaluate(&mut self) -> Result<ScalarValue> {
Ok(ScalarValue::Float64(
self.sum.map(|f| f / self.count as f64),
))
if self.count == 0 {
// If all input are nulls, count will be 0 and we will get null after the division.
// This is consistent with Spark Average implementation.
return Ok(ScalarValue::Float64(None));
} else {
Ok(ScalarValue::Float64(
self.sum.map(|f| f / self.count as f64),
))
}
}

fn size(&self) -> usize {
Expand Down

0 comments on commit 1b5cece

Please sign in to comment.