Skip to content

fix(spark): derive pmod's decimal result type from the declared arguments #24409

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jefffrey merged 7 commits into from
Sep 6, 2026
+228 −25
Merged

fix(spark): derive pmod's decimal result type from the declared arguments #24409

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
1f8de79
fix(spark): derive pmod's decimal result type from the declared argum…
amitvijapur Aug 16, 2026
e363e12
fix(spark): read pmod's result type from ScalarFunctionArgs
amitvijapur Aug 17, 2026
ff84872
fix(spark): use one_of signature and decide pmod's null result types …
amitvijapur Aug 22, 2026
d336e0e
fix(spark): apply review feedback on pmod
amitvijapur Aug 30, 2026
a263911
fix(spark): drop two more unnecessary comments
amitvijapur Aug 30, 2026
3f026a1
fix(spark): drop the leftover duplicate null-handling comment
amitvijapur Sep 5, 2026
ce97a59
Apply suggestion from @Jefffrey
Jefffrey Sep 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
187 changes: 162 additions & 25 deletions datafusion/spark/src/function/math/modulus.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,23 @@
// 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},
cmp::{eq, lt},
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
Expand Down Expand Up @@ -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<ColumnarValue> {
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,
)?));
}

Comment thread
Jefffrey marked this conversation as resolved.
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))
}

Expand Down Expand Up @@ -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,
),
}
}
}
Expand All @@ -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<ColumnarValue> {
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(),
)
}
}

Expand Down Expand Up @@ -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 =
Expand All @@ -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 =
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 =
Expand All @@ -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());
}

Expand All @@ -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 =
Expand All @@ -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());
}

Expand All @@ -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());
}

Expand All @@ -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
Expand All @@ -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());
}

Expand All @@ -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 =
Expand All @@ -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
Expand All @@ -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 =
Expand Down Expand Up @@ -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 =
Expand All @@ -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)
);
}
}
Loading