Skip to content

Commit d12d695

Browse files
osipovartemosipovartem
authored
Merge branch 'main' into upstream-slt-decimal-clippy
2 parents c25d58a + 262936e commit d12d695

2 files changed

Lines changed: 228 additions & 25 deletions

File tree

‎datafusion/spark/src/function/math/modulus.rs‎

Lines changed: 162 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,23 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use std::sync::Arc;
19+
1820
use arrow::array::{ArrayRef, BooleanArray, Scalar, new_null_array};
21+
use arrow::compute::kernels::cast::{CastOptions, cast_with_options};
1922
use arrow::compute::kernels::numeric::add;
2023
use arrow::compute::kernels::{
2124
boolean::{and, is_not_null, or},
2225
cmp::{eq, lt},
2326
numeric::{neg, rem},
2427
zip::zip,
2528
};
26-
use arrow::datatypes::DataType;
29+
use arrow::datatypes::{DECIMAL128_MAX_PRECISION, DataType};
2730
use arrow::error::ArrowError;
28-
use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err};
31+
use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err, exec_err};
2932
use datafusion_expr::{
30-
ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
33+
Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
34+
TypeSignatureClass, Volatility, binary::decimal_coercion,
3135
};
3236

3337
/// Returns a one element array holding negative zero, for the floating point
@@ -101,23 +105,87 @@ pub fn spark_mod(
101105
Ok(ColumnarValue::Array(result))
102106
}
103107

108+
/// Spark derives the decimal result type of `pmod` with `Pmod.resultDecimalType`,
109+
/// which follows the `Remainder` rule:
110+
///
111+
/// ```text
112+
/// scale = max(s1, s2)
113+
/// precision = min(p1 - s1, p2 - s2) + scale
114+
/// ```
115+
///
116+
/// The rule is applied to the input argument types.
117+
fn pmod_decimal_result_type(p1: u8, s1: i8, p2: u8, s2: i8) -> DataType {
118+
let scale = s1.max(s2);
119+
let whole_digits = (i32::from(p1) - i32::from(s1)).min(i32::from(p2) - i32::from(s2));
120+
let precision =
121+
(whole_digits + i32::from(scale)).clamp(1, i32::from(DECIMAL128_MAX_PRECISION));
122+
DataType::Decimal128(precision as u8, scale)
123+
}
124+
104125
/// Spark-compatible `pmod` function
105126
/// In ANSI mode, division by zero throws an error.
106127
/// In legacy mode, division by zero returns NULL (Spark behavior).
107128
pub fn spark_pmod(
108129
args: &[ColumnarValue],
109130
enable_ansi_mode: bool,
131+
result_type: &DataType,
110132
) -> Result<ColumnarValue> {
111133
assert_eq_or_internal_err!(args.len(), 2, "pmod expects exactly two arguments");
112134
let args = ColumnarValue::values_to_arrays(args)?;
113-
let left = &args[0];
114-
let right = &args[1];
135+
136+
// Need to handle nulls separately as they are pass through by the signature
137+
if args.iter().any(|arg| arg.data_type() == &DataType::Null) {
138+
return Ok(ColumnarValue::Scalar(ScalarValue::try_new_null(
139+
result_type,
140+
)?));
141+
}
142+
143+
let (left, right): (ArrayRef, ArrayRef) =
144+
if args[0].data_type() == args[1].data_type() {
145+
(Arc::clone(&args[0]), Arc::clone(&args[1]))
146+
} else {
147+
let Some(computation_type) =
148+
decimal_coercion(args[0].data_type(), args[1].data_type())
149+
else {
150+
return exec_err!(
151+
"pmod does not support ({}, {})",
152+
args[0].data_type(),
153+
args[1].data_type()
154+
);
155+
};
156+
let widen = CastOptions {
157+
safe: false,
158+
..Default::default()
159+
};
160+
(
161+
cast_with_options(&args[0], &computation_type, &widen)?,
162+
cast_with_options(&args[1], &computation_type, &widen)?,
163+
)
164+
};
165+
166+
let left = &left;
167+
let right = &right;
115168
let zero = ScalarValue::new_zero(left.data_type())?.to_array_of_size(left.len())?;
116169
let result = try_rem(left, right, enable_ansi_mode)?;
117170
let neg = lt(&result, &zero)?;
118171
let plus = zip(&neg, right, &zero)?;
119172
let result = add(&plus, &result)?;
120173
let result = try_rem(&result, right, enable_ansi_mode)?;
174+
175+
// The remainder is bounded by the divisor, but the result type only carries
176+
// `min(p1 - s1, p2 - s2)` integer digits, so a remainder approaching a
177+
// divisor wider than the dividend does not always fit. Spark wraps decimal
178+
// arithmetic in `CheckOverflow(nullOnOverflow = !ansiEnabled)`, so an
179+
// overflow here is NULL in legacy mode and an error under ANSI.
180+
let result = if result.data_type() == result_type {
181+
result
182+
} else {
183+
let narrow = CastOptions {
184+
safe: !enable_ansi_mode,
185+
..Default::default()
186+
};
187+
cast_with_options(&result, result_type, &narrow)?
188+
};
121189
Ok(ColumnarValue::Array(result))
122190
}
123191

