Add binary_mut and try_binary_mut

arrow/src/compute/kernels/arithmetic.rs

@@ -27,7 +27,8 @@ use crate::array::*;
 use crate::buffer::MutableBuffer;
 use crate::compute::kernels::arity::unary;
 use crate::compute::{
-    binary, binary_opt, try_binary, try_unary, try_unary_dyn, unary_dyn,
+    binary, binary_mut, binary_opt, try_binary, try_binary_mut, try_unary, try_unary_dyn,
+    unary_dyn,
 };
 use crate::datatypes::{
     ArrowNativeTypeOp, ArrowNumericType, DataType, Date32Type, Date64Type,
@@ -733,6 +734,28 @@ where
     math_op(left, right, |a, b| a.add_wrapping(b))
 }
 
+/// Perform `left + right` operation on two arrays by mutating `left` with operation results.
+/// If either left or right value is null then the result is also null.
+///
+/// This only mutates the array if it is not shared buffers with other arrays. For shared
+/// array, it returns an `Err` which wraps input array.
+///
+/// This doesn't detect overflow. Once overflowing, the result will wrap around.
+/// For an overflow-checking variant, use `add_checked_mut` instead.
+pub fn add_mut<T>(
+    left: PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+) -> std::result::Result<
+    PrimitiveArray<T>,
+    std::result::Result<PrimitiveArray<T>, ArrowError>,
+>
+where
+    T: ArrowNumericType,
+    T::Native: ArrowNativeTypeOp,
+{
+    binary_mut(left, right, |a, b| a.add_wrapping(b))
+}
+
 /// Perform `left + right` operation on two arrays. If either left or right value is null
 /// then the result is also null.
 ///
@@ -749,6 +772,28 @@ where
     try_binary(left, right, |a, b| a.add_checked(b))
 }
 
+/// Perform `left + right` operation on two arrays by mutating `left` with operation results.
+/// If either left or right value is null then the result is also null.
+///
+/// This only mutates the array if it is not shared buffers with other arrays. For shared
+/// array, it returns an `Err` which wraps input array.
+///
+/// This detects overflow and returns an `Err` for that. For an non-overflow-checking variant,
+/// use `add_mut` instead.
+pub fn add_checked_mut<T>(
+    left: PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+) -> std::result::Result<
+    PrimitiveArray<T>,
+    std::result::Result<PrimitiveArray<T>, ArrowError>,
+>
+where
+    T: ArrowNumericType,
+    T::Native: ArrowNativeTypeOp,
+{
+    try_binary_mut(left, right, |a, b| a.add_checked(b))
+}
+
 /// Perform `left + right` operation on two arrays. If either left or right value is null
 /// then the result is also null.
 ///
@@ -3098,4 +3143,29 @@ mod tests {
         assert_eq!(result.len(), 13);
         assert_eq!(result.null_count(), 13);
     }
+
+    #[test]
+    fn test_primitive_array_add_mut() {
+        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
+        let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
+
+        let c = add_mut(a, &b).unwrap();
+        let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
+        assert_eq!(c, expected);
+    }
+
+    #[test]
+    fn test_primitive_add_mut_wrapping_overflow() {
+        let a = Int32Array::from(vec![i32::MAX, i32::MIN]);
+        let b = Int32Array::from(vec![1, 1]);
+
+        let wrapped = add_mut(a, &b).unwrap();
+        let expected = Int32Array::from(vec![-2147483648, -2147483647]);
+        assert_eq!(expected, wrapped);
+
+        let a = Int32Array::from(vec![i32::MAX, i32::MIN]);
+        let b = Int32Array::from(vec![1, 1]);
+        let overflow = add_checked_mut(a, &b);
+        let _ = overflow.expect_err("overflow should be detected");
+    }
 }

arrow/src/compute/kernels/arity.rs

@@ -205,6 +205,72 @@ where
     Ok(unsafe { build_primitive_array(len, buffer, null_count, null_buffer) })
 }
 
