Skip to content

Commit f232887

Browse files
authored
feat: Support ANSI mode sum expr (int inputs) (#2600)
1 parent b5e4290 commit f232887

File tree

7 files changed

+840
-56
lines changed

7 files changed

+840
-56
lines changed

docs/source/user-guide/latest/compatibility.md

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,12 +32,11 @@ Comet has the following limitations when reading Parquet files:
3232

3333
## ANSI Mode
3434

35-
Comet will fall back to Spark for the following expressions when ANSI mode is enabled. Thes expressions can be enabled by setting
35+
Comet will fall back to Spark for the following expressions when ANSI mode is enabled. These expressions can be enabled by setting
3636
`spark.comet.expression.EXPRNAME.allowIncompatible=true`, where `EXPRNAME` is the Spark expression class name. See
3737
the [Comet Supported Expressions Guide](expressions.md) for more information on this configuration setting.
3838

3939
- Average
40-
- Sum
4140
- Cast (in some cases)
4241

4342
There is an [epic](https://github.com/apache/datafusion-comet/issues/313) where we are tracking the work to fully implement ANSI support.

native/core/src/execution/planner.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ use datafusion::{
7171
use datafusion_comet_spark_expr::{
7272
create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle,
7373
BloomFilterAgg, BloomFilterMightContain, EvalMode, SparkHour, SparkMinute, SparkSecond,
74+
SumInteger,
7475
};
7576
use iceberg::expr::Bind;
7677

@@ -1813,6 +1814,12 @@ impl PhysicalPlanner {
18131814
AggregateUDF::new_from_impl(SumDecimal::try_new(datatype, eval_mode)?);
18141815
AggregateExprBuilder::new(Arc::new(func), vec![child])
18151816
}
1817+
DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
1818+
let eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
1819+
let func =
1820+
AggregateUDF::new_from_impl(SumInteger::try_new(datatype, eval_mode)?);
1821+
AggregateExprBuilder::new(Arc::new(func), vec![child])
1822+
}
18161823
_ => {
18171824
// cast to the result data type of SUM if necessary, we should not expect
18181825
// a cast failure since it should have already been checked at Spark side

native/spark-expr/src/agg_funcs/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ mod correlation;
2121
mod covariance;
2222
mod stddev;
2323
mod sum_decimal;
24+
mod sum_int;
2425
mod variance;
2526

2627
pub use avg::Avg;
@@ -29,4 +30,5 @@ pub use correlation::Correlation;
2930
pub use covariance::Covariance;
3031
pub use stddev::Stddev;
3132
pub use sum_decimal::SumDecimal;
33+
pub use sum_int::SumInteger;
3234
pub use variance::Variance;

0 commit comments

Comments
 (0)