[feat] Support unary_dyn_mut in arrow-arth.

arrow-arith/src/arithmetic.rs

@@ -939,6 +939,20 @@ where
     unary_dyn::<_, T>(array, |value| value.add_wrapping(scalar))
 }
 
+pub fn add_scalar_dyn_mut<T>(
+    array: ArrayRef,
+    scalar: T::Native,
+) -> Result<ArrayRef, ArrowError>
+where
+    T: ArrowNumericType,
+    T::Native: ArrowNativeTypeOp,
+{
+    match unary_dyn_mut::<_, T>(array, |value| value.add_wrapping(scalar)) {
+        Ok(array) => Ok(array),
+        Err(array) => add_scalar_dyn::<T>(&array, 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.

arrow-arith/src/arity.rs

@@ -152,6 +152,37 @@ where
     }
 }
 
+/// Applies an infallible unary function to an [`ArrayRef`] with primitive values.
+/// If the buffer of ArrayRef is not shared with other arrays, then func will
+/// mutate the buffer directly without allocating new buffer.
+pub fn unary_dyn_mut<F, T>(array: ArrayRef, op: F) -> Result<ArrayRef, ArrayRef>
+where
+    T: ArrowPrimitiveType,
+    F: Fn(T::Native) -> T::Native,
+{
+    let array_ref = array.as_ref();
+    downcast_dictionary_array! {
+        array_ref => unary_dict::<_, F, T>(array_ref, op).map_err(|_| array),
+        t => {
+            //Todo support unary_dict_mut
+            if PrimitiveArray::<T>::is_compatible(t) {
+                let primitive_array = array.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap().clone();
+                // Need drop the strong ref which clone before in this function.
+                std::mem::drop(array);
+                match unary_mut::<T, F>(
+                    primitive_array,
+                    op,
+                ) {
+                    Ok(arr) =>  Ok(Arc::new(arr)),
+                    Err(arr) => Err(Arc::new(arr)),
+                }
+            } else {
+                Err(array)
+            }
+        }
+    }
+}
+
 /// Applies a fallible unary function to an array with primitive values.
 pub fn try_unary_dyn<F, T>(array: &dyn Array, op: F) -> Result<ArrayRef, ArrowError>
 where
@@ -552,6 +583,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::arithmetic::{add_scalar_dyn, add_scalar_dyn_mut};
     use arrow_array::builder::*;
     use arrow_array::cast::*;
     use arrow_array::types::*;
@@ -576,6 +608,43 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_unary_f64_mut() {
+        // 1. only have one strong ref, use copy on write.
+        let input =
+            Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
+        let result =
+            unary_dyn_mut::<_, Float64Type>(make_array(input.into_data()), |n| n + 1.0)
+                .unwrap();
+        assert_eq!(
+            result.as_any().downcast_ref::<Float64Array>().unwrap(),
+            &Float64Array::from(vec![Some(6.1f64), None, Some(7.8), None, Some(8.2)])
+        );
+
+        // 2. More than one strong ref
+        let input =
+            Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
+        let slice = input.slice(1, 4);
+        let result = unary_dyn_mut::<_, Float64Type>(slice, |n| n + 1.0);
+        assert!(result.is_err())
+    }
+    #[test]
+    fn profile() {
+        let mut vec = vec![];
+        let batch_size = 16384;
+        for i in 0..1024 {
+            let array = PrimitiveArray::<Float32Type>::from_iter_values(
+                (0..batch_size).map(|x| x as f32),
+            );
+            vec.push(make_array(array.into_data()));
+        }
+
+        for i in 0..vec.len() {
+            add_scalar_dyn::<Float32Type>(vec.pop().unwrap().as_ref(), 1 as f32)
+                .expect("panic message");
+        }
+    }
+
     #[test]
     fn test_unary_dict_and_unary_dyn() {
         let mut builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();

arrow/Cargo.toml

@@ -296,3 +296,9 @@ required-features = ["pyarrow"]
 [[test]]
 name = "array_cast"
 required-features = ["chrono-tz"]
+
+[[bench]]
+name = "cow"
+harness = false
+required-features = ["test_utils"]
+

arrow/benches/cow.rs

@@ -0,0 +1,68 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#[macro_use]
+extern crate criterion;
+use arrow::util::bench_util::create_primitive_array;
+use arrow_arith::arithmetic::{add_scalar_dyn, add_scalar_dyn_mut};
+use arrow_array::types::Float32Type;
+use arrow_array::{make_array, Array, ArrayRef, PrimitiveArray};
+use arrow_schema::ArrowError;
+use criterion::Criterion;
+
+extern crate arrow;
+
+const BATCH_SIZE: i32 = 163840;
+const BATCH_NUM: i32 = 1024 * 10;
+fn add_cow() {
+    let mut vec = vec![];
+    for i in 0..BATCH_NUM {
+        let array = PrimitiveArray::<Float32Type>::from_iter_values(
+            (0..BATCH_SIZE).map(|x| x as f32),
+        );
+        vec.push(make_array(array.into_data()));
+    }
+
+    for i in 0..vec.len() {
+        add_scalar_dyn_mut::<Float32Type>(vec.pop().unwrap(), 1 as f32)
+            .expect("panic message");
+    }
+}
+
+fn add_copy() {
+    let mut vec = vec![];
+    for i in 0..BATCH_NUM {
+        let array = PrimitiveArray::<Float32Type>::from_iter_values(
+            (0..BATCH_SIZE).map(|x| x as f32),
+        );
+        vec.push(make_array(array.into_data()));
+    }
+
+    for i in 0..vec.len() {
+        add_scalar_dyn::<Float32Type>(vec.pop().unwrap().as_ref(), 1 as f32)
+            .expect("panic message");
+    }
+}
+
+fn cow_benchmark(c: &mut Criterion) {
+    c.bench_function(&format!("add"), |b| {
+        b.iter(|| criterion::black_box(add_cow()))
+    });
+}
+
+criterion_group!(benches, cow_benchmark);
+criterion_main!(benches);