@@ -182,7 +250,19 @@ impl Default for SparkPmod {
182250
impl SparkPmod {
183251
pub fn new() -> Self {
184252
Self {
185-
signature: Signature::numeric(2, Volatility::Immutable),
253+
signature: Signature::one_of(
254+
vec![
255+
// A decimal pair must reach `return_type` with the
256+
// precision and scale as written, since Spark defines
257+
// `Pmod.resultDecimalType` on the declared arguments.
258+
TypeSignature::Coercible(vec![
259+
Coercion::new_exact(TypeSignatureClass::Decimal),
260+
Coercion::new_exact(TypeSignatureClass::Decimal),
261+
]),
262+
TypeSignature::Numeric(2),
263+
],
264+
Volatility::Immutable,
265+
),
186266
}
187267
}
188268
}
@@ -203,13 +283,24 @@ impl ScalarUDFImpl for SparkPmod {
203283
"pmod expects exactly two arguments"
204284
);
205285

206-
// Return the same type as the first argument for simplicity
207-
// Arrow's rem function handles type promotion internally
208-
Ok(arg_types[0].clone())
286+
match (&arg_types[0], &arg_types[1]) {
287+
(DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => {
288+
Ok(pmod_decimal_result_type(*p1, *s1, *p2, *s2))
289+
}
290+
// Need to handle nulls explicitly, see: https://github.com/apache/datafusion/issues/19458
291+
// We align with the behaviour of `mod`
292+
(DataType::Null, DataType::Null) => Ok(DataType::Float64),
293+
(DataType::Null, other) | (other, DataType::Null) => Ok(other.clone()),
294+
_ => Ok(arg_types[0].clone()),
295+
}
209296
}
210297

211298
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
212-
spark_pmod(&args.args, args.config_options.execution.enable_ansi_mode)
299+
spark_pmod(
300+
&args.args,
301+
args.config_options.execution.enable_ansi_mode,
302+
args.return_type(),
303+
)
213304
}
214305
}
215306