+/// Given two arrays of length `len`, calls `op(a[i], b[i])` for `i` in `0..len`, mutating
+/// the mutable [`PrimitiveArray`] `a`. If any index is null in either `a` or `b`, the
+/// corresponding index in the result will also be null.
+///
+/// Mutable primitive array means that the buffer is not shared with other arrays.
+/// As a result, this mutates the buffer directly without allocating new buffer.
+///
+/// Like [`unary`] the provided function is evaluated for every index, ignoring validity. This
+/// is beneficial when the cost of the operation is low compared to the cost of branching, and
+/// especially when the operation can be vectorised, however, requires `op` to be infallible
+/// for all possible values of its inputs
+///
+/// # Error
+///
+/// This function gives error if the arrays have different lengths.
+/// This function gives error of original [`PrimitiveArray`] `a` if it is not a mutable
+/// primitive array.
+pub fn binary_mut<T, F>(
+    a: PrimitiveArray<T>,
+    b: &PrimitiveArray<T>,
+    op: F,
+) -> std::result::Result<
+    PrimitiveArray<T>,
+    std::result::Result<PrimitiveArray<T>, ArrowError>,
+>
+where
+    T: ArrowPrimitiveType,
+    F: Fn(T::Native, T::Native) -> T::Native,
+{
+    if a.len() != b.len() {
+        return Err(Err(ArrowError::ComputeError(
+            "Cannot perform binary operation on arrays of different length".to_string(),
+        )));
+    }
+    let len = a.len();
+
+    if a.is_empty() {
+        return Ok(PrimitiveArray::from(ArrayData::new_empty(&T::DATA_TYPE)));
+    }
+
+    let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len).unwrap();
+    let null_count = null_buffer
+        .as_ref()
+        .map(|x| len - x.count_set_bits_offset(0, len))
+        .unwrap_or_default();
+
+    let mut builder = a.into_builder().map_err(Ok)?;
+
+    builder
+        .values_slice_mut()
+        .iter_mut()
+        .zip(b.values())
+        .for_each(|(l, r)| *l = op(*l, *r));
+
+    let array_builder = builder
+        .finish()
+        .data()
+        .clone()
+        .into_builder()
+        .null_bit_buffer(null_buffer)
+        .null_count(null_count);
+
+    let array_data = unsafe { array_builder.build_unchecked() };
+    Ok(PrimitiveArray::<T>::from(array_data))
+}
+
 /// Applies the provided fallible binary operation across `a` and `b`, returning any error,
 /// and collecting the results into a [`PrimitiveArray`]. If any index is null in either `a`
 /// or `b`, the corresponding index in the result will also be null
@@ -262,6 +328,79 @@ where
     }
 }
 
