feat: Support Ansi mode in abs function

core/benches/cast_from_string.rs

@@ -17,7 +17,7 @@
 
 use arrow_array::{builder::StringBuilder, RecordBatch};
 use arrow_schema::{DataType, Field, Schema};
-use comet::execution::datafusion::expressions::cast::{Cast, EvalMode};
+use comet::execution::datafusion::expressions::{cast::Cast, EvalMode};
 use criterion::{criterion_group, criterion_main, Criterion};
 use datafusion_physical_expr::{expressions::Column, PhysicalExpr};
 use std::sync::Arc;

core/benches/cast_numeric.rs

@@ -17,7 +17,7 @@
 
 use arrow_array::{builder::Int32Builder, RecordBatch};
 use arrow_schema::{DataType, Field, Schema};
-use comet::execution::datafusion::expressions::cast::{Cast, EvalMode};
+use comet::execution::datafusion::expressions::{cast::Cast, EvalMode};
 use criterion::{criterion_group, criterion_main, Criterion};
 use datafusion_physical_expr::{expressions::Column, PhysicalExpr};
 use std::sync::Arc;

core/src/errors.rs

@@ -88,6 +88,9 @@ pub enum CometError {
         to_type: String,
     },
 
+    #[error("[ARITHMETIC_OVERFLOW] {from_type} overflow. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")]
+    ArithmeticOverflow { from_type: String },
+
     #[error(transparent)]
     Arrow {
         #[from]

core/src/execution/datafusion/expressions/abs.rs

@@ -0,0 +1,87 @@
+// 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 arrow::datatypes::DataType;
+use arrow_schema::ArrowError;
+use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature};
+use datafusion_common::DataFusionError;
+use datafusion_functions::math;
+use std::{any::Any, sync::Arc};
+
+use crate::errors::CometError;
+
+use super::EvalMode;
+
+fn arithmetic_overflow_error(from_type: &str) -> CometError {
+    CometError::ArithmeticOverflow {
+        from_type: from_type.to_string(),
+    }
+}
+
+#[derive(Debug)]
+pub struct CometAbsFunc {
+    inner_abs_func: Arc<dyn ScalarUDFImpl>,
+    eval_mode: EvalMode,
+    data_type_name: String,
+}
+
+impl CometAbsFunc {
+    pub fn new(eval_mode: EvalMode, data_type_name: String) -> Self {
+        Self {
+            inner_abs_func: math::abs().inner(),
+            eval_mode,
+            data_type_name,
+        }
+    }
+}
+
+impl ScalarUDFImpl for CometAbsFunc {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+    fn name(&self) -> &str {
+        "abs"
+    }
+
+    fn signature(&self) -> &Signature {
+        self.inner_abs_func.signature()
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType, DataFusionError> {
+        self.inner_abs_func.return_type(arg_types)
+    }
+
+    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
+        match self.inner_abs_func.invoke(args) {
+            Ok(result) => Ok(result),
+            Err(DataFusionError::ArrowError(ArrowError::ComputeError(msg), trace))
+                if msg.contains("overflow") =>
+            {
+                if self.eval_mode == EvalMode::Legacy {
+                    Ok(args[0].clone())
+                } else {
+                    let msg = arithmetic_overflow_error(&self.data_type_name).to_string();
+                    Err(DataFusionError::ArrowError(
+                        ArrowError::ComputeError(msg),
+                        trace,
+                    ))
+                }
+            }
+            other => other,
+        }
+    }
+}

core/src/execution/datafusion/expressions/cast.rs

@@ -51,6 +51,8 @@ use crate::{
     },
 };
 
