Unify Spark TrySum implementation with DataFusion Sum implementation #19593
Conversation
| DataType::Int64 => Ok(DataType::Int64), | ||
| DataType::UInt64 => Ok(DataType::UInt64), | ||
| DataType::Float64 => Ok(DataType::Float64), | ||
| DataType::Null => Ok(DataType::Float64), |
There was a problem hiding this comment.
Spark try_sum had test case where null inputs will return a null as double, so adding that here
| if args.expr_fields[0].data_type() == &DataType::Null { | ||
| return Ok(Box::new(NoopAccumulator::new(ScalarValue::Float64(None)))); | ||
| } |
| // See COMMENTS.md to understand why nullable is set to true | ||
| Field::new_list_field(args.return_type().clone(), true), | ||
| false, | ||
| match (args.is_distinct, self.try_sum_mode) { |
There was a problem hiding this comment.
It's a bit ugly here, would be nice if we manage to make accumulators own their state fields, see: #14701 (comment)
| // Can overflow into null | ||
| if self.try_sum_mode { | ||
| return SetMonotonicity::NotMonotonic; | ||
| } |
There was a problem hiding this comment.
I'm not certain on this 🤔
| Err(e) => { | ||
| return Err(e.into()); | ||
| } |
There was a problem hiding this comment.
Technically these errs are unreachable as I think only arithmeticoverflow error variant is returned for add_checked but keeping this in case
| // Only difference from TrySumAccumulator is that it verifies the resulting sum | ||
| // can fit within the decimals precision; if Rust had specialization we could unify | ||
| // the two types (╥﹏╥) | ||
| #[derive(Debug)] | ||
| struct TrySumDecimalAccumulator<T: DecimalType + std::fmt::Debug> { | ||
| inner: TrySumAccumulator<T>, | ||
| } |
There was a problem hiding this comment.
Open to any ideas to unify these structs
| /// Thin wrapper over DataFusion native [`Sum`] which is configurable into a try | ||
| /// sum mode to return `null` on overflows. We need this thin wrapper to provide | ||
| /// the `try_sum` named function for use in Spark. | ||
| #[derive(PartialEq, Eq, Hash, Debug)] |
There was a problem hiding this comment.
Would be nice if we had a simpler way to do these types of thin wrappers 🤔
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { |
There was a problem hiding this comment.
These tests seemed to already be covered by SLTs so removed them
1 participant
Which issue does this PR close?
Rationale for this change
Refactor try_sum to reuse existing sum code from native Datafusion version.
What changes are included in this PR?
Add new mode to Sum which allows toggling into try_sum mode, affecting which accumulator is used.
Add new try_sum accumulator for sum which returns null if sum at any point overflows.
Are these changes tested?
Existing tests.
Are there any user-facing changes?
No