@@ -558,7 +649,8 @@ mod test {
558649
let left_value = ColumnarValue::Array(Arc::new(left));
559650
let right_value = ColumnarValue::Array(Arc::new(right));
560651

561-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
652+
let return_type = left_value.data_type();
653+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
562654

563655
if let ColumnarValue::Array(result_array) = result {
564656
let result_int32 =
@@ -581,7 +673,8 @@ mod test {
581673
let left_value = ColumnarValue::Array(Arc::new(left));
582674
let right_value = ColumnarValue::Array(Arc::new(right));
583675

584-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
676+
let return_type = left_value.data_type();
677+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
585678

586679
if let ColumnarValue::Array(result_array) = result {
587680
let result_int64 =
@@ -625,7 +718,8 @@ mod test {
625718
let left_value = ColumnarValue::Array(Arc::new(left));
626719
let right_value = ColumnarValue::Array(Arc::new(right));
627720

628-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
721+
let return_type = left_value.data_type();
722+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
629723

630724
if let ColumnarValue::Array(result_array) = result {
631725
let result_float64 = result_array
@@ -683,7 +777,8 @@ mod test {
683777
let left_value = ColumnarValue::Array(Arc::new(left));
684778
let right_value = ColumnarValue::Array(Arc::new(right));
685779

686-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
780+
let return_type = left_value.data_type();
781+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
687782

688783
if let ColumnarValue::Array(result_array) = result {
689784
let result_float32 = result_array
@@ -718,7 +813,8 @@ mod test {
718813

719814
let left_value = ColumnarValue::Array(Arc::new(left));
720815

721-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
816+
let return_type = left_value.data_type();
817+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
722818

723819
if let ColumnarValue::Array(result_array) = result {
724820
let result_int32 =
@@ -737,7 +833,8 @@ mod test {
737833
let left = Int32Array::from(vec![Some(10)]);
738834
let left_value = ColumnarValue::Array(Arc::new(left));
739835

740-
let result = spark_pmod(&[left_value], false);
836+
let return_type = left_value.data_type();
837+
let result = spark_pmod(&[left_value], false, &return_type);
741838
assert!(result.is_err());
742839
}
743840

@@ -750,7 +847,8 @@ mod test {
750847
let left_value = ColumnarValue::Array(Arc::new(left));
751848
let right_value = ColumnarValue::Array(Arc::new(right));
752849

753-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
850+
let return_type = left_value.data_type();
851+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
754852

755853
if let ColumnarValue::Array(result_array) = result {
756854
let result_int32 =
@@ -772,7 +870,8 @@ mod test {
772870
let left_value = ColumnarValue::Array(Arc::new(left));
773871
let right_value = ColumnarValue::Array(Arc::new(right));
774872

775-
let result = spark_pmod(&[left_value, right_value], true);
873+
let return_type = left_value.data_type();
874+
let result = spark_pmod(&[left_value, right_value], true, &return_type);
776875
assert!(result.is_err());
777876
}
778877

@@ -788,7 +887,8 @@ mod test {
788887
let left_value = ColumnarValue::Array(Arc::new(left));
789888
let right_value = ColumnarValue::Array(Arc::new(right));
790889

791-
let result = spark_pmod(&[left_value, right_value], true);
890+
let return_type = left_value.data_type();
891+
let result = spark_pmod(&[left_value, right_value], true, &return_type);
792892
assert!(result.is_err());
793893
}
794894

@@ -802,7 +902,8 @@ mod test {
802902
let left_value = ColumnarValue::Array(Arc::new(left));
803903
let right_value = ColumnarValue::Array(Arc::new(right));
804904

805-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
905+
let return_type = left_value.data_type();
906+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
806907

807908
if let ColumnarValue::Array(result_array) = result {
808909
let result_float64 = result_array
@@ -825,7 +926,8 @@ mod test {
825926
let left_value = ColumnarValue::Array(Arc::new(left));
826927
let right_value = ColumnarValue::Array(Arc::new(right));
827928

828-
let result = spark_pmod(&[left_value, right_value], true);
929+
let return_type = left_value.data_type();
930+
let result = spark_pmod(&[left_value, right_value], true, &return_type);
829931
assert!(result.is_err());
830932
}
831933

@@ -840,7 +942,8 @@ mod test {
840942
let left_value = ColumnarValue::Array(Arc::new(left));
841943
let right_value = ColumnarValue::Array(Arc::new(right));
842944

843-
let result = spark_pmod(&[left_value, right_value], true).unwrap();
945+
let return_type = left_value.data_type();
946+
let result = spark_pmod(&[left_value, right_value], true, &return_type).unwrap();
844947

845948
if let ColumnarValue::Array(result_array) = result {
846949
let result_int32 =
@@ -858,7 +961,8 @@ mod test {
858961
let left_value = ColumnarValue::Array(Arc::new(left));
859962
let right_value = ColumnarValue::Array(Arc::new(right));
860963

861-
let result = spark_pmod(&[left_value, right_value], true).unwrap();
964+
let return_type = left_value.data_type();
965+
let result = spark_pmod(&[left_value, right_value], true, &return_type).unwrap();
862966

863967
if let ColumnarValue::Array(result_array) = result {
864968
let result_float64 = result_array
@@ -881,7 +985,8 @@ mod test {
881985
let left_value = ColumnarValue::Array(Arc::new(left));
882986
let right_value = ColumnarValue::Array(Arc::new(right));
883987

884-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
988+
let return_type = left_value.data_type();
989+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
885990

886991
if let ColumnarValue::Array(result_array) = result {
887992
let result_int32 =
@@ -919,7 +1024,8 @@ mod test {
9191024
let left_value = ColumnarValue::Array(Arc::new(left));
9201025
let right_value = ColumnarValue::Array(Arc::new(right));
9211026

922-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
1027+
let return_type = left_value.data_type();
1028+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
9231029

9241030
if let ColumnarValue::Array(result_array) = result {
9251031
let result_int32 =
@@ -935,4 +1041,35 @@ mod test {
9351041
panic!("Expected array result");
9361042
}
9371043
}
1044+
1045+
/// Spark's `Pmod.resultDecimalType`: scale = max(s1, s2),
1046+
/// precision = min(p1 - s1, p2 - s2) + scale.
1047+
#[test]
1048+
fn test_pmod_decimal_result_type() {
1049+
// Equal scales: the narrower argument decides the precision.
1050+
assert_eq!(
1051+
pmod_decimal_result_type(3, 1, 2, 1),
1052+
DataType::Decimal128(2, 1)
1053+
);
1054+
// The divisor is wider, so the dividend bounds the result.
1055+
assert_eq!(
1056+
pmod_decimal_result_type(5, 2, 4, 1),
1057+
DataType::Decimal128(5, 2)
1058+
);
1059+
// Differing scales: the wider scale wins.
1060+
assert_eq!(
1061+
pmod_decimal_result_type(6, 3, 4, 1),
1062+
DataType::Decimal128(6, 3)
1063+
);
1064+
// Integral decimals keep a zero scale.
1065+
assert_eq!(
1066+
pmod_decimal_result_type(10, 0, 5, 0),
1067+
DataType::Decimal128(5, 0)
1068+
);
1069+
// The result never exceeds the maximum precision arrow can represent.
1070+
assert_eq!(
1071+
pmod_decimal_result_type(38, 0, 38, 38),
1072+
DataType::Decimal128(38, 38)
1073+
);
1074+
}
9381075
}

0 commit comments

Comments
 (0)