Decimal256 coercion (#7034)

* Add Decimal256 ARM to TypeCoercion is_signed_numeric & is_decimal functions * Add Decimal256 to aggregates.rs as well
apache · Jul 20, 2023 · 8668da6 · 8668da6
1 parent 7e2cca8
commit 8668da6
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/datafusion/expr/src/type_coercion/aggregates.rs b/datafusion/expr/src/type_coercion/aggregates.rs
@@ -467,7 +467,7 @@ pub fn is_sum_support_arg_type(arg_type: &DataType) -> bool {
         _ => matches!(
             arg_type,
             arg_type if NUMERICS.contains(arg_type)
-            || matches!(arg_type, DataType::Decimal128(_, _))
+            || matches!(arg_type, DataType::Decimal128(_, _) | DataType::Decimal256(_, _))
         ),
     }
 }
@@ -480,7 +480,7 @@ pub fn is_avg_support_arg_type(arg_type: &DataType) -> bool {
         _ => matches!(
             arg_type,
             arg_type if NUMERICS.contains(arg_type)
-                || matches!(arg_type, DataType::Decimal128(_, _))
+                || matches!(arg_type, DataType::Decimal128(_, _)| DataType::Decimal256(_, _))
         ),
     }
 }
@@ -579,6 +579,7 @@ mod tests {
         let input_types = vec![
             vec![DataType::Int32],
             vec![DataType::Decimal128(10, 2)],
+            vec![DataType::Decimal256(1, 1)],
             vec![DataType::Utf8],
         ];
         for fun in funs {
@@ -594,6 +595,7 @@ mod tests {
             vec![DataType::Int32],
             vec![DataType::Float32],
             vec![DataType::Decimal128(20, 3)],
+            vec![DataType::Decimal256(20, 3)],
         ];
         for fun in funs {
             for input_type in &input_types {

diff --git a/datafusion/expr/src/type_coercion/mod.rs b/datafusion/expr/src/type_coercion/mod.rs
@@ -49,6 +49,7 @@ pub fn is_signed_numeric(dt: &DataType) -> bool {
             | DataType::Float32
             | DataType::Float64
             | DataType::Decimal128(_, _)
+            | DataType::Decimal256(_, _),
     )
 }
 
@@ -91,5 +92,5 @@ pub fn is_utf8_or_large_utf8(dt: &DataType) -> bool {
 
 /// Determine whether the given data type `dt` is a `Decimal`.
 pub fn is_decimal(dt: &DataType) -> bool {
-    matches!(dt, DataType::Decimal128(_, _))
+    matches!(dt, DataType::Decimal128(_, _) | DataType::Decimal256(_, _))
 }