diff --git a/datafusion/expr-common/src/type_coercion/binary.rs b/datafusion/expr-common/src/type_coercion/binary.rs index 381897ae86fdc..e7c20dde101b0 100644 --- a/datafusion/expr-common/src/type_coercion/binary.rs +++ b/datafusion/expr-common/src/type_coercion/binary.rs @@ -1124,8 +1124,7 @@ fn get_wider_decimal_type_cross_variant( // max(s1, s2) + max(p1-s1, p2-s2), max(s1, s2) let s = s1.max(s2); - let range = (p1 as i8 - s1).max(p2 as i8 - s2); - let required_precision = (range + s) as u8; + let required_precision = required_decimal_precision(p1, s1, p2, s2); // Choose the larger variant between the two input types, while making sure we don't overflow the precision. match (lhs_type, rhs_type) { @@ -1193,33 +1192,52 @@ fn get_wider_decimal_type( ) -> Option { match (lhs_decimal_type, rhs_type) { (DataType::Decimal32(p1, s1), DataType::Decimal32(p2, s2)) => { - // max(s1, s2) + max(p1-s1, p2-s2), max(s1, s2) let s = *s1.max(s2); - let range = (*p1 as i8 - s1).max(*p2 as i8 - s2); - Some(create_decimal32_type((range + s) as u8, s)) + Some(create_decimal32_type( + required_decimal_precision(*p1, *s1, *p2, *s2), + s, + )) } (DataType::Decimal64(p1, s1), DataType::Decimal64(p2, s2)) => { - // max(s1, s2) + max(p1-s1, p2-s2), max(s1, s2) let s = *s1.max(s2); - let range = (*p1 as i8 - s1).max(*p2 as i8 - s2); - Some(create_decimal64_type((range + s) as u8, s)) + Some(create_decimal64_type( + required_decimal_precision(*p1, *s1, *p2, *s2), + s, + )) } (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => { - // max(s1, s2) + max(p1-s1, p2-s2), max(s1, s2) let s = *s1.max(s2); - let range = (*p1 as i8 - s1).max(*p2 as i8 - s2); - Some(create_decimal128_type((range + s) as u8, s)) + Some(create_decimal128_type( + required_decimal_precision(*p1, *s1, *p2, *s2), + s, + )) } (DataType::Decimal256(p1, s1), DataType::Decimal256(p2, s2)) => { - // max(s1, s2) + max(p1-s1, p2-s2), max(s1, s2) let s = *s1.max(s2); - let range = (*p1 as i8 - s1).max(*p2 as i8 - s2); - Some(create_decimal256_type((range + s) as u8, s)) + Some(create_decimal256_type( + required_decimal_precision(*p1, *s1, *p2, *s2), + s, + )) } (_, _) => None, } } +/// Computes `max(s1, s2) + max(p1 - s1, p2 - s2)`: the precision needed to hold +/// any value of either decimal type. +/// +/// The intermediate values do not fit in `i8` (the type of a decimal scale): +/// `Decimal256` allows a precision and a scale of up to 76, so `p1 - s1` can +/// reach 152 and the sum can reach 228. Computing this in `i8` panics with +/// "attempt to add with overflow" in debug builds, so widen to `i32` and +/// saturate into `u8` instead. Callers then either clamp the result to the +/// variant's maximum precision (`create_decimal*_type`) or reject it. +fn required_decimal_precision(p1: u8, s1: i8, p2: u8, s2: i8) -> u8 { + let s = s1.max(s2) as i32; + let range = (p1 as i32 - s1 as i32).max(p2 as i32 - s2 as i32); + (range + s).clamp(0, u8::MAX as i32) as u8 +} + /// Convert the numeric data type to the decimal data type. /// We support signed and unsigned integer types and floating-point type. fn coerce_numeric_type_to_decimal32(numeric_type: &DataType) -> Option { diff --git a/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs b/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs index cfa3bbe189929..cfb0166ed2dc7 100644 --- a/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs +++ b/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs @@ -1093,3 +1093,34 @@ fn test_string_concat_coercion() -> Result<()> { Ok(()) } + +/// `Decimal256` allows a precision and a scale of up to 76, so the required +/// precision `max(s1, s2) + max(p1 - s1, p2 - s2)` can reach 228 and the +/// intermediate `p - s` can reach 152. Neither fits in the `i8` used for +/// decimal scales, which used to panic with "attempt to add with overflow" +/// (or "attempt to subtract with overflow") in debug builds. +#[test] +fn test_decimal256_comparison_coercion_precision_overflow() -> Result<()> { + // required precision = max(0, 52) + max(76 - 0, 76 - 52) = 128 + assert_eq!( + comparison_coercion(&DataType::Decimal256(76, 0), &DataType::Decimal256(76, 52)), + Some(DataType::Decimal256(76, 52)) + ); + + // required precision = max(0, 76) + max(76 - 0, 76 - 76) = 152 + assert_eq!( + comparison_coercion(&DataType::Decimal256(76, 0), &DataType::Decimal256(76, 76)), + Some(DataType::Decimal256(76, 76)) + ); + + // `p1 - s1` alone is 76 - (-76) = 152 before the sum is even computed + assert_eq!( + comparison_coercion( + &DataType::Decimal256(76, -76), + &DataType::Decimal256(76, 76) + ), + Some(DataType::Decimal256(76, 76)) + ); + + Ok(()) +}