Skip to content

Commit e363e12

Browse files
amitvijapuramitvijapur
committed
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.
1 parent 1f8de79 commit e363e12

1 file changed

Lines changed: 46 additions & 34 deletions

File tree

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

Lines changed: 46 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -184,19 +184,11 @@ fn pmod_numeric_coercion(lhs: &DataType, rhs: &DataType) -> Result<Vec<DataType>
184184
pub fn spark_pmod(
185185
args: &[ColumnarValue],
186186
enable_ansi_mode: bool,
187+
result_type: &DataType,
187188
) -> Result<ColumnarValue> {
188189
assert_eq_or_internal_err!(args.len(), 2, "pmod expects exactly two arguments");
189190
let args = ColumnarValue::values_to_arrays(args)?;
190191

191-
// Decimal arguments reach here with their declared types intact, so the
192-
// Spark result type is derived before they are widened for the computation.
193-
let result_type = match (args[0].data_type(), args[1].data_type()) {
194-
(DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => {
195-
Some(pmod_decimal_result_type(*p1, *s1, *p2, *s2))
196-
}
197-
_ => None,
198-
};
199-
200192
let (left, right): (ArrayRef, ArrayRef) =
201193
if args[0].data_type() == args[1].data_type() {
202194
(Arc::clone(&args[0]), Arc::clone(&args[1]))
@@ -230,15 +222,14 @@ pub fn spark_pmod(
230222
// divisor wider than the dividend does not always fit. Spark wraps decimal
231223
// arithmetic in `CheckOverflow(nullOnOverflow = !ansiEnabled)`, so an
232224
// overflow here is NULL in legacy mode and an error under ANSI.
233-
let result = match result_type {
234-
Some(result_type) if result.data_type() != &result_type => {
235-
let narrow = CastOptions {
236-
safe: !enable_ansi_mode,
237-
..Default::default()
238-
};
239-
cast_with_options(&result, &result_type, &narrow)?
240-
}
241-
_ => result,
225+
let result = if result.data_type() == result_type {
226+
result
227+
} else {
228+
let narrow = CastOptions {
229+
safe: !enable_ansi_mode,
230+
..Default::default()
231+
};
232+
cast_with_options(&result, result_type, &narrow)?
242233
};
243234
Ok(ColumnarValue::Array(result))
244235
}
@@ -353,7 +344,13 @@ impl ScalarUDFImpl for SparkPmod {
353344
}
354345

355346
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
356-
spark_pmod(&args.args, args.config_options.execution.enable_ansi_mode)
347+
// The Spark result type was already derived in `return_type`, so it is
348+
// read back here rather than recomputed from the argument arrays.
349+
spark_pmod(
350+
&args.args,
351+
args.config_options.execution.enable_ansi_mode,
352+
args.return_type(),
353+
)
357354
}
358355
}
359356

@@ -702,7 +699,8 @@ mod test {
702699
let left_value = ColumnarValue::Array(Arc::new(left));
703700
let right_value = ColumnarValue::Array(Arc::new(right));
704701

705-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
702+
let return_type = left_value.data_type();
703+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
706704

707705
if let ColumnarValue::Array(result_array) = result {
708706
let result_int32 =
@@ -725,7 +723,8 @@ mod test {
725723
let left_value = ColumnarValue::Array(Arc::new(left));
726724
let right_value = ColumnarValue::Array(Arc::new(right));
727725

728-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
726+
let return_type = left_value.data_type();
727+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
729728

730729
if let ColumnarValue::Array(result_array) = result {
731730
let result_int64 =
@@ -769,7 +768,8 @@ mod test {
769768
let left_value = ColumnarValue::Array(Arc::new(left));
770769
let right_value = ColumnarValue::Array(Arc::new(right));
771770

772-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
771+
let return_type = left_value.data_type();
772+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
773773

774774
if let ColumnarValue::Array(result_array) = result {
775775
let result_float64 = result_array
@@ -827,7 +827,8 @@ mod test {
827827
let left_value = ColumnarValue::Array(Arc::new(left));
828828
let right_value = ColumnarValue::Array(Arc::new(right));
829829

830-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
830+
let return_type = left_value.data_type();
831+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
831832

832833
if let ColumnarValue::Array(result_array) = result {
833834
let result_float32 = result_array
@@ -862,7 +863,8 @@ mod test {
862863

863864
let left_value = ColumnarValue::Array(Arc::new(left));
864865

865-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
866+
let return_type = left_value.data_type();
867+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
866868

867869
if let ColumnarValue::Array(result_array) = result {
868870
let result_int32 =
@@ -881,7 +883,8 @@ mod test {
881883
let left = Int32Array::from(vec![Some(10)]);
882884
let left_value = ColumnarValue::Array(Arc::new(left));
883885

884-
let result = spark_pmod(&[left_value], false);
886+
let return_type = left_value.data_type();
887+
let result = spark_pmod(&[left_value], false, &return_type);
885888
assert!(result.is_err());
886889
}
887890

@@ -894,7 +897,8 @@ mod test {
894897
let left_value = ColumnarValue::Array(Arc::new(left));
895898
let right_value = ColumnarValue::Array(Arc::new(right));
896899

897-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
900+
let return_type = left_value.data_type();
901+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
898902

899903
if let ColumnarValue::Array(result_array) = result {
900904
let result_int32 =
@@ -916,7 +920,8 @@ mod test {
916920
let left_value = ColumnarValue::Array(Arc::new(left));
917921
let right_value = ColumnarValue::Array(Arc::new(right));
918922

919-
let result = spark_pmod(&[left_value, right_value], true);
923+
let return_type = left_value.data_type();
924+
let result = spark_pmod(&[left_value, right_value], true, &return_type);
920925
assert!(result.is_err());
921926
}
922927

@@ -932,7 +937,8 @@ mod test {
932937
let left_value = ColumnarValue::Array(Arc::new(left));
933938
let right_value = ColumnarValue::Array(Arc::new(right));
934939

935-
let result = spark_pmod(&[left_value, right_value], true);
940+
let return_type = left_value.data_type();
941+
let result = spark_pmod(&[left_value, right_value], true, &return_type);
936942
assert!(result.is_err());
937943
}
938944

@@ -946,7 +952,8 @@ mod test {
946952
let left_value = ColumnarValue::Array(Arc::new(left));
947953
let right_value = ColumnarValue::Array(Arc::new(right));
948954

949-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
955+
let return_type = left_value.data_type();
956+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
950957

951958
if let ColumnarValue::Array(result_array) = result {
952959
let result_float64 = result_array
@@ -969,7 +976,8 @@ mod test {
969976
let left_value = ColumnarValue::Array(Arc::new(left));
970977
let right_value = ColumnarValue::Array(Arc::new(right));
971978

972-
let result = spark_pmod(&[left_value, right_value], true);
979+
let return_type = left_value.data_type();
980+
let result = spark_pmod(&[left_value, right_value], true, &return_type);
973981
assert!(result.is_err());
974982
}
975983

@@ -984,7 +992,8 @@ mod test {
984992
let left_value = ColumnarValue::Array(Arc::new(left));
985993
let right_value = ColumnarValue::Array(Arc::new(right));
986994

987-
let result = spark_pmod(&[left_value, right_value], true).unwrap();
995+
let return_type = left_value.data_type();
996+
let result = spark_pmod(&[left_value, right_value], true, &return_type).unwrap();
988997

989998
if let ColumnarValue::Array(result_array) = result {
990999
let result_int32 =
@@ -1002,7 +1011,8 @@ mod test {
10021011
let left_value = ColumnarValue::Array(Arc::new(left));
10031012
let right_value = ColumnarValue::Array(Arc::new(right));
10041013

1005-
let result = spark_pmod(&[left_value, right_value], true).unwrap();
1014+
let return_type = left_value.data_type();
1015+
let result = spark_pmod(&[left_value, right_value], true, &return_type).unwrap();
10061016

10071017
if let ColumnarValue::Array(result_array) = result {
10081018
let result_float64 = result_array
@@ -1025,7 +1035,8 @@ mod test {
10251035
let left_value = ColumnarValue::Array(Arc::new(left));
10261036
let right_value = ColumnarValue::Array(Arc::new(right));
10271037

1028-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
1038+
let return_type = left_value.data_type();
1039+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
10291040

10301041
if let ColumnarValue::Array(result_array) = result {
10311042
let result_int32 =
@@ -1063,7 +1074,8 @@ mod test {
10631074
let left_value = ColumnarValue::Array(Arc::new(left));
10641075
let right_value = ColumnarValue::Array(Arc::new(right));
10651076

1066-
let result = spark_pmod(&[left_value, right_value], false).unwrap();
1077+
let return_type = left_value.data_type();
1078+
let result = spark_pmod(&[left_value, right_value], false, &return_type).unwrap();
10671079

10681080
if let ColumnarValue::Array(result_array) = result {
10691081
let result_int32 =

0 commit comments

Comments
 (0)