+use super::EvalMode;
+
 static TIMESTAMP_FORMAT: Option<&str> = Some("%Y-%m-%d %H:%M:%S%.f");
 
 static CAST_OPTIONS: CastOptions = CastOptions {
@@ -60,13 +62,6 @@ static CAST_OPTIONS: CastOptions = CastOptions {
         .with_timestamp_format(TIMESTAMP_FORMAT),
 };
 
-#[derive(Debug, Hash, PartialEq, Clone, Copy)]
-pub enum EvalMode {
-    Legacy,
-    Ansi,
-    Try,
-}
-
 #[derive(Debug, Hash)]
 pub struct Cast {
     pub child: Arc<dyn PhysicalExpr>,

core/src/execution/datafusion/expressions/mod.rs

@@ -24,6 +24,7 @@ pub mod if_expr;
 mod normalize_nan;
 pub mod scalar_funcs;
 pub use normalize_nan::NormalizeNaNAndZero;
+pub mod abs;
 pub mod avg;
 pub mod avg_decimal;
 pub mod bloom_filter_might_contain;
@@ -37,3 +38,10 @@ pub mod sum_decimal;
 pub mod temporal;
 mod utils;
 pub mod variance;
+
+#[derive(Debug, Hash, PartialEq, Clone, Copy)]
+pub enum EvalMode {
+    Legacy,
+    Ansi,
+    Try,
+}

core/src/execution/datafusion/planner.rs

@@ -24,7 +24,6 @@ use datafusion::{
     arrow::{compute::SortOptions, datatypes::SchemaRef},
     common::DataFusionError,
     execution::FunctionRegistry,
-    functions::math,
     logical_expr::{
         BuiltinScalarFunction, Operator as DataFusionOperator, ScalarFunctionDefinition,
     },
@@ -52,6 +51,7 @@ use datafusion_common::{
     tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter},
     JoinType as DFJoinType, ScalarValue,
 };
+use datafusion_physical_expr::udf::ScalarUDF;
 use itertools::Itertools;
 use jni::objects::GlobalRef;
 use num::{BigInt, ToPrimitive};
@@ -65,7 +65,7 @@ use crate::{
                 avg_decimal::AvgDecimal,
                 bitwise_not::BitwiseNotExpr,
                 bloom_filter_might_contain::BloomFilterMightContain,
-                cast::{Cast, EvalMode},
+                cast::Cast,
                 checkoverflow::CheckOverflow,
                 correlation::Correlation,
                 covariance::Covariance,
@@ -95,6 +95,8 @@ use crate::{
     },
 };
 
+use super::expressions::{abs::CometAbsFunc, EvalMode};
+
 // For clippy error on type_complexity.
 type ExecResult<T> = Result<T, ExecutionError>;
 type PhyAggResult = Result<Vec<Arc<dyn AggregateExpr>>, ExecutionError>;
@@ -149,6 +151,20 @@ impl PhysicalPlanner {
         }
     }
 
+    fn eval_mode_from_str(
+        eval_mode_str: &str,
+        allow_try: bool,
+    ) -> Result<EvalMode, ExecutionError> {
+        match eval_mode_str {
+            "ANSI" => Ok(EvalMode::Ansi),
+            "LEGACY" => Ok(EvalMode::Legacy),
+            "TRY" if allow_try => Ok(EvalMode::Try),
+            other => Err(ExecutionError::GeneralError(format!(
+                "Invalid EvalMode: \"{other}\""
+            ))),
+        }
+    }
+
     /// Create a DataFusion physical expression from Spark physical expression
     fn create_expr(
         &self,
@@ -348,16 +364,7 @@ impl PhysicalPlanner {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema)?;
                 let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap());
                 let timezone = expr.timezone.clone();
-                let eval_mode = match expr.eval_mode.as_str() {
-                    "ANSI" => EvalMode::Ansi,
-                    "TRY" => EvalMode::Try,
-                    "LEGACY" => EvalMode::Legacy,
-                    other => {
-                        return Err(ExecutionError::GeneralError(format!(
-                            "Invalid Cast EvalMode: \"{other}\""
-                        )))
-                    }
-                };
+                let eval_mode = Self::eval_mode_from_str(expr.eval_mode.as_str(), true)?;
                 Ok(Arc::new(Cast::new(child, datatype, eval_mode, timezone)))
             }
             ExprStruct::Hour(expr) => {
@@ -495,7 +502,10 @@ impl PhysicalPlanner {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema.clone())?;
                 let return_type = child.data_type(&input_schema)?;
                 let args = vec![child];
-                let scalar_def = ScalarFunctionDefinition::UDF(math::abs());
+                let eval_mode = Self::eval_mode_from_str(expr.eval_mode.as_str(), false)?;
+                let comet_abs =
+                    ScalarUDF::new_from_impl(CometAbsFunc::new(eval_mode, return_type.to_string()));
+                let scalar_def = ScalarFunctionDefinition::UDF(Arc::new(comet_abs));
 
                 let expr =
                     ScalarFunctionExpr::new("abs", scalar_def, args, return_type, None, false);

core/src/execution/proto/expr.proto

@@ -473,6 +473,7 @@ message BitwiseNot {
 
 message Abs {
   Expr child = 1;
+  string eval_mode = 2;
 }
 
 message Subquery {

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -54,6 +54,13 @@ import org.apache.comet.shims.ShimQueryPlanSerde
  * An utility object for query plan and expression serialization.
  */
 object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim {
+
+  object ExecutionMode {
+    val ANSI = "ANSI"
+    val LEGACY = "LEGACY"
+    val TRY = "TRY"
+  }
+
   def emitWarning(reason: String): Unit = {
     logWarning(s"Comet native execution is disabled due to: $reason")
   }
@@ -691,7 +698,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
         case Cast(child, dt, timeZoneId, evalMode) =>
           val evalModeStr = if (evalMode.isInstanceOf[Boolean]) {
             // Spark 3.2 & 3.3 has ansiEnabled boolean
-            if (evalMode.asInstanceOf[Boolean]) "ANSI" else "LEGACY"
+            if (evalMode.asInstanceOf[Boolean]) ExecutionMode.ANSI else ExecutionMode.LEGACY
           } else {
             // Spark 3.4+ has EvalMode enum with values LEGACY, ANSI, and TRY
             evalMode.toString
@@ -1474,15 +1481,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
             None
           }
 
-        case Abs(child, _) =>
+        case Abs(child, failOnErr) =>
           val childExpr = exprToProtoInternal(child, inputs)
           if (childExpr.isDefined) {
-            val abs =
-              ExprOuterClass.Abs
-                .newBuilder()
-                .setChild(childExpr.get)
-                .build()
-            Some(Expr.newBuilder().setAbs(abs).build())
+            val evalModeStr = if (failOnErr) ExecutionMode.ANSI else ExecutionMode.LEGACY
+            val absBuilder = ExprOuterClass.Abs.newBuilder()
+            absBuilder.setChild(childExpr.get)
+            absBuilder.setEvalMode(evalModeStr)
+            Some(Expr.newBuilder().setAbs(absBuilder).build())
           } else {
             withInfo(expr, child)
             None

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -850,6 +850,34 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("abs Overflow ansi mode") {
+    val data: Seq[(Int, Int)] = Seq((Int.MaxValue, Int.MinValue))
+    withSQLConf(
+      SQLConf.ANSI_ENABLED.key -> "true",
+      CometConf.COMET_ANSI_MODE_ENABLED.key -> "true") {
+      withParquetTable(data, "tbl") {
+        checkSparkMaybeThrows(sql("select abs(_1), abs(_2) from tbl")) match {
+          case (Some(sparkExc), Some(cometExc)) =>
+            val cometErrorPattern =
+              """.+[ARITHMETIC_OVERFLOW].+overflow. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error.""".r
+            val sparkErrorPattern = ".*integer overflow.*".r
+            assert(cometErrorPattern.findFirstIn(cometExc.getMessage).isDefined)
+            assert(sparkErrorPattern.findFirstIn(sparkExc.getMessage).isDefined)
+          case _ => fail("Exception should be thrown")
+        }
+      }
+    }
+  }
+
+  test("abs Overflow legacy mode") {
+    val data: Seq[(Int, Int)] = Seq((Int.MaxValue, Int.MinValue), (1, -1))
+    withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") {
+      withParquetTable(data, "tbl") {
+        checkSparkAnswerAndOperator("select abs(_1), abs(_2) from tbl")
+      }
+    }
+  }
+
   test("ceil and floor") {
     Seq("true", "false").foreach { dictionary =>
       withSQLConf(