diff --git a/datafusion/spark/src/function/math/modulus.rs b/datafusion/spark/src/function/math/modulus.rs index c37513c12cd6b..92f037dfff558 100644 --- a/datafusion/spark/src/function/math/modulus.rs +++ b/datafusion/spark/src/function/math/modulus.rs @@ -15,7 +15,10 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use arrow::array::{ArrayRef, BooleanArray, Scalar, new_null_array}; +use arrow::compute::kernels::cast::{CastOptions, cast_with_options}; use arrow::compute::kernels::numeric::add; use arrow::compute::kernels::{ boolean::{and, is_not_null, or}, @@ -23,11 +26,12 @@ use arrow::compute::kernels::{ numeric::{neg, rem}, zip::zip, }; -use arrow::datatypes::DataType; +use arrow::datatypes::{DECIMAL128_MAX_PRECISION, DataType}; use arrow::error::ArrowError; -use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err}; +use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err, exec_err}; use datafusion_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + TypeSignatureClass, Volatility, binary::decimal_coercion, }; /// Returns a one element array holding negative zero, for the floating point @@ -101,23 +105,87 @@ pub fn spark_mod( Ok(ColumnarValue::Array(result)) } +/// Spark derives the decimal result type of `pmod` with `Pmod.resultDecimalType`, +/// which follows the `Remainder` rule: +/// +/// ```text +/// scale = max(s1, s2) +/// precision = min(p1 - s1, p2 - s2) + scale +/// ``` +/// +/// The rule is applied to the input argument types. +fn pmod_decimal_result_type(p1: u8, s1: i8, p2: u8, s2: i8) -> DataType { + let scale = s1.max(s2); + let whole_digits = (i32::from(p1) - i32::from(s1)).min(i32::from(p2) - i32::from(s2)); + let precision = + (whole_digits + i32::from(scale)).clamp(1, i32::from(DECIMAL128_MAX_PRECISION)); + DataType::Decimal128(precision as u8, scale) +} + /// Spark-compatible `pmod` function /// In ANSI mode, division by zero throws an error. /// In legacy mode, division by zero returns NULL (Spark behavior). pub fn spark_pmod( args: &[ColumnarValue], enable_ansi_mode: bool, + result_type: &DataType, ) -> Result { assert_eq_or_internal_err!(args.len(), 2, "pmod expects exactly two arguments"); let args = ColumnarValue::values_to_arrays(args)?; - let left = &args[0]; - let right = &args[1]; + + // Need to handle nulls separately as they are pass through by the signature + if args.iter().any(|arg| arg.data_type() == &DataType::Null) { + return Ok(ColumnarValue::Scalar(ScalarValue::try_new_null( + result_type, + )?)); + } + + let (left, right): (ArrayRef, ArrayRef) = + if args[0].data_type() == args[1].data_type() { + (Arc::clone(&args[0]), Arc::clone(&args[1])) + } else { + let Some(computation_type) = + decimal_coercion(args[0].data_type(), args[1].data_type()) + else { + return exec_err!( + "pmod does not support ({}, {})", + args[0].data_type(), + args[1].data_type() + ); + }; + let widen = CastOptions { + safe: false, + ..Default::default() + }; + ( + cast_with_options(&args[0], &computation_type, &widen)?, + cast_with_options(&args[1], &computation_type, &widen)?, + ) + }; + + let left = &left; + let right = &right; let zero = ScalarValue::new_zero(left.data_type())?.to_array_of_size(left.len())?; let result = try_rem(left, right, enable_ansi_mode)?; let neg = lt(&result, &zero)?; let plus = zip(&neg, right, &zero)?; let result = add(&plus, &result)?; let result = try_rem(&result, right, enable_ansi_mode)?; + + // The remainder is bounded by the divisor, but the result type only carries + // `min(p1 - s1, p2 - s2)` integer digits, so a remainder approaching a + // divisor wider than the dividend does not always fit. Spark wraps decimal + // arithmetic in `CheckOverflow(nullOnOverflow = !ansiEnabled)`, so an + // overflow here is NULL in legacy mode and an error under ANSI. + let result = if result.data_type() == result_type { + result + } else { + let narrow = CastOptions { + safe: !enable_ansi_mode, + ..Default::default() + }; + cast_with_options(&result, result_type, &narrow)? + }; Ok(ColumnarValue::Array(result)) } @@ -182,7 +250,19 @@ impl Default for SparkPmod { impl SparkPmod { pub fn new() -> Self { Self { - signature: Signature::numeric(2, Volatility::Immutable), + signature: Signature::one_of( + vec![ + // A decimal pair must reach `return_type` with the + // precision and scale as written, since Spark defines + // `Pmod.resultDecimalType` on the declared arguments. + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Decimal), + Coercion::new_exact(TypeSignatureClass::Decimal), + ]), + TypeSignature::Numeric(2), + ], + Volatility::Immutable, + ), } } } @@ -203,13 +283,24 @@ impl ScalarUDFImpl for SparkPmod { "pmod expects exactly two arguments" ); - // Return the same type as the first argument for simplicity - // Arrow's rem function handles type promotion internally - Ok(arg_types[0].clone()) + match (&arg_types[0], &arg_types[1]) { + (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => { + Ok(pmod_decimal_result_type(*p1, *s1, *p2, *s2)) + } + // Need to handle nulls explicitly, see: https://github.com/apache/datafusion/issues/19458 + // We align with the behaviour of `mod` + (DataType::Null, DataType::Null) => Ok(DataType::Float64), + (DataType::Null, other) | (other, DataType::Null) => Ok(other.clone()), + _ => Ok(arg_types[0].clone()), + } } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - spark_pmod(&args.args, args.config_options.execution.enable_ansi_mode) + spark_pmod( + &args.args, + args.config_options.execution.enable_ansi_mode, + args.return_type(), + ) } } @@ -558,7 +649,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_int32 = @@ -581,7 +673,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_int64 = @@ -625,7 +718,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_float64 = result_array @@ -683,7 +777,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_float32 = result_array @@ -718,7 +813,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_int32 = @@ -737,7 +833,8 @@ mod test { let left = Int32Array::from(vec![Some(10)]); let left_value = ColumnarValue::Array(Arc::new(left)); - let result = spark_pmod(&[left_value], false); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value], false, &return_type); assert!(result.is_err()); } @@ -750,7 +847,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_int32 = @@ -772,7 +870,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], true); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], true, &return_type); assert!(result.is_err()); } @@ -788,7 +887,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], true); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], true, &return_type); assert!(result.is_err()); } @@ -802,7 +902,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_float64 = result_array @@ -825,7 +926,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], true); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], true, &return_type); assert!(result.is_err()); } @@ -840,7 +942,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], true).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], true, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_int32 = @@ -858,7 +961,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], true).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], true, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_float64 = result_array @@ -881,7 +985,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_int32 = @@ -919,7 +1024,8 @@ mod test { let left_value = ColumnarValue::Array(Arc::new(left)); let right_value = ColumnarValue::Array(Arc::new(right)); - let result = spark_pmod(&[left_value, right_value], false).unwrap(); + let return_type = left_value.data_type(); + let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap(); if let ColumnarValue::Array(result_array) = result { let result_int32 = @@ -935,4 +1041,35 @@ mod test { panic!("Expected array result"); } } + + /// Spark's `Pmod.resultDecimalType`: scale = max(s1, s2), + /// precision = min(p1 - s1, p2 - s2) + scale. + #[test] + fn test_pmod_decimal_result_type() { + // Equal scales: the narrower argument decides the precision. + assert_eq!( + pmod_decimal_result_type(3, 1, 2, 1), + DataType::Decimal128(2, 1) + ); + // The divisor is wider, so the dividend bounds the result. + assert_eq!( + pmod_decimal_result_type(5, 2, 4, 1), + DataType::Decimal128(5, 2) + ); + // Differing scales: the wider scale wins. + assert_eq!( + pmod_decimal_result_type(6, 3, 4, 1), + DataType::Decimal128(6, 3) + ); + // Integral decimals keep a zero scale. + assert_eq!( + pmod_decimal_result_type(10, 0, 5, 0), + DataType::Decimal128(5, 0) + ); + // The result never exceeds the maximum precision arrow can represent. + assert_eq!( + pmod_decimal_result_type(38, 0, 38, 38), + DataType::Decimal128(38, 38) + ); + } } diff --git a/datafusion/sqllogictest/test_files/spark/math/pmod.slt b/datafusion/sqllogictest/test_files/spark/math/pmod.slt index 1165fdfba3f7c..7f07e4666e8ba 100644 --- a/datafusion/sqllogictest/test_files/spark/math/pmod.slt +++ b/datafusion/sqllogictest/test_files/spark/math/pmod.slt @@ -98,6 +98,10 @@ SELECT pmod(10.5::float8, 0.0::float8); statement error DataFusion error: Arrow error: Divide by zero error SELECT pmod(10.5::float8, -0.0::float8); +# A decimal result that does not fit raises instead of returning NULL +statement error DataFusion error: Arrow error: Invalid argument error: 9999\.8 is too large to store in a Decimal128 of precision 3 +SELECT pmod(-0.1::decimal(3,1), 9999.9::decimal(5,1)); + # A NULL dividend short-circuits to NULL before the divisor is validated query I SELECT pmod(NULL::int, 0::int) as pmod_null_dividend_ansi; @@ -128,6 +132,21 @@ SELECT pmod(NULL::int, NULL::int) as pmod_null_3; ---- NULL +query T +SELECT arrow_typeof(pmod(NULL, NULL)); +---- +Float64 + +query T +SELECT arrow_typeof(pmod(2.5::decimal(3,1), NULL)); +---- +Decimal128(3, 1) + +# An untyped NULL beside a typed non-decimal argument takes the Numeric path, +# which cannot coerce the pair. `mod` rejects it the same way. +statement error DataFusion error: Error during planning: Internal error: Function 'pmod' failed to match any signature +SELECT pmod(NULL, 3::int); + # PMOD tests with large integers query I SELECT pmod(100::int, 30::int) as pmod_large_1; @@ -305,6 +324,53 @@ SELECT pmod(-10.0::decimal(3,1), 3.0::decimal(2,1)) as pmod_decimal_4; ---- 2 +# The decimal result type follows Spark's Pmod.resultDecimalType, which applies +# the Remainder rule to the declared argument types: +# scale = max(s1, s2) +# precision = min(p1 - s1, p2 - s2) + scale +query T +SELECT arrow_typeof(pmod(2.5::decimal(3,1), 1.2::decimal(2,1))); +---- +Decimal128(2, 1) + +# The divisor bounds the result, so the narrower argument decides the precision +query T +SELECT arrow_typeof(pmod(10.0::decimal(5,2), 3.0::decimal(4,1))); +---- +Decimal128(5, 2) + +# Differing scales: the wider scale wins +query T +SELECT arrow_typeof(pmod(1.234::decimal(6,3), 2.1::decimal(4,1))); +---- +Decimal128(6, 3) + +# The dividend does not fit the result type, so it must not be narrowed before +# the remainder is taken +query T +SELECT arrow_typeof(pmod(99.9::decimal(3,1), 2.5::decimal(2,1))); +---- +Decimal128(2, 1) + +query R +SELECT pmod(99.9::decimal(3,1), 2.5::decimal(2,1)); +---- +2.4 + +# A remainder is bounded by the divisor, not by the result type, so a divisor +# wider than the dividend can produce a value the result type cannot hold. Spark +# wraps decimal arithmetic in CheckOverflow(nullOnOverflow = !ansiEnabled), so +# this is NULL in legacy mode; the ANSI error is asserted in the ANSI block above. +query T +SELECT arrow_typeof(pmod(-0.1::decimal(3,1), 9999.9::decimal(5,1))); +---- +Decimal128(3, 1) + +query R +SELECT pmod(-0.1::decimal(3,1), 9999.9::decimal(5,1)); +---- +NULL + # PMOD tests with different integer types query I SELECT pmod(10::int8, 3::int8) as pmod_int8_1;