You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by mb...@apache.org on 2025/07/11 18:18:26 UTC
(datafusion-comet) branch main updated: chore: Introduce ANSI support for remainder operation (#1971)
This is an automated email from the ASF dual-hosted git repository.
mbutrovich pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 496cad9ce chore: Introduce ANSI support for remainder operation (#1971)
496cad9ce is described below
commit 496cad9ce89642f5f6b0977935e178dc886fd384
Author: Rishab Joshi <81...@users.noreply.github.com>
AuthorDate: Fri Jul 11 12:18:23 2025 -0600
chore: Introduce ANSI support for remainder operation (#1971)
---
native/core/src/execution/planner.rs | 47 ++-
native/spark-expr/src/comet_scalar_funcs.rs | 16 +
native/spark-expr/src/error.rs | 3 +
native/spark-expr/src/lib.rs | 10 +-
native/spark-expr/src/math_funcs/mod.rs | 2 +
native/spark-expr/src/math_funcs/modulo_expr.rs | 430 +++++++++++++++++++++
native/spark-expr/tests/spark_expr_reg.rs | 1 +
.../org/apache/comet/serde/QueryPlanSerde.scala | 4 +-
.../org/apache/comet/CometExpressionSuite.scala | 74 ++++
9 files changed, 561 insertions(+), 26 deletions(-)
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 7003b2651..327e57bcf 100644
--- a/native/core/src/execution/planner.rs
+++ b/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()
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs
index 6177ef498..d3f50b2fa 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/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,6 +54,15 @@ 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.
@@ -60,6 +70,7 @@ 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}",
diff --git a/native/spark-expr/src/error.rs b/native/spark-expr/src/error.rs
index 9f6df6460..4b00b70eb 100644
--- a/native/spark-expr/src/error.rs
+++ b/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),
diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs
index 1ebf9842e..4b29b6177 100644
--- a/native/spark-expr/src/lib.rs
+++ b/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
+}
diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs
index 03eb9a76c..e7e203de5 100644
--- a/native/spark-expr/src/math_funcs/mod.rs
+++ b/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;
diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs
new file mode 100644
index 000000000..16e8af414
--- /dev/null
+++ b/native/spark-expr/src/math_funcs/modulo_expr.rs
@@ -0,0 +1,430 @@
+// 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::{create_comet_physical_fun, IfExpr};
+use crate::{divide_by_zero_error, Cast, EvalMode, SparkCastOptions};
+use arrow::compute::kernels::numeric::rem;
+use arrow::datatypes::*;
+use datafusion::common::{exec_err, internal_err, DataFusionError, Result, ScalarValue};
+use datafusion::execution::FunctionRegistry;
+use datafusion::physical_expr::expressions::{lit, BinaryExpr};
+use datafusion::physical_expr::ScalarFunctionExpr;
+use datafusion::physical_expr_common::datum::{apply, apply_cmp_for_nested};
+use datafusion::{
+ logical_expr::{ColumnarValue, Operator},
+ physical_expr::PhysicalExpr,
+};
+use std::cmp::max;
+use std::sync::Arc;
+
+/// Spark-compliant modulo function. If `fail_on_error` is true, then this function computes modulo
+/// in ANSI mode and returns an error on division by zero, otherwise it returns `NULL` for such
+/// cases.
+pub fn spark_modulo(args: &[ColumnarValue], fail_on_error: bool) -> Result<ColumnarValue> {
+ if args.len() != 2 {
+ return exec_err!("modulo expects exactly two arguments");
+ }
+
+ let lhs = &args[0];
+ let rhs = &args[1];
+
+ let left_data_type = lhs.data_type();
+ let right_data_type = rhs.data_type();
+
+ if left_data_type.is_nested() {
+ if right_data_type != left_data_type {
+ return internal_err!("Type mismatch for spark modulo operation");
+ }
+ return apply_cmp_for_nested(Operator::Modulo, lhs, rhs);
+ }
+
+ match apply(lhs, rhs, rem) {
+ Ok(result) => Ok(result),
+ Err(e) if e.to_string().contains("Divide by zero") && fail_on_error => {
+ // Return Spark-compliant divide by zero error.
+ Err(divide_by_zero_error().into())
+ }
+ Err(e) => Err(e),
+ }
+}
+
+pub fn create_modulo_expr(
+ left: Arc<dyn PhysicalExpr>,
+ right: Arc<dyn PhysicalExpr>,
+ data_type: DataType,
+ input_schema: SchemaRef,
+ fail_on_error: bool,
+ registry: &dyn FunctionRegistry,
+) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
+ // For non-ANSI mode, wrap the right expression such that any zero value is replaced with `NULL`
+ // to prevent divide by zero error.
+ let right_non_ansi_safe = if !fail_on_error {
+ null_if_zero_primitive(right, &input_schema)?
+ } else {
+ right
+ };
+
+ // If the data type is `Decimal128` and the (scale + integral part) exceeds the maximum allowed
+ // for `Decimal128`, then cast both operands to `Decimal256` before creating the modulo scalar
+ // expression, otherwise, create the modulo scalar expression directly.
+ match (
+ left.data_type(&input_schema),
+ right_non_ansi_safe.data_type(&input_schema),
+ ) {
+ (Ok(DataType::Decimal128(p1, s1)), Ok(DataType::Decimal128(p2, s2)))
+ if max(s1, s2) as u8 + max(p1 - s1 as u8, p2 - s2 as u8) > DECIMAL128_MAX_PRECISION =>
+ {
+ let left_256 = Arc::new(Cast::new(
+ left,
+ DataType::Decimal256(p1, s1),
+ SparkCastOptions::new_without_timezone(EvalMode::Legacy, false),
+ ));
+ let right_256 = Arc::new(Cast::new(
+ right_non_ansi_safe,
+ DataType::Decimal256(p2, s2),
+ SparkCastOptions::new_without_timezone(EvalMode::Legacy, false),
+ ));
+
+ let modulo_scalar_func = create_modulo_scalar_function(
+ left_256,
+ right_256,
+ &data_type,
+ registry,
+ fail_on_error,
+ )?;
+
+ Ok(Arc::new(Cast::new(
+ modulo_scalar_func,
+ data_type,
+ SparkCastOptions::new_without_timezone(EvalMode::Legacy, false),
+ )))
+ }
+ _ => create_modulo_scalar_function(
+ left,
+ right_non_ansi_safe,
+ &data_type,
+ registry,
+ fail_on_error,
+ ),
+ }
+}
+
+fn null_if_zero_primitive(
+ expression: Arc<dyn PhysicalExpr>,
+ input_schema: &Schema,
+) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
+ let expr_data_type = expression.data_type(input_schema)?;
+
+ if is_primitive_datatype(&expr_data_type) {
+ let zero = match expr_data_type {
+ DataType::Int8 => ScalarValue::Int8(Some(0)),
+ DataType::Int16 => ScalarValue::Int16(Some(0)),
+ DataType::Int32 => ScalarValue::Int32(Some(0)),
+ DataType::Int64 => ScalarValue::Int64(Some(0)),
+ DataType::UInt8 => ScalarValue::UInt8(Some(0)),
+ DataType::UInt16 => ScalarValue::UInt16(Some(0)),
+ DataType::UInt32 => ScalarValue::UInt32(Some(0)),
+ DataType::UInt64 => ScalarValue::UInt64(Some(0)),
+ DataType::Float32 => ScalarValue::Float32(Some(0.0)),
+ DataType::Float64 => ScalarValue::Float64(Some(0.0)),
+ DataType::Decimal128(s, p) => ScalarValue::Decimal128(Some(0), s, p),
+ DataType::Decimal256(s, p) => ScalarValue::Decimal256(Some(i256::from(0)), s, p),
+ _ => return Ok(expression),
+ };
+
+ // Create an expression like - `if (eval(expr) == Literal(0)) then NULL else eval(expr)`.
+ // This expression evaluates to null for rows with zero values to prevent divide by zero
+ // error.
+ let eq_expr = Arc::new(BinaryExpr::new(
+ Arc::<dyn PhysicalExpr>::clone(&expression),
+ Operator::Eq,
+ lit(zero),
+ ));
+ let null_literal = lit(ScalarValue::try_new_null(&expr_data_type)?);
+ let if_expr = Arc::new(IfExpr::new(eq_expr, null_literal, expression));
+ Ok(if_expr)
+ } else {
+ Ok(expression)
+ }
+}
+
+fn is_primitive_datatype(dt: &DataType) -> bool {
+ matches!(
+ dt,
+ DataType::Int8
+ | DataType::Int16
+ | DataType::Int32
+ | DataType::Int64
+ | DataType::UInt8
+ | DataType::UInt16
+ | DataType::UInt32
+ | DataType::UInt64
+ | DataType::Float32
+ | DataType::Float64
+ | DataType::Decimal128(_, _)
+ | DataType::Decimal256(_, _)
+ )
+}
+
+fn create_modulo_scalar_function(
+ left: Arc<dyn PhysicalExpr>,
+ right: Arc<dyn PhysicalExpr>,
+ data_type: &DataType,
+ registry: &dyn FunctionRegistry,
+ fail_on_error: bool,
+) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
+ let func_name = "spark_modulo";
+ let modulo_expr =
+ create_comet_physical_fun(func_name, data_type.clone(), registry, Some(fail_on_error))?;
+ Ok(Arc::new(ScalarFunctionExpr::new(
+ func_name,
+ modulo_expr,
+ vec![left, right],
+ Arc::new(Field::new(func_name, data_type.clone(), true)),
+ )))
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{
+ Array, ArrayRef, Decimal128Array, Decimal128Builder, Int32Array, PrimitiveArray,
+ RecordBatch,
+ };
+ use datafusion::logical_expr::ColumnarValue;
+ use datafusion::physical_expr::expressions::{Column, Literal};
+ use datafusion::prelude::SessionContext;
+
+ fn with_fail_on_error<F: Fn(bool)>(test_fn: F) {
+ for fail_on_error in [true, false] {
+ test_fn(fail_on_error);
+ }
+ }
+
+ pub fn verify_result<T>(
+ expr: Arc<dyn PhysicalExpr>,
+ batch: RecordBatch,
+ should_fail: bool,
+ expected_result: Option<Arc<PrimitiveArray<T>>>,
+ ) where
+ T: ArrowPrimitiveType,
+ {
+ let actual_result = expr.evaluate(&batch);
+
+ if should_fail {
+ match actual_result {
+ Err(error) => {
+ assert!(
+ error
+ .to_string()
+ .contains("[DIVIDE_BY_ZERO] Division by zero"),
+ "Error message did not match. Actual message: {error}"
+ );
+ }
+ Ok(value) => {
+ panic!("Expected error, but got: {value:?}");
+ }
+ }
+ } else {
+ match (actual_result, expected_result) {
+ (Ok(ColumnarValue::Array(ref actual)), Some(expected)) => {
+ assert_eq!(actual.len(), expected.len(), "Array length mismatch");
+
+ let actual_arr = actual.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
+ let expected_arr = expected
+ .as_any()
+ .downcast_ref::<PrimitiveArray<T>>()
+ .unwrap();
+
+ for i in 0..actual_arr.len() {
+ assert_eq!(
+ actual_arr.is_null(i),
+ expected_arr.is_null(i),
+ "Nullity mismatch at index {i}"
+ );
+ if !actual_arr.is_null(i) {
+ let actual_value = actual_arr.value(i);
+ let expected_value = expected_arr.value(i);
+ assert_eq!(
+ actual_value, expected_value,
+ "Mismatch at index {i}, actual {actual_value:?}, expected {expected_value:?}"
+ );
+ }
+ }
+ }
+ (actual, expected) => {
+ panic!("Actual: {actual:?}, expected: {expected:?}");
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn test_modulo_basic_int() {
+ with_fail_on_error(|fail_on_error| {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Int32, false),
+ ]));
+
+ let a_array = Arc::new(Int32Array::from(vec![3, 2, i32::MIN]));
+ let b_array = Arc::new(Int32Array::from(vec![1, 5, -1]));
+ let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a_array, b_array]).unwrap();
+
+ let left_expr = Arc::new(Column::new("a", 0));
+ let right_expr = Arc::new(Column::new("b", 1));
+
+ let session_ctx = SessionContext::new();
+ let modulo_expr = create_modulo_expr(
+ left_expr,
+ right_expr,
+ DataType::Int32,
+ schema,
+ fail_on_error,
+ &session_ctx.state(),
+ )
+ .unwrap();
+
+ // This test case should not fail as there is no division by zero.
+ let should_fail = false;
+ let expected_result = Arc::new(Int32Array::from(vec![0, 2, 0]));
+ verify_result(modulo_expr, batch, should_fail, Some(expected_result));
+ })
+ }
+
+ #[test]
+ fn test_modulo_basic_decimal() {
+ with_fail_on_error(|fail_on_error| {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Decimal128(18, 4), false),
+ Field::new("b", DataType::Decimal128(18, 4), false),
+ ]));
+
+ let mut a_builder =
+ Decimal128Builder::with_capacity(2).with_data_type(DataType::Decimal128(18, 4));
+ a_builder.append_value(3000000000000000000);
+ a_builder.append_value(2000000000000000000);
+ let a_array: ArrayRef = Arc::new(a_builder.finish());
+
+ let mut b_builder =
+ Decimal128Builder::with_capacity(2).with_data_type(DataType::Decimal128(18, 4));
+ b_builder.append_value(1000000000000000000);
+ b_builder.append_value(5000000000000000000);
+ let b_array: ArrayRef = Arc::new(b_builder.finish());
+
+ let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a_array, b_array]).unwrap();
+
+ let left_expr = Arc::new(Column::new("a", 0));
+ let right_expr = Arc::new(Column::new("b", 1));
+
+ let session_ctx = SessionContext::new();
+ let modulo_expr = create_modulo_expr(
+ left_expr,
+ right_expr,
+ DataType::Decimal128(18, 4),
+ schema,
+ fail_on_error,
+ &session_ctx.state(),
+ )
+ .unwrap();
+
+ // This test case should not fail as there is no division by zero.
+ let should_fail = false;
+ let expected_result = Arc::new(Decimal128Array::from(vec![
+ Some(0),
+ Some(2000000000000000000),
+ ]));
+ verify_result(modulo_expr, batch, should_fail, Some(expected_result));
+ })
+ }
+
+ #[test]
+ fn test_modulo_divide_by_zero_int() {
+ with_fail_on_error(|fail_on_error| {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Int32, false),
+ ]));
+
+ let a_array = Arc::new(Int32Array::from(vec![3]));
+ let b_array = Arc::new(Int32Array::from(vec![0]));
+ let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a_array, b_array]).unwrap();
+
+ let left_expr = Arc::new(Column::new("a", 0));
+ let right_expr = Arc::new(Column::new("b", 1));
+
+ let session_ctx = SessionContext::new();
+ let modulo_expr = create_modulo_expr(
+ left_expr,
+ right_expr,
+ DataType::Int32,
+ schema,
+ fail_on_error,
+ &session_ctx.state(),
+ )
+ .unwrap();
+
+ // Expected result in non-ANSI mode.
+ let expected_result = Arc::new(Int32Array::from(vec![None]));
+ verify_result(modulo_expr, batch, fail_on_error, Some(expected_result));
+ })
+ }
+
+ #[test]
+ fn test_division_by_zero_with_complex_int_expr() {
+ with_fail_on_error(|fail_on_error| {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Int32, false),
+ Field::new("c", DataType::Int32, false),
+ ]));
+
+ let a_array = Arc::new(Int32Array::from(vec![3, 0]));
+ let b_array = Arc::new(Int32Array::from(vec![2, 4]));
+ let c_array = Arc::new(Int32Array::from(vec![4, 5]));
+ let batch =
+ RecordBatch::try_new(Arc::clone(&schema), vec![a_array, b_array, c_array]).unwrap();
+
+ let left_expr = Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("a", 0)),
+ Operator::Divide,
+ Arc::new(Column::new("b", 1)),
+ ));
+ let right_expr = Arc::new(BinaryExpr::new(
+ Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
+ Operator::Divide,
+ Arc::new(Column::new("c", 2)),
+ ));
+
+ // Computes modulo of (a / b) % (0 / c).
+ let session_ctx = SessionContext::new();
+ let modulo_expr = create_modulo_expr(
+ left_expr,
+ right_expr,
+ DataType::Int32,
+ schema,
+ fail_on_error,
+ &session_ctx.state(),
+ )
+ .unwrap();
+
+ // Expected result in non-ANSI mode.
+ let expected_result = Arc::new(Int32Array::from(vec![None, None]));
+ verify_result(modulo_expr, batch, fail_on_error, Some(expected_result));
+ })
+ }
+}
diff --git a/native/spark-expr/tests/spark_expr_reg.rs b/native/spark-expr/tests/spark_expr_reg.rs
index 6c94b458f..88e34ae2b 100644
--- a/native/spark-expr/tests/spark_expr_reg.rs
+++ b/native/spark-expr/tests/spark_expr_reg.rs
@@ -32,6 +32,7 @@ mod tests {
"xxhash64",
DataType::Int64,
&session_state,
+ None,
)?);
let ctx = SessionContext::new_with_state(session_state);
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 43c03d6b4..077faeb41 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -750,12 +750,10 @@ object QueryPlanSerde extends Logging with CometExprShim {
None
case rem @ Remainder(left, right, _) if supportedDataType(left.dataType) =>
- val rightExpr = nullIfWhenPrimitive(right)
-
createMathExpression(
expr,
left,
- rightExpr,
+ right,
inputs,
binding,
rem.dataType,
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index a30d25f87..6f29d4887 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -1868,6 +1868,80 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("remainder function") {
+ def withAnsiMode(enabled: Boolean)(f: => Unit): Unit = {
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> enabled.toString,
+ CometConf.COMET_ANSI_MODE_ENABLED.key -> enabled.toString,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true")(f)
+ }
+
+ def verifyResult(query: String): Unit = {
+ val expectedDivideByZeroError =
+ "[DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead."
+
+ checkSparkMaybeThrows(sql(query)) match {
+ case (Some(sparkException), Some(cometException)) =>
+ assert(sparkException.getMessage.contains(expectedDivideByZeroError))
+ assert(cometException.getMessage.contains(expectedDivideByZeroError))
+ case (None, None) => checkSparkAnswerAndOperator(sql(query))
+ case (None, Some(ex)) =>
+ fail("Comet threw an exception but Spark did not. Comet exception: " + ex.getMessage)
+ case (Some(sparkException), None) =>
+ fail(
+ "Spark threw an exception but Comet did not. Spark exception: " +
+ sparkException.getMessage)
+ }
+ }
+
+ val tableName = "remainder_table"
+ withTable(tableName) {
+ sql(s"""
+ |create table $tableName (
+ |a_int int,
+ |b_int int,
+ |a_float float,
+ |b_float float,
+ |a_decimal decimal(18,4),
+ |b_decimal decimal(18,4)
+ |) using parquet
+ """.stripMargin)
+
+ val minInt = Int.MinValue
+
+ sql(s"""
+ |insert into $tableName values
+ |(3, 0, 3.0, 0.0, cast(3 as decimal(18,4)), cast(0 as decimal(18,4))),
+ |(10, 3, 10.5, 3.0, cast(10 as decimal(18,4)), cast(3 as decimal(18,4))),
+ |(-10, 3, -10.5, 3.0, cast(-10 as decimal(18,4)), cast(3 as decimal(18,4))),
+ |($minInt, -1, $minInt, -1.0, cast($minInt as decimal(18,4)), cast(-1 as decimal(18,4))),
+ |(null, 3, null, 3.0, null, cast(3 as decimal(18,4))),
+ |(3, null, 3.0, null, cast(3 as decimal(18,4)), null)
+ """.stripMargin)
+
+ Seq(true, false).foreach(enabled => {
+ withAnsiMode(enabled = enabled) {
+ verifyResult(s"""
+ |select
+ |a_int, b_int, a_float, b_float, a_decimal, b_decimal,
+ |a_int % b_int as int_int_modulo,
+ |a_int % b_float as int_float_modulo,
+ |mod(a_int, b_int) as int_int_mod,
+ |mod(a_int, b_float) as int_float_mod,
+ |a_float % b_float as float_float_modulo,
+ |a_float % b_decimal as float_decimal_modulo,
+ |mod(a_float, b_float) as float_float_mod,
+ |mod(a_float, b_decimal) as float_decimal_mod,
+ |a_decimal % b_decimal as decimal_decimal_modulo,
+ |mod(a_decimal, b_decimal) as decimal_decimal_mod
+ |from $tableName
+ """.stripMargin)
+ }
+ })
+ }
+ }
+
test("hash functions with random input") {
val dataGen = DataGenerator.DEFAULT
// sufficient number of rows to create dictionary encoded ArrowArray.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org