+/// Applies the provided fallible binary operation across `a` and `b` by mutating the mutable
+/// [`PrimitiveArray`] `a` with the results, returning any error. If any index is null in
+/// either `a` or `b`, the corresponding index in the result will also be null
+///
+/// Like [`try_unary`] the function is only evaluated for non-null indices
+///
+/// Mutable primitive array means that the buffer is not shared with other arrays.
+/// As a result, this mutates the buffer directly without allocating new buffer.
+///
+/// # Error
+///
+/// Return an error if the arrays have different lengths or
+/// the operation is under erroneous.
+/// This function gives error of original [`PrimitiveArray`] `a` if it is not a mutable
+/// primitive array.
+pub fn try_binary_mut<T, F>(
+    a: PrimitiveArray<T>,
+    b: &PrimitiveArray<T>,
+    op: F,
+) -> std::result::Result<
+    PrimitiveArray<T>,
+    std::result::Result<PrimitiveArray<T>, ArrowError>,
+>
+where
+    T: ArrowPrimitiveType,
+    F: Fn(T::Native, T::Native) -> Result<T::Native>,
+{
+    if a.len() != b.len() {
+        return Err(Err(ArrowError::ComputeError(
+            "Cannot perform binary operation on arrays of different length".to_string(),
+        )));
+    }
+    let len = a.len();
+
+    if a.is_empty() {
+        return Ok(PrimitiveArray::from(ArrayData::new_empty(&T::DATA_TYPE)));
+    }
+
+    if a.null_count() == 0 && b.null_count() == 0 {
+        try_binary_no_nulls_mut(len, a, b, op)
+    } else {
+        let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len).unwrap();
+        let null_count = null_buffer
+            .as_ref()
+            .map(|x| len - x.count_set_bits_offset(0, len))
+            .unwrap_or_default();
+
+        let mut builder = a.into_builder().map_err(Ok)?;
+
+        let slice = builder.values_slice_mut();
+
+        try_for_each_valid_idx(len, 0, null_count, null_buffer.as_deref(), |idx| {
+            unsafe {
+                *slice.get_unchecked_mut(idx) =
+                    op(*slice.get_unchecked(idx), b.value_unchecked(idx))?
+            };
+            Ok::<_, ArrowError>(())
+        })
+        .map_err(Err)?;
+
+        let array_builder = builder
+            .finish()
+            .data()
+            .clone()
+            .into_builder()
+            .null_bit_buffer(null_buffer)
+            .null_count(null_count);
+
+        let array_data = unsafe { array_builder.build_unchecked() };
+        Ok(PrimitiveArray::<T>::from(array_data))
+    }
+}
+
 /// This intentional inline(never) attribute helps LLVM optimize the loop.
 #[inline(never)]
 fn try_binary_no_nulls<A: ArrayAccessor, B: ArrayAccessor, F, O>(
@@ -283,6 +422,33 @@ where
     Ok(unsafe { build_primitive_array(len, buffer.into(), 0, None) })
 }
 
+/// This intentional inline(never) attribute helps LLVM optimize the loop.
+#[inline(never)]
+fn try_binary_no_nulls_mut<T, F>(
+    len: usize,
+    a: PrimitiveArray<T>,
+    b: &PrimitiveArray<T>,
+    op: F,
+) -> std::result::Result<
+    PrimitiveArray<T>,
+    std::result::Result<PrimitiveArray<T>, ArrowError>,
+>
+where
+    T: ArrowPrimitiveType,
+    F: Fn(T::Native, T::Native) -> Result<T::Native>,
+{
+    let mut builder = a.into_builder().map_err(Ok)?;
+    let slice = builder.values_slice_mut();
+
+    for idx in 0..len {
+        unsafe {
+            *slice.get_unchecked_mut(idx) =
+                op(*slice.get_unchecked(idx), b.value_unchecked(idx)).map_err(Err)?;
+        };
+    }
+    Ok(builder.finish())
+}
+
 #[inline(never)]
 fn try_binary_opt_no_nulls<A: ArrayAccessor, B: ArrayAccessor, F, O>(
     len: usize,
@@ -358,6 +524,7 @@ mod tests {
     use super::*;
     use crate::array::{as_primitive_array, Float64Array, PrimitiveDictionaryBuilder};
     use crate::datatypes::{Float64Type, Int32Type, Int8Type};
+    use arrow_array::Int32Array;
 
     #[test]
     fn test_unary_f64_slice() {
@@ -417,4 +584,37 @@ mod tests {
             &expected
         );
     }
+
+    #[test]
+    fn test_binary_mut() {
+        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
+        let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
+        let c = binary_mut(a, &b, |l, r| l + r).unwrap();
+
+        let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
+        assert_eq!(c, expected);
+    }
+
+    #[test]
+    fn test_try_binary_mut() {
+        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
+        let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
+        let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap();
+
+        let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
+        assert_eq!(c, expected);
+
+        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
+        let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
+        let _ = try_binary_mut(a, &b, |l, r| {
+            if l == 1 {
+                Err(ArrowError::InvalidArgumentError(
+                    "got error".parse().unwrap(),
+                ))
+            } else {
+                Ok(l + r)
+            }
+        })
+        .expect_err("should got error");
+    }
 }