[arrow-rs] branch master updated: Support DictionaryArray in add_scalar kernel (#2018)

[vi...@apache.org]

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

viirya 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 f8fa9849c Support DictionaryArray in add_scalar kernel (#2018)
f8fa9849c is described below

commit f8fa9849c6da689e5dff128be3ac8f11758a3e8e
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Wed Jul 6 18:22:41 2022 -0700

    Support DictionaryArray in add_scalar kernel (#2018)
    
    * Add add_scalar_dyn
    
    * Trigger Build
    
    * Trigger Build
---
 arrow/src/compute/kernels/arithmetic.rs | 52 +++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs
index 0189dade2..f64038c19 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -30,6 +30,7 @@ use crate::buffer::Buffer;
 #[cfg(feature = "simd")]
 use crate::buffer::MutableBuffer;
 use crate::compute::kernels::arity::unary;
+use crate::compute::unary_dyn;
 use crate::compute::util::combine_option_bitmap;
 use crate::datatypes;
 use crate::datatypes::{ArrowNumericType, DataType};
@@ -707,6 +708,17 @@ where
     Ok(unary(array, |value| value + scalar))
 }
 
+/// Add every value in an array by a scalar. If any value in the array is null then the
+/// result is also null. The given array must be a `PrimitiveArray` of the type same as
+/// the scalar, or a `DictionaryArray` of the value type same as the scalar.
+pub fn add_scalar_dyn<T>(array: &dyn Array, scalar: T::Native) -> Result<ArrayRef>
+where
+    T: datatypes::ArrowNumericType,
+    T::Native: Add<Output = T::Native>,
+{
+    unary_dyn::<_, T>(array, |value| value + scalar)
+}
+
 /// Perform `left - right` operation on two arrays. If either left or right value is null
 /// then the result is also null.
 pub fn subtract<T>(
@@ -958,6 +970,46 @@ mod tests {
         assert_eq!(19, c.value(4));
     }
 
+    #[test]
+    fn test_primitive_array_add_scalar_dyn() {
+        let a = Int32Array::from(vec![Some(5), Some(6), Some(7), None, Some(9)]);
+        let b = 1_i32;
+        let c = add_scalar_dyn::<Int32Type>(&a, b).unwrap();
+        let c = c.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(6, c.value(0));
+        assert_eq!(7, c.value(1));
+        assert_eq!(8, c.value(2));
+        assert!(c.is_null(3));
+        assert_eq!(10, c.value(4));
+
+        let key_builder = PrimitiveBuilder::<Int8Type>::new(3);
+        let value_builder = PrimitiveBuilder::<Int32Type>::new(2);
+        let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder);
+        builder.append(5).unwrap();
+        builder.append_null().unwrap();
+        builder.append(7).unwrap();
+        builder.append(8).unwrap();
+        builder.append(9).unwrap();
+        let a = builder.finish();
+        let b = -1_i32;
+
+        let c = add_scalar_dyn::<Int32Type>(&a, b).unwrap();
+        let c = c
+            .as_any()
+            .downcast_ref::<DictionaryArray<Int8Type>>()
+            .unwrap();
+        let values = c
+            .values()
+            .as_any()
+            .downcast_ref::<PrimitiveArray<Int32Type>>()
+            .unwrap();
+        assert_eq!(4, values.value(c.key(0).unwrap()));
+        assert!(c.is_null(1));
+        assert_eq!(6, values.value(c.key(2).unwrap()));
+        assert_eq!(7, values.value(c.key(3).unwrap()));
+        assert_eq!(8, values.value(c.key(4).unwrap()));
+    }
+
     #[test]
     fn test_primitive_array_subtract_dyn() {
         let a = Int32Array::from(vec![Some(51), Some(6), Some(15), Some(8), Some(9)]);