[arrow-rs] branch master updated: Make arithmetic kernels supports dictionary of decimal array (#3255)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 9abdb5572 Make arithmetic kernels supports dictionary of decimal array 
(#3255)
9abdb5572 is described below

commit 9abdb5572bcaf4d594514d58410a812f3a490c55
Author: Liang-Chi Hsieh <vii...@gmail.com>
AuthorDate: Fri Dec 2 10:49:46 2022 -0800

    Make arithmetic kernels supports dictionary of decimal array (#3255)
---
 arrow/src/compute/kernels/arithmetic.rs | 80 ++++++++++++++++++++++++++++++++-
 1 file changed, 78 insertions(+), 2 deletions(-)

diff --git a/arrow/src/compute/kernels/arithmetic.rs 
b/arrow/src/compute/kernels/arithmetic.rs
index c57e27095..cafd63620 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -35,8 +35,8 @@ use crate::datatypes::{
 };
 #[cfg(feature = "dyn_arith_dict")]
 use crate::datatypes::{
-    Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, 
UInt16Type,
-    UInt32Type, UInt64Type, UInt8Type,
+    Decimal128Type, Decimal256Type, Float32Type, Float64Type, Int16Type, 
Int32Type,
+    Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
 };
 use crate::error::{ArrowError, Result};
 use crate::{datatypes, downcast_primitive_array};
@@ -461,6 +461,14 @@ macro_rules! typed_dict_op {
                 let array = $MATH_OP::<$KT, Float64Type, _>($LEFT, $RIGHT, 
$OP)?;
                 Ok(Arc::new(array))
             }
+            (DataType::Decimal128(_, s1), DataType::Decimal128(_, s2)) if s1 
== s2 => {
+                let array = $MATH_OP::<$KT, Decimal128Type, _>($LEFT, $RIGHT, 
$OP)?;
+                Ok(Arc::new(array))
+            }
+            (DataType::Decimal256(_, s1), DataType::Decimal256(_, s2)) if s1 
== s2 => {
+                let array = $MATH_OP::<$KT, Decimal256Type, _>($LEFT, $RIGHT, 
$OP)?;
+                Ok(Arc::new(array))
+            }
             (t1, t2) => Err(ArrowError::CastError(format!(
                 "Cannot perform arithmetic operation on two dictionary arrays 
of different value types ({} and {})",
                 t1, t2
@@ -3150,4 +3158,72 @@ mod tests {
         let overflow = try_unary_mut(a, |value| value.add_checked(1));
         let _ = overflow.unwrap().expect_err("overflow should be detected");
     }
+
+    #[test]
+    #[cfg(feature = "dyn_arith_dict")]
+    fn test_dict_decimal() {
+        let values = Decimal128Array::from_iter_values([0, 1, 2, 3, 4, 5]);
+        let keys = Int8Array::from_iter_values([1_i8, 2, 5, 4, 3, 0]);
+        let array1 = DictionaryArray::try_new(&keys, &values).unwrap();
+
+        let values = Decimal128Array::from_iter_values([7, -3, 4, 3, 5]);
+        let keys = Int8Array::from_iter_values([0_i8, 0, 1, 2, 3, 4]);
+        let array2 = DictionaryArray::try_new(&keys, &values).unwrap();
+
+        let result = add_dyn(&array1, &array2).unwrap();
+        let expected =
+            Arc::new(Decimal128Array::from(vec![8, 9, 2, 8, 6, 5])) as 
ArrayRef;
+        assert_eq!(&result, &expected);
+
+        let result = subtract_dyn(&array1, &array2).unwrap();
+        let expected =
+            Arc::new(Decimal128Array::from(vec![-6, -5, 8, 0, 0, -5])) as 
ArrayRef;
+        assert_eq!(&result, &expected);
+
+        let values = Decimal256Array::from_iter_values([
+            i256::from_i128(0),
+            i256::from_i128(1),
+            i256::from_i128(2),
+            i256::from_i128(3),
+            i256::from_i128(4),
+            i256::from_i128(5),
+        ]);
+        let keys =
+            Int8Array::from(vec![Some(1_i8), None, Some(5), Some(4), Some(3), 
None]);
+        let array1 = DictionaryArray::try_new(&keys, &values).unwrap();
+
+        let values = Decimal256Array::from_iter_values([
+            i256::from_i128(7),
+            i256::from_i128(-3),
+            i256::from_i128(4),
+            i256::from_i128(3),
+            i256::from_i128(5),
+        ]);
+        let keys =
+            Int8Array::from(vec![Some(0_i8), Some(0), None, Some(2), Some(3), 
Some(4)]);
+        let array2 = DictionaryArray::try_new(&keys, &values).unwrap();
+
+        let result = add_dyn(&array1, &array2).unwrap();
+        let expected = Arc::new(Decimal256Array::from(vec![
+            Some(i256::from_i128(8)),
+            None,
+            None,
+            Some(i256::from_i128(8)),
+            Some(i256::from_i128(6)),
+            None,
+        ])) as ArrayRef;
+
+        assert_eq!(&result, &expected);
+
+        let result = subtract_dyn(&array1, &array2).unwrap();
+        let expected = Arc::new(Decimal256Array::from(vec![
+            Some(i256::from_i128(-6)),
+            None,
+            None,
+            Some(i256::from_i128(0)),
+            Some(i256::from_i128(0)),
+            None,
+        ])) as ArrayRef;
+        assert_eq!(&result, &expected);
+    }
 }