chore: Introduce ANSI support for remainder operation (#1971)

Author: rishvin

native/core/src/execution/planner.rs

@@ -17,7 +17,6 @@
 
 //! Converts Spark physical plan to DataFusion physical plan
 
-use super::expressions::EvalMode;
 use crate::execution::operators::CopyMode;
 use crate::execution::operators::FilterExec as CometFilterExec;
 use crate::{
@@ -62,8 +61,8 @@ use datafusion::{
     prelude::SessionContext,
 };
 use datafusion_comet_spark_expr::{
-    create_comet_physical_fun, create_negate_expr, BloomFilterAgg, BloomFilterMightContain,
-    SparkHour, SparkMinute, SparkSecond,
+    create_comet_physical_fun, create_modulo_expr, create_negate_expr, BloomFilterAgg,
+    BloomFilterMightContain, EvalMode, SparkHour, SparkMinute, SparkSecond,
 };
 
 use crate::execution::operators::ExecutionError::GeneralError;
@@ -268,13 +267,22 @@ impl PhysicalPlanner {
                     is_integral_div: true,
                 },
             ),
-            ExprStruct::Remainder(expr) => self.create_binary_expr(
-                expr.left.as_ref().unwrap(),
-                expr.right.as_ref().unwrap(),
-                expr.return_type.as_ref(),
-                DataFusionOperator::Modulo,
-                input_schema,
-            ),
+            ExprStruct::Remainder(expr) => {
+                let left =
+                    self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let right =
+                    self.create_expr(expr.right.as_ref().unwrap(), Arc::clone(&input_schema))?;
+
+                let result = create_modulo_expr(
+                    left,
+                    right,
+                    expr.return_type.as_ref().map(to_arrow_datatype).unwrap(),
+                    input_schema,
+                    expr.fail_on_error,
+                    &self.session_ctx.state(),
+                );
+                result.map_err(|e| GeneralError(e.to_string()))
+            }
             ExprStruct::Eq(expr) => {
                 let left =
                     self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
@@ -851,19 +859,13 @@ impl PhysicalPlanner {
             right.data_type(&input_schema),
         ) {
             (
-                DataFusionOperator::Plus
-                | DataFusionOperator::Minus
-                | DataFusionOperator::Multiply
-                | DataFusionOperator::Modulo,
+                DataFusionOperator::Plus | DataFusionOperator::Minus | DataFusionOperator::Multiply,
                 Ok(DataType::Decimal128(p1, s1)),
                 Ok(DataType::Decimal128(p2, s2)),
             ) if ((op == DataFusionOperator::Plus || op == DataFusionOperator::Minus)
                 && max(s1, s2) as u8 + max(p1 - s1 as u8, p2 - s2 as u8)
                     >= DECIMAL128_MAX_PRECISION)
-                || (op == DataFusionOperator::Multiply && p1 + p2 >= DECIMAL128_MAX_PRECISION)
-                || (op == DataFusionOperator::Modulo
-                    && max(s1, s2) as u8 + max(p1 - s1 as u8, p2 - s2 as u8)
-                        > DECIMAL128_MAX_PRECISION) =>
+                || (op == DataFusionOperator::Multiply && p1 + p2 >= DECIMAL128_MAX_PRECISION) =>
             {
                 let data_type = return_type.map(to_arrow_datatype).unwrap();
                 // For some Decimal128 operations, we need wider internal digits.
@@ -903,6 +905,7 @@ impl PhysicalPlanner {
                     func_name,
                     data_type.clone(),
                     &self.session_ctx.state(),
+                    None,
                 )?;
                 Ok(Arc::new(ScalarFunctionExpr::new(
                     func_name,
@@ -2305,8 +2308,12 @@ impl PhysicalPlanner {
                 }
             };
 
-        let fun_expr =
-            create_comet_physical_fun(fun_name, data_type.clone(), &self.session_ctx.state())?;
+        let fun_expr = create_comet_physical_fun(
+            fun_name,
+            data_type.clone(),
+            &self.session_ctx.state(),
+            None,
+        )?;
 
         let args = args
             .into_iter()

native/spark-expr/src/comet_scalar_funcs.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::hash_funcs::*;
+use crate::math_funcs::modulo_expr::spark_modulo;
 use crate::{
     spark_array_repeat, spark_ceil, spark_date_add, spark_date_sub, spark_decimal_div,
     spark_decimal_integral_div, spark_floor, spark_hex, spark_isnan, spark_make_decimal,
@@ -53,13 +54,23 @@ macro_rules! make_comet_scalar_udf {
         );
         Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
     }};
+    ($name:expr, $func:ident, without $data_type:ident, $fail_on_error:ident) => {{
+        let scalar_func = CometScalarFunction::new(
+            $name.to_string(),
+            Signature::variadic_any(Volatility::Immutable),
+            $data_type,
+            Arc::new(move |args| $func(args, $fail_on_error)),
+        );
+        Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
+    }};
 }
 
 /// Create a physical scalar function.
 pub fn create_comet_physical_fun(
     fun_name: &str,
     data_type: DataType,
     registry: &dyn FunctionRegistry,
+    fail_on_error: Option<bool>,
 ) -> Result<Arc<ScalarUDF>, DataFusionError> {
     match fun_name {
         "ceil" => {
@@ -144,6 +155,11 @@ pub fn create_comet_physical_fun(
             let func = Arc::new(spark_array_repeat);
             make_comet_scalar_udf!("array_repeat", func, without data_type)
         }
+        "spark_modulo" => {
+            let func = Arc::new(spark_modulo);
+            let fail_on_error = fail_on_error.unwrap_or(false);
+            make_comet_scalar_udf!("spark_modulo", func, without data_type, fail_on_error)
+        }
         _ => registry.udf(fun_name).map_err(|e| {
             DataFusionError::Execution(format!(
                 "Function {fun_name} not found in the registry: {e}",

native/spark-expr/src/error.rs

@@ -51,6 +51,9 @@ pub enum SparkError {
     #[error("[ARITHMETIC_OVERFLOW] {from_type} overflow. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")]
     ArithmeticOverflow { from_type: String },
 
+    #[error("[DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")]
+    DivideByZero,
+
     #[error("ArrowError: {0}.")]
     Arrow(ArrowError),
 

native/spark-expr/src/lib.rs

@@ -73,9 +73,9 @@ pub use error::{SparkError, SparkResult};
 pub use hash_funcs::*;
 pub use json_funcs::ToJson;
 pub use math_funcs::{
-    create_negate_expr, spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor,
-    spark_hex, spark_make_decimal, spark_round, spark_unhex, spark_unscaled_value, CheckOverflow,
-    NegativeExpr, NormalizeNaNAndZero,
+    create_modulo_expr, create_negate_expr, spark_ceil, spark_decimal_div,
+    spark_decimal_integral_div, spark_floor, spark_hex, spark_make_decimal, spark_round,
+    spark_unhex, spark_unscaled_value, CheckOverflow, NegativeExpr, NormalizeNaNAndZero,
 };
 pub use string_funcs::*;
 
@@ -103,3 +103,7 @@ pub(crate) fn arithmetic_overflow_error(from_type: &str) -> SparkError {
         from_type: from_type.to_string(),
     }
 }
+
+pub(crate) fn divide_by_zero_error() -> SparkError {
+    SparkError::DivideByZero
+}

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

@@ -20,6 +20,7 @@ mod div;
 mod floor;
 pub(crate) mod hex;
 pub mod internal;
+pub mod modulo_expr;
 mod negative;
 mod round;
 pub(crate) mod unhex;
@@ -31,6 +32,7 @@ pub use div::spark_decimal_integral_div;
 pub use floor::spark_floor;
 pub use hex::spark_hex;
 pub use internal::*;
+pub use modulo_expr::create_modulo_expr;
 pub use negative::{create_negate_expr, NegativeExpr};
 pub use round::spark_round;
 pub use unhex::spark_unhex;