impl_ansi

Author: coderfender

native/core/src/execution/planner.rs

@@ -1086,7 +1086,8 @@ impl PhysicalPlanner {
             }
             _ => {
                 let data_type = return_type.map(to_arrow_datatype).unwrap();
-                if [EvalMode::Try, EvalMode::Ansi].contains(&eval_mode) && (data_type.is_numeric())
+                if [EvalMode::Try, EvalMode::Ansi].contains(&eval_mode)
+                    && (data_type.is_integer() || data_type.is_floating())
                 {
                     let op_str = match op {
                         DataFusionOperator::Plus => "checked_add",

native/spark-expr/src/comet_scalar_funcs.rs

@@ -136,10 +136,12 @@ pub fn create_comet_physical_fun_with_eval_mode(
             make_comet_scalar_udf!("decimal_div", spark_decimal_div, data_type)
         }
         "decimal_integral_div" => {
+            let is_ansi_mode = eval_mode == EvalMode::Ansi;
             make_comet_scalar_udf!(
                 "decimal_integral_div",
                 spark_decimal_integral_div,
-                data_type
+                data_type,
+                is_ansi_mode
             )
         }
         "checked_add" => {

native/spark-expr/src/math_funcs/div.rs

@@ -15,30 +15,33 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::divide_by_zero_error;
 use crate::math_funcs::utils::get_precision_scale;
 use arrow::array::{Array, Decimal128Array};
 use arrow::datatypes::{DataType, DECIMAL128_MAX_PRECISION};
+use arrow::error::ArrowError;
 use arrow::{
     array::{ArrayRef, AsArray},
     datatypes::Decimal128Type,
 };
 use datafusion::common::DataFusionError;
 use datafusion::physical_plan::ColumnarValue;
-use num::{BigInt, Signed, ToPrimitive};
+use num::{BigInt, Signed, ToPrimitive, Zero};
 use std::sync::Arc;
 
 pub fn spark_decimal_div(
     args: &[ColumnarValue],
     data_type: &DataType,
 ) -> Result<ColumnarValue, DataFusionError> {
-    spark_decimal_div_internal(args, data_type, false)
+    spark_decimal_div_internal(args, data_type, false, false)
 }
 
 pub fn spark_decimal_integral_div(
     args: &[ColumnarValue],
     data_type: &DataType,
+    is_ansi_enabled: bool,
 ) -> Result<ColumnarValue, DataFusionError> {
-    spark_decimal_div_internal(args, data_type, true)
+    spark_decimal_div_internal(args, data_type, true, is_ansi_enabled)
 }
 
 // Let Decimal(p3, s3) as return type i.e. Decimal(p1, s1) / Decimal(p2, s2) = Decimal(p3, s3).
@@ -50,6 +53,7 @@ fn spark_decimal_div_internal(
     args: &[ColumnarValue],
     data_type: &DataType,
     is_integral_div: bool,
+    is_ansi_enabled: bool,
 ) -> Result<ColumnarValue, DataFusionError> {
     let left = &args[0];
     let right = &args[1];
@@ -96,9 +100,12 @@ fn spark_decimal_div_internal(
     } else {
         let l_mul = 10_i128.pow(l_exp);
         let r_mul = 10_i128.pow(r_exp);
-        arrow::compute::kernels::arity::binary(left, right, |l, r| {
+        arrow::compute::kernels::arity::try_binary(left, right, |l, r| {
             let l = l * l_mul;
             let r = r * r_mul;
+            if is_ansi_enabled && r.is_zero() {
+                return Err(ArrowError::ComputeError(divide_by_zero_error().to_string()));
+            }
             let div = if r == 0 { 0 } else { l / r };
             let res = if is_integral_div {
                 div
@@ -107,7 +114,7 @@ fn spark_decimal_div_internal(
             } else {
                 div + 5
             } / 10;
-            res.to_i128().unwrap_or(i128::MAX)
+            Ok(res.to_i128().unwrap_or(i128::MAX))
         })?
     };
     let result = result.with_data_type(DataType::Decimal128(p3, s3));

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

@@ -180,14 +180,6 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase {
 
 object CometIntegralDivide extends CometExpressionSerde[IntegralDivide] with MathBase {
 
-  override def getSupportLevel(expr: IntegralDivide): SupportLevel = {
-    if (expr.evalMode == EvalMode.ANSI) {
-      Incompatible(Some("ANSI mode is not supported"))
-    } else {
-      Compatible(None)
-    }
-  }
-
   override def convert(
       expr: IntegralDivide,
       inputs: Seq[Attribute],
@@ -206,7 +198,8 @@ object CometIntegralDivide extends CometExpressionSerde[IntegralDivide] with Mat
       if (expr.right.dataType.isInstanceOf[DecimalType]) expr.right
       else Cast(expr.right, DecimalType(19, 0))
 
-    val rightExpr = nullIfWhenPrimitive(expr.right)
+    val rightExpr =
+      if (expr.evalMode != EvalMode.ANSI) nullIfWhenPrimitive(expr.right) else expr.right
 
     val dataType = (left.dataType, right.dataType) match {
       case (l: DecimalType, r: DecimalType) =>