You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "coderfender (via GitHub)" <gi...@apache.org> on 2025/10/17 19:13:17 UTC
[PR] [WIP]: Support ANSI mode sum expr [datafusion-comet]
coderfender opened a new pull request, #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes #.
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## How are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501556265
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
Review Comment:
Thank you .I made sure it is all error driven and not panic
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3618783525
@andygrove I rebased with main and fixed merge conflicts . Please review whenever you get a chance
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546645885
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -55,6 +55,8 @@ import org.apache.comet.shims.CometExprShim
*/
object QueryPlanSerde extends Logging with CometExprShim {
+ private val integerTypes = Seq(ByteType, ShortType, IntegerType, LongType)
Review Comment:
Perhaps it would be better to add an `isIntegerType` method to `object DataTypeSupport`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2547766132
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
Review Comment:
Sure. We are essentially checking the `has_null` arrays for early termination in Try mode and also explicitly checking if the sums are nulls in non-try mode (for early termination) which should make these unwraps safe. However, I do agree that we could add exception handling here to handle corrupt state situations.
In try mode :
`
let that_has_all_nulls = states[1].as_boolean().value(0);
let that_overflowed = !that_has_all_nulls && that_sum.is_none();
let this_overflowed = !self.has_all_nulls && self.sum.is_none();
if that_overflowed || this_overflowed {
self.sum = None;
self.has_all_nulls = false;
return Ok(());
}
self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
if that_has_all_nulls {
return Ok(());
}
if self.has_all_nulls {
self.sum = that_sum;
return Ok(());
}
`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3676122730
Fixed clippy errors
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2641462759
##########
spark/src/main/scala/org/apache/comet/serde/aggregates.scala:
##########
@@ -213,17 +213,6 @@ object CometAverage extends CometAggregateExpressionSerde[Average] {
object CometSum extends CometAggregateExpressionSerde[Sum] {
- override def getSupportLevel(sum: Sum): SupportLevel = {
- sum.evalMode match {
- case EvalMode.ANSI if !sum.dataType.isInstanceOf[DecimalType] =>
- Incompatible(Some("ANSI mode for non decimal inputs is not supported"))
Review Comment:
Fixed documentation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3505254319
```
- Windows support *** FAILED *** (3 seconds, 973 milliseconds)
Expected only Comet native operators, but found Project.
plan: Project
+- Window
+- CometExchange
+- CometScan [native_iceberg_compat] parquet
```
Thats actually interesting why would window fall back, we planning to fallback on windows in https://github.com/apache/datafusion-comet/pull/2726 but the test would still preserve windows, so I think we need to check a fallback reason
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2520308349
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,573 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
Review Comment:
This is basically an internal DF API to be leveraged during planning phase
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2520309728
##########
spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala:
##########
@@ -3004,6 +3002,227 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
+ "null_tbl") {
+ val res = sql("SELECT sum(_1) FROM null_tbl")
+ checkSparkAnswerAndOperator(res)
+ assert(res.collect() === Array(Row(null)))
+ }
+ }
+ }
+ }
+
+ test("ANSI support for try_sum - null test") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
+ "null_tbl") {
+ val res = sql("SELECT try_sum(_1) FROM null_tbl")
+ checkSparkAnswerAndOperator(res)
+ assert(res.collect() === Array(Row(null)))
+ }
+ }
+ }
+ }
+
+ test("ANSI support for sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
+ }
+ }
+ }
+ }
+
+ test("ANSI support for try_sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
Review Comment:
Thank you . I fixed the test
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,573 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
Review Comment:
Thank you
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] [WIP]: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3416780935
Draft PR to support sum function - WIP
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2625064276
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
Review Comment:
I think the check here might be an overkill. We are guaranteeing that the accumulator's state is holding `sum` and `has_all_nulls` for `Try` eval mode as per Spark's semantics :
https://github.com/apache/datafusion-comet/pull/2600/changes#diff-21a43f466d9ae04ca6b3f7f64049a595f6705e9e4325965ec980ecd5f77ad036R231
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ // safe to unwrap (since we checked nulls above) but handling error just in case state is corrupt
+ let left = self.sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Current batch's sum is None".to_string(),
+ )
+ })?;
+ let right = that_sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Incoming sum is None".to_string(),
+ )
+ })?;
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+
+ fn resize_helper(&mut self, total_num_groups: usize) {
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ self.resize_helper(total_num_groups);
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int32 => update_groups_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int16 => update_groups_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int8 => update_groups_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "Unsupported data type for SumIntGroupsAccumulator: {:?}",
+ values.data_type()
+ )))
+ }
+ };
+ Ok(())
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> DFResult<ArrayRef> {
+ match emit_to {
+ EmitTo::All => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .iter()
+ .zip(self.has_all_nulls.iter())
+ .map(|(&sum, &is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+
+ self.sums.clear();
+ self.has_all_nulls.clear();
+ Ok(result)
+ }
+ EmitTo::First(n) => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .drain(..n)
+ .zip(self.has_all_nulls.drain(..n))
+ .map(|(sum, is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+ Ok(result)
+ }
+ }
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> DFResult<Vec<ArrayRef>> {
+ let sums = emit_to.take_needed(&mut self.sums);
+
+ if self.eval_mode == EvalMode::Try {
+ let has_all_nulls = emit_to.take_needed(&mut self.has_all_nulls);
+ Ok(vec![
+ Arc::new(Int64Array::from(sums)),
+ Arc::new(BooleanArray::from(has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Int64Array::from(sums))])
+ }
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+
+ let that_sums = values[0].as_primitive::<Int64Type>();
+
+ self.resize_helper(total_num_groups);
+
+ let that_sums_is_all_nulls = if self.eval_mode == EvalMode::Try {
+ Some(values[1].as_boolean())
Review Comment:
I think the check here might be an overkill. We are guaranteeing that the group accumulator's state is holding `sum` and `has_all_nulls` for `Try` eval mode as per Spark's semantics
https://github.com/apache/datafusion-comet/pull/2600/changes#diff-21a43f466d9ae04ca6b3f7f64049a595f6705e9e4325965ec980ecd5f77ad036R467
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3684480902
Thank you for the comments @comphead . I addressed the review comments and pushed a fix. Please take a look whenever you get a chance
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2496105709
##########
native/.DS_Store:
##########
Review Comment:
Please remove this file and add `.DS_Store` to .gitignore
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int32 => update_groups_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int16 => update_groups_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int8 => update_groups_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "Unsupported data type for SumIntGroupsAccumulator: {:?}",
+ values.data_type()
+ )))
+ }
+ };
+ Ok(())
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> DFResult<ArrayRef> {
+ match emit_to {
+ EmitTo::All => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .iter()
+ .zip(self.has_all_nulls.iter())
+ .map(|(&sum, &is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+
+ self.sums.clear();
+ self.has_all_nulls.clear();
+ Ok(result)
+ }
+ EmitTo::First(n) => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .drain(..n)
+ .zip(self.has_all_nulls.drain(..n))
+ .map(|(sum, is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+ Ok(result)
+ }
+ }
+ }
+
+ fn state(&mut self, _emit_to: EmitTo) -> DFResult<Vec<ArrayRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Int64Array::from(self.sums.clone())),
+ Arc::new(BooleanArray::from(self.has_all_nulls.clone())),
+ ])
+ } else {
+ Ok(vec![Arc::new(Int64Array::from(self.sums.clone()))])
+ }
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+
+ let that_sums = values[0].as_primitive::<Int64Type>();
+
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+
+ let that_sums_is_all_nulls = if self.eval_mode == EvalMode::Try {
+ Some(values[1].as_boolean())
+ } else {
+ None
+ };
+
+ for (idx, &group_index) in group_indices.iter().enumerate() {
+ let that_sum = if that_sums.is_null(idx) {
+ None
+ } else {
+ Some(that_sums.value(idx))
+ };
+
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = that_sums_is_all_nulls.unwrap().value(idx);
+
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed =
+ !self.has_all_nulls[group_index] && self.sums[group_index].is_none();
+
+ if that_overflowed || this_overflowed {
+ self.sums[group_index] = None;
+ self.has_all_nulls[group_index] = false;
+ continue;
+ }
+
+ self.has_all_nulls[group_index] =
+ self.has_all_nulls[group_index] && that_has_all_nulls;
+
+ if that_has_all_nulls {
+ continue;
+ }
+
+ if self.has_all_nulls[group_index] {
+ self.sums[group_index] = that_sum;
+ continue;
+ }
+ } else {
+ if that_sum.is_none() {
+ continue;
+ }
+ if self.sums[group_index].is_none() {
+ self.sums[group_index] = that_sum;
+ continue;
+ }
+ }
+
+ // Both sides have non-null. Update sums now
+ let left = self.sums[group_index].unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sums[group_index] = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match left.add_checked(right) {
+ Ok(v) => self.sums[group_index] = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )));
+ } else {
+ // overflow. update flag accordingly
+ self.sums[group_index] = None;
+ self.has_all_nulls[group_index] = false;
+ }
+ }
+ }
+ }
+ }
+ }
+ Ok(())
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
Review Comment:
`use std::mem::size_of_val;`
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
Review Comment:
At line 138 it returns an Err when the conversion fails. Would it make sense to return an Err here too instead of panicking ?
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
Review Comment:
It would be helpful to print the problematic value in the error message.
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
Review Comment:
Why only the `0` item here and at line 245 ?
Shouldn't it iterate over all items ?
```rust
let rows = that_sum_array.len();
for row in 0..rows {
let that_sum = if that_sum_array.is_null(row) {
None
} else {
Some(that_sum_array.value(row))
};
// accumulate here
}
```
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
Review Comment:
This needs to be fully qualified or imported - `use std::mem::size_of_val;`
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ assert!(opt_filter.is_none(), "opt_filter is not supported yet");
Review Comment:
```suggestion
debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
```
?
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
Review Comment:
Print the problematic value too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3676059487
@comphead could you also review this PR when you have time?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3513482104
@comphead , great point. The reason is the same as the special implementation as @sunchao for `Decimal` addition (`SumDecimal` implementation) . If we truly want a nartive window expression, we might want to implement sliding window accumulator (in [df](https://docs.rs/datafusion-functions-aggregate/latest/src/datafusion_functions_aggregate/sum.rs.html)) . I plan to support native sliding windows for both sum decimal and sum int functions in subsequent PRs (once I finish ANSI support for `avg` function)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2641487897
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1555,11 +1612,63 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
(1 to 50).flatMap(_ => Seq((maxDec38_0, 1)))
}
+ test("ANSI support - SUM function") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ // Test long overflow
+ withParquetTable(Seq((Long.MaxValue, 1L), (100L, 1L)), "tbl") {
+ val res = sql("SELECT SUM(_1) FROM tbl")
+ if (ansiEnabled) {
+ checkSparkAnswerMaybeThrows(res) match {
+ case (Some(sparkExc), Some(cometExc)) =>
+ assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
Review Comment:
Thank you . I added explicit comment
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1602,35 +1768,68 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("try_sum overflow - with GROUP BY") {
+ // Test Long overflow with GROUP BY - some groups overflow while some don't
+ withParquetTable(Seq((Long.MaxValue, 1), (100L, 1), (200L, 2), (300L, 2)), "tbl") {
+ val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640521230
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,589 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
Review Comment:
what is the name would be for sum of doubles? 🤔
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640514639
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1555,11 +1612,63 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
(1 to 50).flatMap(_ => Seq((maxDec38_0, 1)))
}
+ test("ANSI support - SUM function") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ // Test long overflow
+ withParquetTable(Seq((Long.MaxValue, 1L), (100L, 1L)), "tbl") {
+ val res = sql("SELECT SUM(_1) FROM tbl")
+ if (ansiEnabled) {
+ checkSparkAnswerMaybeThrows(res) match {
+ case (Some(sparkExc), Some(cometExc)) =>
+ assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
Review Comment:
it would prob good to comment why both of exceptions having this token?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501556590
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
Review Comment:
thank you :) I changed it to fully qualified import
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501555396
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
Review Comment:
Thank you for the review . I believe we wouldnt be needing these checks at all given how we check the data types multiple times (from planning phase , internal functions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3500518505
@andygrove , I believe this should be ready for review (pending CI)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2641454971
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,589 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
Review Comment:
Sure . Let me file an issue once this is merged to break down the implementation into separate impl then
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3558768646
Thank you @andygrove . Let me rebase Ave fix the merge conflicts
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546645885
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -55,6 +55,8 @@ import org.apache.comet.shims.CometExprShim
*/
object QueryPlanSerde extends Logging with CometExprShim {
+ private val integerTypes = Seq(ByteType, ShortType, IntegerType, LongType)
Review Comment:
Perhaps it would be better to add an `isIntegerType` method to `object DataTypeSupport`, or add something in
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3643720049
@andygrove , @martin-g could you please take a (hopefully final) look at the PR and let me know if you think we need further changes ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3643919463
> @andygrove , @martin-g could you please take a (hopefully final) look at the PR and let me know if you think we need further changes ?
The changes in the PR are pretty substantial (~800 new lines of code) so there is a lot to review. I don't see anything that concerns me at this point, but I would like to run some benchmarks to check that there are no regressions in the non ANSI case. Unfortunately, we don't have any existing microbenchmarks that work with ANSI enabled, so we need to create something.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3524738665
@andygrove @comphead . All the tests passed and the PR is rest for review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2547766132
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
Review Comment:
Sure. We are essentially checking the `has_null` arrays for early termination in Try mode and also explicitly checking if the sums are nulls in non-try mode (for early termination) which should make these unwraps safe. However, I do agree that we could add exception handling here to handle corrupt state situations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546880200
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
Review Comment:
Could you add a comment here explaining why the unwraps are safe, and also use `unwrap_or_else` to add error handling rather than panic if these are ever `None`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2547835524
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
Review Comment:
Thank you . I added error handling to make sure that we throw an error when the input value cannot be casted / up-casted to i64
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3558899296
Here are microbenchmark results:
### Legacy mode
```
OpenJDK 64-Bit Server VM 17.0.16+8-Ubuntu-0ubuntu122.04.1 on Linux 6.8.0-87-generic
AMD Ryzen 9 7950X3D 16-Core Processor
TPCDS Micro Benchmarks: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
agg_sum_integers_no_grouping 1720 1812 129 167.4 6.0 1.0X
agg_sum_integers_no_grouping: Comet 1171 1179 11 246.0 4.1 1.5X
Running benchmark: TPCDS Micro Benchmarks
Running case: agg_sum_integers_with_grouping
Stopped after 2 iterations, 5105 ms
Running case: agg_sum_integers_with_grouping: Comet
Stopped after 2 iterations, 2803 ms
OpenJDK 64-Bit Server VM 17.0.16+8-Ubuntu-0ubuntu122.04.1 on Linux 6.8.0-87-generic
AMD Ryzen 9 7950X3D 16-Core Processor
TPCDS Micro Benchmarks: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
agg_sum_integers_with_grouping 2529 2553 34 113.9 8.8 1.0X
agg_sum_integers_with_grouping: Comet 1396 1402 8 206.3 4.8 1.8X
```
### ANSI mode
```
OpenJDK 64-Bit Server VM 17.0.16+8-Ubuntu-0ubuntu122.04.1 on Linux 6.8.0-87-generic
AMD Ryzen 9 7950X3D 16-Core Processor
TPCDS Micro Benchmarks: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
agg_sum_integers_no_grouping 1751 1802 72 164.5 6.1 1.0X
agg_sum_integers_no_grouping: Comet 1358 1360 3 212.1 4.7 1.3X
Running benchmark: TPCDS Micro Benchmarks
Running case: agg_sum_integers_with_grouping
Stopped after 2 iterations, 5563 ms
Running case: agg_sum_integers_with_grouping: Comet
Stopped after 2 iterations, 3668 ms
OpenJDK 64-Bit Server VM 17.0.16+8-Ubuntu-0ubuntu122.04.1 on Linux 6.8.0-87-generic
AMD Ryzen 9 7950X3D 16-Core Processor
TPCDS Micro Benchmarks: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
agg_sum_integers_with_grouping 2772 2782 14 103.9 9.6 1.0X
agg_sum_integers_with_grouping: Comet 1834 1834 1 157.1 6.4 1.5X
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640510387
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1508,11 +1531,28 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
Review Comment:
ditto
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640511128
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1529,11 +1569,28 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for try_sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, try_sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
Review Comment:
same
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2641486614
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1508,11 +1531,28 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
Review Comment:
Thank you . I removed `assert` statements
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501555748
##########
native/.DS_Store:
##########
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2547832428
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
Review Comment:
Thank you . I made the changes to use arrow-rs macro to downcast
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+
+ fn resize_helper(&mut self, total_num_groups: usize) {
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ self.resize_helper(total_num_groups);
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ values
Review Comment:
Thank you . I made the changes to use arrow-rs macro to downcast
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501555396
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
Review Comment:
Thank you for the review . I believe we wouldnt be needing these checks at all given how we check the data types multiple times (from planning phase , internal functions) but let me go ahead and add them just in case
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3465565929
## [Codecov](https://app.codecov.io/gh/apache/datafusion-comet/pull/2600?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `0%` with `2 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 10.88%. Comparing base ([`f09f8af`](https://app.codecov.io/gh/apache/datafusion-comet/commit/f09f8af64c6599255e116a376f4f008f2fd63b43?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`4bbe21a`](https://app.codecov.io/gh/apache/datafusion-comet/commit/4bbe21a74dc86f25c9293e5d6e690c743386247d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 648 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion-comet/pull/2600?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...main/scala/org/apache/comet/serde/aggregates.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/2600?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fcomet%2Fserde%2Faggregates.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9jb21ldC9zZXJkZS9hZ2dyZWdhdGVzLnNjYWxh) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2600?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #2600 +/- ##
=============================================
- Coverage 56.12% 10.88% -45.25%
+ Complexity 976 232 -744
=============================================
Files 119 147 +28
Lines 11743 13733 +1990
Branches 2251 2342 +91
=============================================
- Hits 6591 1495 -5096
- Misses 4012 12063 +8051
+ Partials 1140 175 -965
```
</details>
[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/datafusion-comet/pull/2600?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501894774
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,573 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
Review Comment:
What is the purpose of this method ?
This returns the size of the pointer/reference `&self`
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,573 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
Review Comment:
```suggestion
Ok(None)
```
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,573 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
Review Comment:
```suggestion
// Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
```
##########
spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala:
##########
@@ -3004,6 +3002,227 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
+ "null_tbl") {
+ val res = sql("SELECT sum(_1) FROM null_tbl")
+ checkSparkAnswerAndOperator(res)
+ assert(res.collect() === Array(Row(null)))
+ }
+ }
+ }
+ }
+
+ test("ANSI support for try_sum - null test") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
+ "null_tbl") {
+ val res = sql("SELECT try_sum(_1) FROM null_tbl")
+ checkSparkAnswerAndOperator(res)
+ assert(res.collect() === Array(Row(null)))
+ }
+ }
+ }
+ }
+
+ test("ANSI support for sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
+ }
+ }
+ }
+ }
+
+ test("ANSI support for try_sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
Review Comment:
```suggestion
val res = sql("SELECT _2, try_sum(_1) FROM tbl group by 1")
```
The name of the test says `try_sum`
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -414,7 +416,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
}
case s: Sum =>
if (AggSerde.sumDataTypeSupported(s.dataType) && !s.dataType
- .isInstanceOf[DecimalType]) {
+ .isInstanceOf[DecimalType] && !integerTypes.contains(s.dataType)) {
Review Comment:
Why the integer types are excluded now ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2623932829
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
Review Comment:
DF docs for `merge_batch` state:
```
/// Updates the accumulator's state from an `Array` containing one
/// or more intermediate values.
```
should there be a check that the input has more than one element before accessing `state[1]`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3676666496
Thank you for the approval @and ygrove
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640518075
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1602,35 +1768,68 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("try_sum overflow - with GROUP BY") {
+ // Test Long overflow with GROUP BY - some groups overflow while some don't
+ withParquetTable(Seq((Long.MaxValue, 1), (100L, 1), (200L, 2), (300L, 2)), "tbl") {
+ val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
Review Comment:
why repartition needed here? if its important for the test case it would be good to comment the reason
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3688101216
Thank you @andygrove , @comphead for the guidance and mentorship :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead merged PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2511667913
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -414,7 +416,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
}
case s: Sum =>
if (AggSerde.sumDataTypeSupported(s.dataType) && !s.dataType
- .isInstanceOf[DecimalType]) {
+ .isInstanceOf[DecimalType] && !integerTypes.contains(s.dataType)) {
Review Comment:
Good question. The reason is the same as the special implementation as @sunchao for Decimal addition (SumDecimal implementation) . If we truly want a nartive window expression, we might want to implement sliding window accumulator (in [df](https://docs.rs/datafusion-functions-aggregate/latest/src/datafusion_functions_aggregate/sum.rs.html)) . I plan to support native sliding windows for both sum decimal and sum int functions in subsequent PRs (once I finish ANSI support for avg function)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501562507
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
Review Comment:
Thank you . DataFusion has only two batches at any given point of time (in case of distributed processing) . We have our state (self) and the incoming batch to be merged . Reference : https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udaf.rs
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546795038
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
Review Comment:
I'm not convinced that these unwraps are safe in `TRY` mode.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2631886812
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ // safe to unwrap (since we checked nulls above) but handling error just in case state is corrupt
+ let left = self.sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Current batch's sum is None".to_string(),
+ )
+ })?;
+ let right = that_sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Incoming sum is None".to_string(),
+ )
+ })?;
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+
+ fn resize_helper(&mut self, total_num_groups: usize) {
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ self.resize_helper(total_num_groups);
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int32 => update_groups_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int16 => update_groups_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int8 => update_groups_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "Unsupported data type for SumIntGroupsAccumulator: {:?}",
+ values.data_type()
+ )))
+ }
+ };
+ Ok(())
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> DFResult<ArrayRef> {
+ match emit_to {
+ EmitTo::All => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .iter()
+ .zip(self.has_all_nulls.iter())
+ .map(|(&sum, &is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+
+ self.sums.clear();
+ self.has_all_nulls.clear();
+ Ok(result)
+ }
+ EmitTo::First(n) => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .drain(..n)
+ .zip(self.has_all_nulls.drain(..n))
+ .map(|(sum, is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+ Ok(result)
+ }
+ }
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> DFResult<Vec<ArrayRef>> {
+ let sums = emit_to.take_needed(&mut self.sums);
+
+ if self.eval_mode == EvalMode::Try {
+ let has_all_nulls = emit_to.take_needed(&mut self.has_all_nulls);
+ Ok(vec![
+ Arc::new(Int64Array::from(sums)),
+ Arc::new(BooleanArray::from(has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Int64Array::from(sums))])
+ }
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+
+ let that_sums = values[0].as_primitive::<Int64Type>();
+
+ self.resize_helper(total_num_groups);
+
+ let that_sums_is_all_nulls = if self.eval_mode == EvalMode::Try {
+ Some(values[1].as_boolean())
Review Comment:
Sorry @andygrove , the Github URL I attempted to paste was referring to the state function. However, I pushed a commit to add safe check before accessing elements by index
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
Review Comment:
Sorry @andygrove , the Github URL I attempted to paste was referring to the state function. However, I pushed a commit to add safe check before accessing elements by index
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2631291542
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ // safe to unwrap (since we checked nulls above) but handling error just in case state is corrupt
+ let left = self.sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Current batch's sum is None".to_string(),
+ )
+ })?;
+ let right = that_sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Incoming sum is None".to_string(),
+ )
+ })?;
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+
+ fn resize_helper(&mut self, total_num_groups: usize) {
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ self.resize_helper(total_num_groups);
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int32 => update_groups_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int16 => update_groups_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int8 => update_groups_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "Unsupported data type for SumIntGroupsAccumulator: {:?}",
+ values.data_type()
+ )))
+ }
+ };
+ Ok(())
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> DFResult<ArrayRef> {
+ match emit_to {
+ EmitTo::All => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .iter()
+ .zip(self.has_all_nulls.iter())
+ .map(|(&sum, &is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+
+ self.sums.clear();
+ self.has_all_nulls.clear();
+ Ok(result)
+ }
+ EmitTo::First(n) => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .drain(..n)
+ .zip(self.has_all_nulls.drain(..n))
+ .map(|(sum, is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+ Ok(result)
+ }
+ }
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> DFResult<Vec<ArrayRef>> {
+ let sums = emit_to.take_needed(&mut self.sums);
+
+ if self.eval_mode == EvalMode::Try {
+ let has_all_nulls = emit_to.take_needed(&mut self.has_all_nulls);
+ Ok(vec![
+ Arc::new(Int64Array::from(sums)),
+ Arc::new(BooleanArray::from(has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Int64Array::from(sums))])
+ }
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+
+ let that_sums = values[0].as_primitive::<Int64Type>();
+
+ self.resize_helper(total_num_groups);
+
+ let that_sums_is_all_nulls = if self.eval_mode == EvalMode::Try {
+ Some(values[1].as_boolean())
Review Comment:
Sorry, I still don't understand. The URL seems to link to this PR. Why could the values array not have length zero, or two, or any other length?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640506621
##########
spark/src/main/scala/org/apache/comet/serde/aggregates.scala:
##########
@@ -213,17 +213,6 @@ object CometAverage extends CometAggregateExpressionSerde[Average] {
object CometSum extends CometAggregateExpressionSerde[Sum] {
- override def getSupportLevel(sum: Sum): SupportLevel = {
- sum.evalMode match {
- case EvalMode.ANSI if !sum.dataType.isInstanceOf[DecimalType] =>
- Incompatible(Some("ANSI mode for non decimal inputs is not supported"))
Review Comment:
nice we should probably reflect it in https://datafusion.apache.org/comet/user-guide/latest/compatibility.html#ansi-mode
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640530313
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,589 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
Review Comment:
I'm thinking this match in loop can be expensive? WDYT of having 3 internal implementations
- Legacy sum
- ANSI sum
- Try mode
And call needed impl depending on eval_mode?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2623939256
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Failed to convert value {:?} to i64",
+ int_array.value(i)
+ ))
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ // safe to unwrap (since we checked nulls above) but handling error just in case state is corrupt
+ let left = self.sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Current batch's sum is None".to_string(),
+ )
+ })?;
+ let right = that_sum.ok_or_else(|| {
+ DataFusionError::Internal(
+ "Invalid state in merging batch. Incoming sum is None".to_string(),
+ )
+ })?;
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+
+ fn resize_helper(&mut self, total_num_groups: usize) {
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ self.resize_helper(total_num_groups);
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ as_primitive_array::<Int64Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int32 => update_groups_sum_internal(
+ as_primitive_array::<Int32Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int16 => update_groups_sum_internal(
+ as_primitive_array::<Int16Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ DataType::Int8 => update_groups_sum_internal(
+ as_primitive_array::<Int8Type>(values),
+ group_indices,
+ &mut self.sums,
+ &mut self.has_all_nulls,
+ self.eval_mode,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "Unsupported data type for SumIntGroupsAccumulator: {:?}",
+ values.data_type()
+ )))
+ }
+ };
+ Ok(())
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> DFResult<ArrayRef> {
+ match emit_to {
+ EmitTo::All => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .iter()
+ .zip(self.has_all_nulls.iter())
+ .map(|(&sum, &is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+
+ self.sums.clear();
+ self.has_all_nulls.clear();
+ Ok(result)
+ }
+ EmitTo::First(n) => {
+ let result = Arc::new(Int64Array::from_iter(
+ self.sums
+ .drain(..n)
+ .zip(self.has_all_nulls.drain(..n))
+ .map(|(sum, is_null)| if is_null { None } else { sum }),
+ )) as ArrayRef;
+ Ok(result)
+ }
+ }
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> DFResult<Vec<ArrayRef>> {
+ let sums = emit_to.take_needed(&mut self.sums);
+
+ if self.eval_mode == EvalMode::Try {
+ let has_all_nulls = emit_to.take_needed(&mut self.has_all_nulls);
+ Ok(vec![
+ Arc::new(Int64Array::from(sums)),
+ Arc::new(BooleanArray::from(has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Int64Array::from(sums))])
+ }
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+
+ let that_sums = values[0].as_primitive::<Int64Type>();
+
+ self.resize_helper(total_num_groups);
+
+ let that_sums_is_all_nulls = if self.eval_mode == EvalMode::Try {
+ Some(values[1].as_boolean())
Review Comment:
Same question as earlier. Is `values` guaranteed to have length > 1 at this point?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640509652
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1490,11 +1501,23 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for try_sum - null test") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
+ "null_tbl") {
+ val res = sql("SELECT try_sum(_1) FROM null_tbl")
+ checkSparkAnswerAndOperator(res)
+ assert(res.collect() === Array(Row(null)))
Review Comment:
do we really need this assertion after `checkSparkAnswerAndOperator` ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2640708288
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1490,11 +1501,23 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for try_sum - null test") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
+ "null_tbl") {
+ val res = sql("SELECT try_sum(_1) FROM null_tbl")
+ checkSparkAnswerAndOperator(res)
+ assert(res.collect() === Array(Row(null)))
Review Comment:
@comphead , Yeah , I found some inconsistencies while writing these tests as (partially) discussed here : https://github.com/apache/datafusion-comet/issues/2692 . Also ,I wanted to be completely sure in terms of output given the core functionality of this PR
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1508,11 +1531,28 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
Review Comment:
replied above
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1602,35 +1768,68 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("try_sum overflow - with GROUP BY") {
+ // Test Long overflow with GROUP BY - some groups overflow while some don't
+ withParquetTable(Seq((Long.MaxValue, 1), (100L, 1), (200L, 2), (300L, 2)), "tbl") {
+ val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
Review Comment:
Sure let me add a comment here
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1555,11 +1612,63 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
(1 to 50).flatMap(_ => Seq((maxDec38_0, 1)))
}
+ test("ANSI support - SUM function") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ // Test long overflow
+ withParquetTable(Seq((Long.MaxValue, 1L), (100L, 1L)), "tbl") {
+ val res = sql("SELECT SUM(_1) FROM tbl")
+ if (ansiEnabled) {
+ checkSparkAnswerMaybeThrows(res) match {
+ case (Some(sparkExc), Some(cometExc)) =>
+ assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
Review Comment:
Sure . I thought that the `ansiEnabled` should give away why the overflow message is expected but let me add a comment to make things explicit
##########
spark/src/main/scala/org/apache/comet/serde/aggregates.scala:
##########
@@ -213,17 +213,6 @@ object CometAverage extends CometAggregateExpressionSerde[Average] {
object CometSum extends CometAggregateExpressionSerde[Sum] {
- override def getSupportLevel(sum: Sum): SupportLevel = {
- sum.evalMode match {
- case EvalMode.ANSI if !sum.dataType.isInstanceOf[DecimalType] =>
- Incompatible(Some("ANSI mode for non decimal inputs is not supported"))
Review Comment:
@comphead thank you . I generally follow up the main PR with a separate doc PR and I hope that is okay to do the same thing with this PR too ?
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1529,11 +1569,28 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for try_sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, try_sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
Review Comment:
replied above
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,589 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ as_primitive_array, cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType,
+ BooleanArray, Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
Review Comment:
It would be the same :) . We even have the same name for the decimal sum (which is a custom implementation) as well .
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2641486614
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1508,11 +1531,28 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test (group by)") {
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
+ withParquetTable(
+ Seq(
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "a"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b"),
+ (null.asInstanceOf[java.lang.Long], "b")),
+ "tbl") {
+ val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
+ checkSparkAnswerAndOperator(res)
+ assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
Review Comment:
Removed `assert` statements
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2501557183
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
Review Comment:
Thank you . (replied above)
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,570 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ return Ok(None);
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ panic!("Unsupported data type {}", values.data_type())
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ assert!(opt_filter.is_none(), "opt_filter is not supported yet");
Review Comment:
Done . Thank you
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546659091
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
Review Comment:
nit: arrow-rs has a macro for downcasting that would reduce boilerplate:
```suggestion
as_primitive_array::<Int64Type>(values),
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2547830479
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -55,6 +55,8 @@ import org.apache.comet.shims.CometExprShim
*/
object QueryPlanSerde extends Logging with CometExprShim {
+ private val integerTypes = Seq(ByteType, ShortType, IntegerType, LongType)
Review Comment:
Thank you . I believe this is no longer relevant because we have fundamentally disabled window function support in comet though
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2547926883
##########
spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala:
##########
@@ -3003,6 +3003,227 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test") {
Review Comment:
Thank you . I moved the tests to Aggregate test suites from Expression tests suite
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3644112323
Thank you @andygrove let me go ahead and work on generating benchmarks for ANSI supported aggregate expressions
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3644388368
Created #2883 to support benchmarking with ANSI mode
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3618875215
Resolved scalafix error
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3500530248
@andygrove , @comphead I believe this should be ready for review (pending CI)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3561308028
@andygrove , Please take a look whenever you get a chance .
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546662149
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+
+ fn resize_helper(&mut self, total_num_groups: usize) {
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ self.resize_helper(total_num_groups);
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ values
Review Comment:
same commet as above about using arrow-rs macro for downcasting
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
+
+ match self.eval_mode {
+ EvalMode::Legacy => {
+ self.sum = Some(left.add_wrapping(right));
+ }
+ EvalMode::Ansi | EvalMode::Try => match left.add_checked(right) {
+ Ok(v) => self.sum = Some(v),
+ Err(_) => {
+ if self.eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(arithmetic_overflow_error("integer")));
+ } else {
+ self.sum = None;
+ self.has_all_nulls = false;
+ }
+ }
+ },
+ }
+ Ok(())
+ }
+}
+
+struct SumIntGroupsAccumulator {
+ sums: Vec<Option<i64>>,
+ has_all_nulls: Vec<bool>,
+ eval_mode: EvalMode,
+}
+
+impl SumIntGroupsAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ Self {
+ sums: Vec::new(),
+ eval_mode,
+ has_all_nulls: Vec::new(),
+ }
+ }
+
+ fn resize_helper(&mut self, total_num_groups: usize) {
+ if self.eval_mode == EvalMode::Try {
+ self.sums.resize(total_num_groups, Some(0));
+ self.has_all_nulls.resize(total_num_groups, true);
+ } else {
+ self.sums.resize(total_num_groups, None);
+ self.has_all_nulls.resize(total_num_groups, false);
+ }
+ }
+}
+
+impl GroupsAccumulator for SumIntGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ fn update_groups_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ group_indices: &[usize],
+ sums: &mut [Option<i64>],
+ has_all_nulls: &mut [bool],
+ eval_mode: EvalMode,
+ ) -> DFResult<()>
+ where
+ T: ArrowPrimitiveType,
+ T::Native: ArrowNativeType,
+ {
+ for (i, &group_index) in group_indices.iter().enumerate() {
+ if !int_array.is_null(i) {
+ // there is an overflow in prev group in try eval. Skip processing
+ if eval_mode == EvalMode::Try
+ && !has_all_nulls[group_index]
+ && sums[group_index].is_none()
+ {
+ continue;
+ }
+ let v = int_array.value(i).to_i64().ok_or_else(|| {
+ DataFusionError::Internal("Failed to convert value to i64".to_string())
+ })?;
+ match eval_mode {
+ EvalMode::Legacy => {
+ sums[group_index] =
+ Some(sums[group_index].unwrap_or(0).add_wrapping(v));
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match sums[group_index].unwrap_or(0).add_checked(v) {
+ Ok(new_sum) => {
+ sums[group_index] = Some(new_sum);
+ }
+ Err(_) => {
+ if eval_mode == EvalMode::Ansi {
+ return Err(DataFusionError::from(
+ arithmetic_overflow_error("integer"),
+ ));
+ } else {
+ sums[group_index] = None;
+ }
+ }
+ };
+ }
+ }
+ has_all_nulls[group_index] = false
+ }
+ }
+ Ok(())
+ }
+
+ debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ let values = &values[0];
+ self.resize_helper(total_num_groups);
+
+ match values.data_type() {
+ DataType::Int64 => update_groups_sum_internal(
+ values
Review Comment:
same comment as above about using arrow-rs macro for downcasting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546785073
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
Review Comment:
```suggestion
let v = int_array.value(i).to_i64().ok_or_else(|| {
DataFusionError::Internal("Failed to convert value to i64".to_string())
})?;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2546808679
##########
spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala:
##########
@@ -3003,6 +3003,227 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("ANSI support for sum - null test") {
Review Comment:
These tests should probably be in `CometAggregateSuite`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on code in PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#discussion_r2547766132
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -0,0 +1,576 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{arithmetic_overflow_error, EvalMode};
+use arrow::array::{
+ cast::AsArray, Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray,
+ Int64Array, PrimitiveArray,
+};
+use arrow::datatypes::{
+ ArrowNativeType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type,
+};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::Volatility::Immutable;
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::{any::Any, sync::Arc};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SumInteger {
+ signature: Signature,
+ eval_mode: EvalMode,
+}
+
+impl SumInteger {
+ pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
+ match data_type {
+ DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
+ signature: Signature::user_defined(Immutable),
+ eval_mode,
+ }),
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for SumInteger".into(),
+ )),
+ }
+ }
+}
+
+impl AggregateUDFImpl for SumInteger {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "sum"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> {
+ Ok(Box::new(SumIntegerAccumulator::new(self.eval_mode)))
+ }
+
+ fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ Arc::new(Field::new("sum", DataType::Int64, true)),
+ Arc::new(Field::new("has_all_nulls", DataType::Boolean, false)),
+ ])
+ } else {
+ Ok(vec![Arc::new(Field::new("sum", DataType::Int64, true))])
+ }
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ _args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(SumIntGroupsAccumulator::new(self.eval_mode)))
+ }
+
+ fn reverse_expr(&self) -> ReversedUDAF {
+ ReversedUDAF::Identical
+ }
+}
+
+#[derive(Debug)]
+struct SumIntegerAccumulator {
+ sum: Option<i64>,
+ eval_mode: EvalMode,
+ has_all_nulls: bool,
+}
+
+impl SumIntegerAccumulator {
+ fn new(eval_mode: EvalMode) -> Self {
+ if eval_mode == EvalMode::Try {
+ Self {
+ // Try mode starts with 0 (because if this is init to None we cant say if it is none due to all nulls or due to an overflow)
+ sum: Some(0),
+ has_all_nulls: true,
+ eval_mode,
+ }
+ } else {
+ Self {
+ sum: None,
+ has_all_nulls: false,
+ eval_mode,
+ }
+ }
+ }
+}
+
+impl Accumulator for SumIntegerAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> {
+ // accumulator internal to add sum and return null sum (and has_nulls false) if there is an overflow in Try Eval mode
+ fn update_sum_internal<T>(
+ int_array: &PrimitiveArray<T>,
+ eval_mode: EvalMode,
+ mut sum: i64,
+ ) -> Result<Option<i64>, DataFusionError>
+ where
+ T: ArrowPrimitiveType,
+ {
+ for i in 0..int_array.len() {
+ if !int_array.is_null(i) {
+ let v = int_array.value(i).to_i64().unwrap();
+ match eval_mode {
+ EvalMode::Legacy => {
+ sum = v.add_wrapping(sum);
+ }
+ EvalMode::Ansi | EvalMode::Try => {
+ match v.add_checked(sum) {
+ Ok(v) => sum = v,
+ Err(_e) => {
+ return if eval_mode == EvalMode::Ansi {
+ Err(DataFusionError::from(arithmetic_overflow_error(
+ "integer",
+ )))
+ } else {
+ Ok(None)
+ };
+ }
+ };
+ }
+ }
+ }
+ }
+ Ok(Some(sum))
+ }
+
+ if self.eval_mode == EvalMode::Try && !self.has_all_nulls && self.sum.is_none() {
+ // we saw an overflow earlier (Try eval mode). Skip processing
+ return Ok(());
+ }
+ let values = &values[0];
+ if values.len() == values.null_count() {
+ Ok(())
+ } else {
+ // No nulls so there should be a non-null sum / null incase overflow in Try eval
+ let running_sum = self.sum.unwrap_or(0);
+ let sum = match values.data_type() {
+ DataType::Int64 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int32 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int32Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int16 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int16Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ DataType::Int8 => update_sum_internal(
+ values
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int8Type>>()
+ .unwrap(),
+ self.eval_mode,
+ running_sum,
+ )?,
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "unsupported data type: {:?}",
+ values.data_type()
+ )));
+ }
+ };
+ self.sum = sum;
+ self.has_all_nulls = false;
+ Ok(())
+ }
+ }
+
+ fn evaluate(&mut self) -> DFResult<ScalarValue> {
+ if self.has_all_nulls {
+ Ok(ScalarValue::Int64(None))
+ } else {
+ Ok(ScalarValue::Int64(self.sum))
+ }
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self)
+ }
+
+ fn state(&mut self) -> DFResult<Vec<ScalarValue>> {
+ if self.eval_mode == EvalMode::Try {
+ Ok(vec![
+ ScalarValue::Int64(self.sum),
+ ScalarValue::Boolean(Some(self.has_all_nulls)),
+ ])
+ } else {
+ Ok(vec![ScalarValue::Int64(self.sum)])
+ }
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> DFResult<()> {
+ let that_sum_array = states[0].as_primitive::<Int64Type>();
+ let that_sum = if that_sum_array.is_null(0) {
+ None
+ } else {
+ Some(that_sum_array.value(0))
+ };
+
+ // Check for overflow for early termination
+ if self.eval_mode == EvalMode::Try {
+ let that_has_all_nulls = states[1].as_boolean().value(0);
+ let that_overflowed = !that_has_all_nulls && that_sum.is_none();
+ let this_overflowed = !self.has_all_nulls && self.sum.is_none();
+ if that_overflowed || this_overflowed {
+ self.sum = None;
+ self.has_all_nulls = false;
+ return Ok(());
+ }
+ self.has_all_nulls = self.has_all_nulls && that_has_all_nulls;
+ if that_has_all_nulls {
+ return Ok(());
+ }
+ if self.has_all_nulls {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ } else {
+ if that_sum.is_none() {
+ return Ok(());
+ }
+ if self.sum.is_none() {
+ self.sum = that_sum;
+ return Ok(());
+ }
+ }
+
+ let left = self.sum.unwrap();
+ let right = that_sum.unwrap();
Review Comment:
Sure. We are essentially checking the `has_null` arrays for early termination in Try mode and also explicitly checking if the sums are nulls in non-try mode (for early termination) which should make these unwraps safe. However, I do agree that we could add exception handling here to handle corrupt state situations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3645020170
Ran benchmarking with ANSI mode enabled for SUM aggregated function only and here are the results
```
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM, ansi mode enabled : true
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : true
Stopped after 32 iterations, 2013 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : true
Stopped after 38 iterations, 2032 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM, ansi mode enabled : true: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : true 48 63 17 219.2 4.6 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : true 36 53 17 289.7 3.5 1.3X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM, ansi mode enabled : true
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : true
Stopped after 30 iterations, 2071 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : true
Stopped after 36 iterations, 2011 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM, ansi mode enabled : true: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : true 47 69 13 221.2 4.5 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : true 35 56 33 303.5 3.3 1.4X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM, ansi mode enabled : true
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : true
Stopped after 33 iterations, 2049 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : true
Stopped after 33 iterations, 2035 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM, ansi mode enabled : true: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : true 50 62 17 211.4 4.7 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : true 37 62 33 281.9 3.5 1.3X
Running benchmark: Grouped HashAgg Exec: multiple group keys (cardinality 100), single aggregate SUM
Running case: SQL Parquet - Spark (SUM) isANSIMode: true
Stopped after 32 iterations, 2026 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: true
Stopped after 34 iterations, 2010 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: multiple group keys (cardinality 100), single aggregate SUM: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: true 49 63 23 213.9 4.7 1.0X
SQL Parquet - Comet (SUM) isANSIMode: true 36 59 23 287.7 3.5 1.3X
Running benchmark: Grouped HashAgg Exec: multiple group keys (cardinality 1024), single aggregate SUM
Running case: SQL Parquet - Spark (SUM) isANSIMode: true
Stopped after 29 iterations, 2131 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: true
Stopped after 43 iterations, 2014 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: multiple group keys (cardinality 1024), single aggregate SUM: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: true 48 73 32 217.1 4.6 1.0X
SQL Parquet - Comet (SUM) isANSIMode: true 36 47 18 290.1 3.4 1.3X
Running benchmark: Grouped HashAgg Exec: multiple group keys (cardinality 1048576), single aggregate SUM
Running case: SQL Parquet - Spark (SUM) isANSIMode: true
Stopped after 2 iterations, 2952 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: true
Stopped after 5 iterations, 2276 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: multiple group keys (cardinality 1048576), single aggregate SUM: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: true 1463 1476 18 7.2 139.5 1.0X
SQL Parquet - Comet (SUM) isANSIMode: true 396 455 92 26.5 37.8 3.7X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 100), multiple aggregates SUM isANSIMode: true
Running case: SQL Parquet - Spark (SUM) isANSIMode: true
Stopped after 13 iterations, 2075 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: true
Stopped after 35 iterations, 2051 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 100), multiple aggregates SUM isANSIMode: true: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: true 56 160 78 188.0 5.3 1.0X
SQL Parquet - Comet (SUM) isANSIMode: true 34 59 45 309.5 3.2 1.6X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1024), multiple aggregates SUM isANSIMode: true
Running case: SQL Parquet - Spark (SUM) isANSIMode: true
Stopped after 16 iterations, 2103 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: true
Stopped after 37 iterations, 2016 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1024), multiple aggregates SUM isANSIMode: true: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: true 55 131 97 189.6 5.3 1.0X
SQL Parquet - Comet (SUM) isANSIMode: true 38 54 30 273.0 3.7 1.4X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1048576), multiple aggregates SUM isANSIMode: true
Running case: SQL Parquet - Spark (SUM) isANSIMode: true
Stopped after 20 iterations, 2118 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: true
Stopped after 22 iterations, 2025 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1048576), multiple aggregates SUM isANSIMode: true: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: true 55 106 71 189.8 5.3 1.0X
SQL Parquet - Comet (SUM) isANSIMode: true 35 92 73 301.9 3.3 1.6X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM on decimal
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : true
Stopped after 2 iterations, 2658 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : true
Stopped after 9 iterations, 2233 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM on decimal: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : true 1257 1329 102 8.3 119.9 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : true 237 248 15 44.2 22.6 5.3X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM on decimal
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : true
Stopped after 2 iterations, 2615 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : true
Stopped after 8 iterations, 2222 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM on decimal: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : true 1307 1308 0 8.0 124.7 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : true 257 278 29 40.7 24.6 5.1X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM on decimal
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : true
Stopped after 2 iterations, 8222 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : true
Stopped after 2 iterations, 2853 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM on decimal: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : true 4082 4111 42 2.6 389.3 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : true 1412 1427 20 7.4 134.7 2.9X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM, ansi mode enabled : false
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : false
Stopped after 9 iterations, 2064 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : false
Stopped after 12 iterations, 2025 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM, ansi mode enabled : false: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : false 221 229 12 47.5 21.0 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : false 153 169 25 68.5 14.6 1.4X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM, ansi mode enabled : false
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : false
Stopped after 9 iterations, 2207 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : false
Stopped after 12 iterations, 2053 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM, ansi mode enabled : false: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : false 226 245 30 46.4 21.6 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : false 165 171 4 63.4 15.8 1.4X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM, ansi mode enabled : false
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : false
Stopped after 2 iterations, 4153 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : false
Stopped after 3 iterations, 2645 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM, ansi mode enabled : false: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : false 2069 2077 11 5.1 197.3 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : false 841 882 49 12.5 80.2 2.5X
Running benchmark: Grouped HashAgg Exec: multiple group keys (cardinality 100), single aggregate SUM
Running case: SQL Parquet - Spark (SUM) isANSIMode: false
Stopped after 4 iterations, 2340 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: false
Stopped after 8 iterations, 2320 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: multiple group keys (cardinality 100), single aggregate SUM: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: false 574 585 12 18.3 54.8 1.0X
SQL Parquet - Comet (SUM) isANSIMode: false 269 290 33 39.0 25.6 2.1X
Running benchmark: Grouped HashAgg Exec: multiple group keys (cardinality 1024), single aggregate SUM
Running case: SQL Parquet - Spark (SUM) isANSIMode: false
Stopped after 2 iterations, 5877 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: false
Stopped after 2 iterations, 2924 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: multiple group keys (cardinality 1024), single aggregate SUM: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: false 2896 2939 60 3.6 276.2 1.0X
SQL Parquet - Comet (SUM) isANSIMode: false 1400 1462 88 7.5 133.5 2.1X
Running benchmark: Grouped HashAgg Exec: multiple group keys (cardinality 1048576), single aggregate SUM
Running case: SQL Parquet - Spark (SUM) isANSIMode: false
Stopped after 2 iterations, 13239 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: false
Stopped after 2 iterations, 10603 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: multiple group keys (cardinality 1048576), single aggregate SUM: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: false 6580 6620 57 1.6 627.5 1.0X
SQL Parquet - Comet (SUM) isANSIMode: false 5265 5302 52 2.0 502.1 1.2X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 100), multiple aggregates SUM isANSIMode: false
Running case: SQL Parquet - Spark (SUM) isANSIMode: false
Stopped after 6 iterations, 2011 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: false
Stopped after 9 iterations, 2165 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 100), multiple aggregates SUM isANSIMode: false: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: false 290 335 74 36.2 27.7 1.0X
SQL Parquet - Comet (SUM) isANSIMode: false 207 241 59 50.6 19.8 1.4X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1024), multiple aggregates SUM isANSIMode: false
Running case: SQL Parquet - Spark (SUM) isANSIMode: false
Stopped after 6 iterations, 2031 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: false
Stopped after 9 iterations, 2281 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1024), multiple aggregates SUM isANSIMode: false: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: false 303 339 54 34.6 28.9 1.0X
SQL Parquet - Comet (SUM) isANSIMode: false 227 253 32 46.2 21.7 1.3X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1048576), multiple aggregates SUM isANSIMode: false
Running case: SQL Parquet - Spark (SUM) isANSIMode: false
Stopped after 2 iterations, 6469 ms
Running case: SQL Parquet - Comet (SUM) isANSIMode: false
Stopped after 2 iterations, 3608 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1048576), multiple aggregates SUM isANSIMode: false: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) isANSIMode: false 3222 3235 18 3.3 307.3 1.0X
SQL Parquet - Comet (SUM) isANSIMode: false 1562 1804 343 6.7 148.9 2.1X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM on decimal
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : false
Stopped after 2 iterations, 2433 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : false
Stopped after 8 iterations, 2034 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 100), single aggregate SUM on decimal: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : false 1210 1217 10 8.7 115.4 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : false 249 254 4 42.0 23.8 4.9X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM on decimal
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : false
Stopped after 2 iterations, 2483 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : false
Stopped after 7 iterations, 2099 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1024), single aggregate SUM on decimal: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : false 1235 1242 10 8.5 117.7 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : false 255 300 54 41.2 24.3 4.8X
Running benchmark: Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM on decimal
Running case: SQL Parquet - Spark (SUM) ansi mode enabled : false
Stopped after 2 iterations, 8194 ms
Running case: SQL Parquet - Comet (SUM) ansi mode enabled : false
Stopped after 2 iterations, 2820 ms
OpenJDK 64-Bit Server VM 17.0.16+8-LTS on Mac OS X 16.0
Apple M2 Max
Grouped HashAgg Exec: single group key (cardinality 1048576), single aggregate SUM on decimal: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------
SQL Parquet - Spark (SUM) ansi mode enabled : false 4069 4097 40 2.6 388.0 1.0X
SQL Parquet - Comet (SUM) ansi mode enabled : false 1381 1410 42 7.6 131.7 2.9X
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr (int inputs) [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3657935293
Rebased with main
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Support ANSI mode sum expr [datafusion-comet]
Posted by "coderfender (via GitHub)" <gi...@apache.org>.
coderfender commented on PR #2600:
URL: https://github.com/apache/datafusion-comet/pull/2600#issuecomment-3524419279
> ```
> - Windows support *** FAILED *** (3 seconds, 973 milliseconds)
> Expected only Comet native operators, but found Project.
> plan: Project
> +- Window
> +- CometExchange
> +- CometScan [native_iceberg_compat] parquet
> ```
>
> Thats actually interesting why would window fall back, we planning to fallback on windows in #2726 but the test would still preserve windows, so I think we need to check a fallback reason
This is no longer a failing test after I rebased with main branch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org