From 1f8de79ccddca1ba7743f132bbfcac00526e92cb Mon Sep 17 00:00:00 2001 From: Amit Vijapur Date: Sun, 16 Aug 2026 13:32:39 +0800 Subject: [PATCH 1/3] fix(spark): derive pmod's decimal result type from the declared arguments Spark derives pmod's decimal result with `Pmod.resultDecimalType`, applying the `Remainder` rule to the declared argument types: scale = max(s1, s2) precision = min(p1 - s1, p2 - s2) + scale `SparkPmod` used `Signature::numeric`, which collapses both arguments to a common decimal before `return_type` runs. The two precisions it saw were already equal, so the rule degenerated to the input precision and `pmod(decimal(3,1), decimal(2,1))` reported `Decimal128(3, 1)` where Spark reports `decimal(2,1)`. Leave decimal arguments intact through coercion, as `try_sum` already does, and apply the rule in `return_type`. Every other argument combination keeps the coercion `Signature::numeric` performed, null handling included, so only the decimal pair changes behaviour. The result type is narrower than the dividend, so the operands cannot be cast to it up front without overflowing it; `spark_pmod` widens them to a common computation type instead and narrows the result afterwards. The remainder is bounded by the divisor rather than by the result type, so that narrowing can overflow when the divisor is wider than the dividend. Spark wraps decimal arithmetic in `CheckOverflow(nullOnOverflow = !ansiEnabled)`, so the narrowing cast returns NULL in legacy mode and raises under ANSI, and the widening cast raises in either mode because the computation type always fits both operands. Values that fit the Spark result type are unchanged; values that do not were previously returned at the wider type and are now NULL or an error, which is the reported-type bug itself rather than a separate behaviour change. Github-Issue:#23895 --- datafusion/spark/src/function/math/modulus.rs | 191 +++++++++++++++++- .../test_files/spark/math/pmod.slt | 68 +++++++ 2 files changed, 251 insertions(+), 8 deletions(-) diff --git a/datafusion/spark/src/function/math/modulus.rs b/datafusion/spark/src/function/math/modulus.rs index c37513c12cd6b..e8f48351ba331 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,15 @@ 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::types::NativeType; +use datafusion_common::{ + Result, ScalarValue, assert_eq_or_internal_err, exec_err, plan_err, +}; use datafusion_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + binary::binary_numeric_coercion, }; /// Returns a one element array holding negative zero, for the floating point @@ -101,6 +108,76 @@ 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 *declared* argument types. Collapsing both +/// arguments to a common decimal first would make the two precisions equal and +/// the rule would degenerate to the input precision, which is why +/// [`SparkPmod::coerce_types`] leaves decimal arguments intact. +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) +} + +/// The type `pmod` computes in, which is not always the type it returns. +/// +/// Spark's result type is narrower than the dividend, so the operands cannot be +/// cast to it before the remainder is taken without overflowing the dividend. +/// The computation therefore runs in a common type wide enough for both, and +/// the result is narrowed afterwards. +fn pmod_computation_type(lhs: &DataType, rhs: &DataType) -> Result { + match binary_numeric_coercion(lhs, rhs) { + Some(computation_type) => Ok(computation_type), + None => exec_err!("pmod does not support ({lhs}, {rhs})"), + } +} + +/// The coercion `Signature::numeric` applied before `pmod` moved to +/// [`Signature::user_defined`], reproduced so that only the decimal pair below +/// changes behaviour. +/// +/// A null argument is skipped rather than coerced, and a call still typed null +/// afterwards falls back to `Float64`; both match `TypeSignature::Numeric` in +/// `datafusion_expr::type_coercion::functions`. +fn pmod_numeric_coercion(lhs: &DataType, rhs: &DataType) -> Result> { + let mut valid_type = lhs.clone(); + + let rhs_native: NativeType = rhs.into(); + if rhs_native != NativeType::Null { + if !rhs_native.is_numeric() { + return plan_err!( + "Function 'pmod' expects Numeric but received {rhs_native}" + ); + } + match binary_numeric_coercion(&valid_type, rhs) { + Some(coerced_type) => valid_type = coerced_type, + None => { + return plan_err!( + "For function 'pmod' {valid_type} and {rhs} are not coercible to a common numeric type" + ); + } + } + } + + let valid_native: NativeType = valid_type.clone().into(); + if valid_native == NativeType::Null { + valid_type = DataType::Float64; + } else if !valid_native.is_numeric() { + return plan_err!("Function 'pmod' expects Numeric but received {valid_native}"); + } + + Ok(vec![valid_type.clone(), valid_type]) +} + /// Spark-compatible `pmod` function /// In ANSI mode, division by zero throws an error. /// In legacy mode, division by zero returns NULL (Spark behavior). @@ -110,14 +187,59 @@ pub fn spark_pmod( ) -> 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]; + + // Decimal arguments reach here with their declared types intact, so the + // Spark result type is derived before they are widened for the computation. + let result_type = match (args[0].data_type(), args[1].data_type()) { + (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => { + Some(pmod_decimal_result_type(*p1, *s1, *p2, *s2)) + } + _ => None, + }; + + let (left, right): (ArrayRef, ArrayRef) = + if args[0].data_type() == args[1].data_type() { + (Arc::clone(&args[0]), Arc::clone(&args[1])) + } else { + let computation_type = + pmod_computation_type(args[0].data_type(), args[1].data_type())?; + // The computation type is wide enough for both operands by + // construction, so widening must not silently null on overflow the + // way arrow's default (`safe: true`) cast would. + 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 = match result_type { + Some(result_type) if result.data_type() != &result_type => { + let narrow = CastOptions { + safe: !enable_ansi_mode, + ..Default::default() + }; + cast_with_options(&result, &result_type, &narrow)? + } + _ => result, + }; Ok(ColumnarValue::Array(result)) } @@ -182,7 +304,7 @@ impl Default for SparkPmod { impl SparkPmod { pub fn new() -> Self { Self { - signature: Signature::numeric(2, Volatility::Immutable), + signature: Signature::user_defined(Volatility::Immutable), } } } @@ -203,9 +325,31 @@ 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)) + } + // Arrow's rem function handles type promotion for the rest + _ => Ok(arg_types[0].clone()), + } + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.len() != 2 { + return plan_err!( + "Function 'pmod' expects 2 arguments but received {}", + arg_types.len() + ); + } + + match (&arg_types[0], &arg_types[1]) { + // Spark applies resultDecimalType to the declared argument types, so + // these are left alone; spark_pmod widens them for the computation. + (DataType::Decimal128(_, _), DataType::Decimal128(_, _)) => { + Ok(arg_types.to_vec()) + } + (lhs, rhs) => pmod_numeric_coercion(lhs, rhs), + } } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -935,4 +1079,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..e58d8542c6533 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; @@ -305,6 +309,70 @@ 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 and an error under ANSI (asserted further down). +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 + +# An untyped NULL pair falls back to Float64, as it did under Signature::numeric +query T +SELECT arrow_typeof(pmod(NULL, NULL)); +---- +Float64 + +query R +SELECT pmod(NULL, NULL); +---- +NULL + +# A null divisor takes the dividend's type +query T +SELECT arrow_typeof(pmod(2.5::decimal(3,1), NULL)); +---- +Decimal128(3, 1) + # PMOD tests with different integer types query I SELECT pmod(10::int8, 3::int8) as pmod_int8_1; From e363e12bb5b7baacba9d76ddd032bdf344b320c7 Mon Sep 17 00:00:00 2001 From: Amit Vijapur Date: Mon, 17 Aug 2026 23:32:01 +0800 Subject: [PATCH 2/3] fix(spark): read pmod's result type from ScalarFunctionArgs `spark_pmod` re-derived the Spark decimal result type from the argument arrays, duplicating the rule `return_type` had already applied. Pass the computed type in instead, so it is derived in exactly one place. Per review feedback on #24409. --- datafusion/spark/src/function/math/modulus.rs | 80 +++++++++++-------- 1 file changed, 46 insertions(+), 34 deletions(-) diff --git a/datafusion/spark/src/function/math/modulus.rs b/datafusion/spark/src/function/math/modulus.rs index e8f48351ba331..993b83b002469 100644 --- a/datafusion/spark/src/function/math/modulus.rs +++ b/datafusion/spark/src/function/math/modulus.rs @@ -184,19 +184,11 @@ fn pmod_numeric_coercion(lhs: &DataType, rhs: &DataType) -> Result 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)?; - // Decimal arguments reach here with their declared types intact, so the - // Spark result type is derived before they are widened for the computation. - let result_type = match (args[0].data_type(), args[1].data_type()) { - (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => { - Some(pmod_decimal_result_type(*p1, *s1, *p2, *s2)) - } - _ => None, - }; - let (left, right): (ArrayRef, ArrayRef) = if args[0].data_type() == args[1].data_type() { (Arc::clone(&args[0]), Arc::clone(&args[1])) @@ -230,15 +222,14 @@ pub fn spark_pmod( // 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 = match result_type { - Some(result_type) if result.data_type() != &result_type => { - let narrow = CastOptions { - safe: !enable_ansi_mode, - ..Default::default() - }; - cast_with_options(&result, &result_type, &narrow)? - } - _ => result, + 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)) } @@ -353,7 +344,13 @@ impl ScalarUDFImpl for SparkPmod { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - spark_pmod(&args.args, args.config_options.execution.enable_ansi_mode) + // The Spark result type was already derived in `return_type`, so it is + // read back here rather than recomputed from the argument arrays. + spark_pmod( + &args.args, + args.config_options.execution.enable_ansi_mode, + args.return_type(), + ) } } @@ -702,7 +699,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 = @@ -725,7 +723,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 = @@ -769,7 +768,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 @@ -827,7 +827,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 @@ -862,7 +863,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 = @@ -881,7 +883,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()); } @@ -894,7 +897,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 = @@ -916,7 +920,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()); } @@ -932,7 +937,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()); } @@ -946,7 +952,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 @@ -969,7 +976,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()); } @@ -984,7 +992,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 = @@ -1002,7 +1011,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 @@ -1025,7 +1035,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 = @@ -1063,7 +1074,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 = From ff84872da5243c805a2ee8f1ceef9c6fe23470a7 Mon Sep 17 00:00:00 2001 From: Amit Vijapur Date: Sat, 22 Aug 2026 17:24:48 +0800 Subject: [PATCH 3/3] fix(spark): use one_of signature and decide pmod's null result types explicitly Replaces the hand-written `coerce_types` with the `one_of` signature suggested in review: Coercible([Decimal, Decimal]) Numeric(2) `Coercible` matches a null argument and passes it through uncoerced (#19458), so a decimal pair still reaches `return_type` with its declared precision and scale, while everything else falls to `Numeric(2)`. That drops `pmod_numeric_coercion`, which existed only to reproduce `Numeric`'s fold. The null cases the coercion API cannot express are handled in `return_type` instead: two untyped nulls answer `Float64`, and one untyped null keeps the other side's type. Both match what `mod` returns, verified directly: SELECT arrow_typeof(mod(NULL, NULL)); -- Float64 SELECT arrow_typeof(mod(2.5::decimal(3,1), NULL)); -- Decimal128(3, 1) Because the null is no longer coerced away, `spark_pmod` can now receive `DataType::Null` arrays, so it short-circuits to a null array of the result type rather than failing to build a zero scalar. `pmod(NULL, 3::int)` does not plan under `Numeric(2)`. That matches both `mod` today and `pmod` before this PR, which used `Signature::numeric(2)`, so it is covered by a `statement error` test rather than treated as a regression. --- datafusion/spark/src/function/math/modulus.rs | 95 +++++++------------ .../test_files/spark/math/pmod.slt | 29 ++++++ 2 files changed, 62 insertions(+), 62 deletions(-) diff --git a/datafusion/spark/src/function/math/modulus.rs b/datafusion/spark/src/function/math/modulus.rs index 993b83b002469..110462c86afe5 100644 --- a/datafusion/spark/src/function/math/modulus.rs +++ b/datafusion/spark/src/function/math/modulus.rs @@ -28,13 +28,10 @@ use arrow::compute::kernels::{ }; use arrow::datatypes::{DECIMAL128_MAX_PRECISION, DataType}; use arrow::error::ArrowError; -use datafusion_common::types::NativeType; -use datafusion_common::{ - Result, ScalarValue, assert_eq_or_internal_err, exec_err, plan_err, -}; +use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err, exec_err}; use datafusion_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, - binary::binary_numeric_coercion, + Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + TypeSignatureClass, Volatility, binary::binary_numeric_coercion, }; /// Returns a one element array holding negative zero, for the floating point @@ -141,43 +138,6 @@ fn pmod_computation_type(lhs: &DataType, rhs: &DataType) -> Result { } } -/// The coercion `Signature::numeric` applied before `pmod` moved to -/// [`Signature::user_defined`], reproduced so that only the decimal pair below -/// changes behaviour. -/// -/// A null argument is skipped rather than coerced, and a call still typed null -/// afterwards falls back to `Float64`; both match `TypeSignature::Numeric` in -/// `datafusion_expr::type_coercion::functions`. -fn pmod_numeric_coercion(lhs: &DataType, rhs: &DataType) -> Result> { - let mut valid_type = lhs.clone(); - - let rhs_native: NativeType = rhs.into(); - if rhs_native != NativeType::Null { - if !rhs_native.is_numeric() { - return plan_err!( - "Function 'pmod' expects Numeric but received {rhs_native}" - ); - } - match binary_numeric_coercion(&valid_type, rhs) { - Some(coerced_type) => valid_type = coerced_type, - None => { - return plan_err!( - "For function 'pmod' {valid_type} and {rhs} are not coercible to a common numeric type" - ); - } - } - } - - let valid_native: NativeType = valid_type.clone().into(); - if valid_native == NativeType::Null { - valid_type = DataType::Float64; - } else if !valid_native.is_numeric() { - return plan_err!("Function 'pmod' expects Numeric but received {valid_native}"); - } - - Ok(vec![valid_type.clone(), valid_type]) -} - /// Spark-compatible `pmod` function /// In ANSI mode, division by zero throws an error. /// In legacy mode, division by zero returns NULL (Spark behavior). @@ -189,6 +149,16 @@ pub fn spark_pmod( assert_eq_or_internal_err!(args.len(), 2, "pmod expects exactly two arguments"); let args = ColumnarValue::values_to_arrays(args)?; + // A null argument is passed through uncoerced by `Coercible` (#19458), so + // it still carries `DataType::Null` here. Every operation below needs a + // concrete numeric type, and the answer is null regardless. + if args.iter().any(|arg| arg.data_type() == &DataType::Null) { + return Ok(ColumnarValue::Array(new_null_array( + result_type, + args[0].len(), + ))); + } + let (left, right): (ArrayRef, ArrayRef) = if args[0].data_type() == args[1].data_type() { (Arc::clone(&args[0]), Arc::clone(&args[1])) @@ -295,7 +265,19 @@ impl Default for SparkPmod { impl SparkPmod { pub fn new() -> Self { Self { - signature: Signature::user_defined(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, + ), } } } @@ -320,29 +302,18 @@ impl ScalarUDFImpl for SparkPmod { (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => { Ok(pmod_decimal_result_type(*p1, *s1, *p2, *s2)) } + // `Coercible` matches a null argument and passes it through + // uncoerced (#19458), so an untyped NULL reaches here rather than + // being folded by `Numeric(2)`. `mod` answers `Float64` for two + // untyped nulls and the other side's type when only one is null; + // `pmod` did too, so that behaviour is kept explicitly here. + (DataType::Null, DataType::Null) => Ok(DataType::Float64), + (DataType::Null, other) | (other, DataType::Null) => Ok(other.clone()), // Arrow's rem function handles type promotion for the rest _ => Ok(arg_types[0].clone()), } } - fn coerce_types(&self, arg_types: &[DataType]) -> Result> { - if arg_types.len() != 2 { - return plan_err!( - "Function 'pmod' expects 2 arguments but received {}", - arg_types.len() - ); - } - - match (&arg_types[0], &arg_types[1]) { - // Spark applies resultDecimalType to the declared argument types, so - // these are left alone; spark_pmod widens them for the computation. - (DataType::Decimal128(_, _), DataType::Decimal128(_, _)) => { - Ok(arg_types.to_vec()) - } - (lhs, rhs) => pmod_numeric_coercion(lhs, rhs), - } - } - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { // The Spark result type was already derived in `return_type`, so it is // read back here rather than recomputed from the argument arrays. diff --git a/datafusion/sqllogictest/test_files/spark/math/pmod.slt b/datafusion/sqllogictest/test_files/spark/math/pmod.slt index e58d8542c6533..4c38de5b81af6 100644 --- a/datafusion/sqllogictest/test_files/spark/math/pmod.slt +++ b/datafusion/sqllogictest/test_files/spark/math/pmod.slt @@ -132,6 +132,35 @@ SELECT pmod(NULL::int, NULL::int) as pmod_null_3; ---- NULL +# An untyped NULL matches the decimal signature and is passed through +# uncoerced (apache/datafusion#19458), so these types are decided explicitly +# rather than by coercion. `mod` answers Float64 for two untyped nulls and +# keeps the other side's type when only one is null; pmod matches it. +query T +SELECT arrow_typeof(pmod(NULL, NULL)); +---- +Float64 + +query R +SELECT pmod(NULL, NULL); +---- +NULL + +query T +SELECT arrow_typeof(pmod(2.5::decimal(3,1), NULL)); +---- +Decimal128(3, 1) + +query R +SELECT pmod(2.5::decimal(3,1), NULL); +---- +NULL + +# 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;