CR comments

Co-authored-by: Jeffrey Vo <jeffrey.vo.australia@gmail.com>

Author: AdamGS

datafusion/expr-common/src/type_coercion/binary.rs

@@ -327,6 +327,16 @@ impl<'a> BinaryTypeCoercer<'a> {
 
 // TODO Move the rest inside of BinaryTypeCoercer
 
+fn is_decimal(data_type: &DataType) -> bool {
+    matches!(
+        data_type,
+        DataType::Decimal32(..)
+            | DataType::Decimal64(..)
+            | DataType::Decimal128(..)
+            | DataType::Decimal256(..)
+    )
+}
+
 /// Coercion rules for mathematics operators between decimal and non-decimal types.
 fn math_decimal_coercion(
     lhs_type: &DataType,
@@ -358,10 +368,11 @@ fn math_decimal_coercion(
             Some((lhs_type.clone(), rhs_type.clone()))
         }
         // Cross-variant decimal coercion - choose larger variant with appropriate precision/scale
-        (Decimal32(_, _), Decimal64(_, _) | Decimal128(_, _) | Decimal256(_, _))
-        | (Decimal64(_, _), Decimal32(_, _) | Decimal128(_, _) | Decimal256(_, _))
-        | (Decimal128(_, _), Decimal32(_, _) | Decimal64(_, _) | Decimal256(_, _))
-        | (Decimal256(_, _), Decimal32(_, _) | Decimal64(_, _) | Decimal128(_, _)) => {
+        (lhs, rhs)
+            if is_decimal(lhs)
+                && is_decimal(rhs)
+                && std::mem::discriminant(lhs) != std::mem::discriminant(rhs) =>
+        {
             let coerced_type = get_wider_decimal_type_cross_variant(lhs_type, rhs_type)?;
             Some((coerced_type.clone(), coerced_type))
         }
@@ -1023,6 +1034,11 @@ fn get_wider_decimal_type_cross_variant(
     let range = (p1 as i8 - s1).max(p2 as i8 - s2);
     let required_precision = (range + s) as u8;
 
+    // We currently don't handle cases where the required percision overflows
+    if required_precision > DECIMAL256_MAX_PRECISION {
+        return None;
+    }
+
     // Choose the larger variant between the two input types
     match (lhs_type, rhs_type) {
         (Decimal32(_, _), Decimal64(_, _)) | (Decimal64(_, _), Decimal32(_, _)) => {

datafusion/physical-plan/src/joins/sort_merge_join/stream.rs

@@ -1999,6 +1999,7 @@ fn is_join_arrays_equal(
             DataType::Decimal32(..) => compare_value!(Decimal32Array),
             DataType::Decimal64(..) => compare_value!(Decimal64Array),
             DataType::Decimal128(..) => compare_value!(Decimal128Array),
+            DataType::Decimal256(..) => compare_value!(Decimal256Array),
             DataType::Timestamp(time_unit, None) => match time_unit {
                 TimeUnit::Second => compare_value!(TimestampSecondArray),
                 TimeUnit::Millisecond => compare_value!(TimestampMillisecondArray),

datafusion/spark/src/function/math/width_bucket.rs

@@ -32,6 +32,7 @@ use datafusion_common::cast::{
 };
 use datafusion_common::{exec_err, Result};
 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_expr::type_coercion::is_signed_numeric;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
 use datafusion_functions::utils::make_scalar_function;
 
@@ -93,23 +94,11 @@ impl ScalarUDFImpl for SparkWidthBucket {
 
         let (v, lo, hi, n) = (&types[0], &types[1], &types[2], &types[3]);
 
-        let is_num = |t: &DataType| {
-            matches!(
-                t,
-                Int8 | Int16
-                    | Int32
-                    | Int64
-                    | Float32
-                    | Float64
-                    | Decimal32(_, _)
-                    | Decimal64(_, _)
-                    | Decimal128(_, _)
-            )
-        };
-
         match (v, lo, hi, n) {
             (a, b, c, &(Int8 | Int16 | Int32 | Int64))
-                if is_num(a) && is_num(b) && is_num(c) =>
+                if is_signed_numeric(a)
+                    && is_signed_numeric(b)
+                    && is_signed_numeric(c) =>
             {
                 Ok(vec![Float64, Float64, Float64, Int32])
             }