From 1c6504c3a772442e9c5c0f79dfca222622929236 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 19 Nov 2022 00:23:39 -0800 Subject: [PATCH 1/6] Add add_mut --- arrow/src/compute/kernels/arithmetic.rs | 24 +++++++- arrow/src/compute/kernels/arity.rs | 77 +++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 1 deletion(-) diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index a99a90204b7f..b44147afa3ca 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -27,7 +27,7 @@ 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_unary, try_unary_dyn, unary_dyn, }; use crate::datatypes::{ ArrowNativeTypeOp, ArrowNumericType, DataType, Date32Type, Date64Type, @@ -733,6 +733,28 @@ where math_op(left, right, |a, b| a.add_wrapping(b)) } +/// Perform `left + right` operation on two arrays while 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( + left: PrimitiveArray, + right: &PrimitiveArray, +) -> std::result::Result< + PrimitiveArray, + std::result::Result, 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. /// diff --git a/arrow/src/compute/kernels/arity.rs b/arrow/src/compute/kernels/arity.rs index c99d2b727b8d..6967c85fd4e7 100644 --- a/arrow/src/compute/kernels/arity.rs +++ b/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( + a: PrimitiveArray, + b: &PrimitiveArray, + op: F, +) -> std::result::Result< + PrimitiveArray, + std::result::Result, 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(|arr| Ok(arr))?; + + 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::::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 @@ -358,6 +424,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 +484,14 @@ 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); + } } From 526abd578ac66082625505810a1f0fd63630ad3a Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 19 Nov 2022 11:05:27 -0800 Subject: [PATCH 2/6] Add try_binary_mut --- arrow/src/compute/kernels/arity.rs | 124 +++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) diff --git a/arrow/src/compute/kernels/arity.rs b/arrow/src/compute/kernels/arity.rs index 6967c85fd4e7..0cedcce64049 100644 --- a/arrow/src/compute/kernels/arity.rs +++ b/arrow/src/compute/kernels/arity.rs @@ -328,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( + a: PrimitiveArray, + b: &PrimitiveArray, + op: F, +) -> std::result::Result< + PrimitiveArray, + std::result::Result, ArrowError>, +> +where + T: ArrowPrimitiveType, + F: Fn(T::Native, T::Native) -> Result, +{ + 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(|arr| Ok(arr))?; + + 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| 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::::from(array_data)) + } +} + /// This intentional inline(never) attribute helps LLVM optimize the loop. #[inline(never)] fn try_binary_no_nulls( @@ -349,6 +422,34 @@ 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( + len: usize, + a: PrimitiveArray, + b: &PrimitiveArray, + op: F, +) -> std::result::Result< + PrimitiveArray, + std::result::Result, ArrowError>, +> +where + T: ArrowPrimitiveType, + F: Fn(T::Native, T::Native) -> Result, +{ + let mut builder = a.into_builder().map_err(|arr| Ok(arr))?; + 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| Err(err))?; + }; + } + Ok(builder.finish()) +} + #[inline(never)] fn try_binary_opt_no_nulls( len: usize, @@ -494,4 +595,27 @@ mod tests { 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"); + } } From c1e3f4d092622ab303f938a306ad2c499a6fde66 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 19 Nov 2022 11:13:12 -0800 Subject: [PATCH 3/6] Add test --- arrow/src/compute/kernels/arithmetic.rs | 52 ++++++++++++++++++++++++- arrow/src/compute/kernels/arity.rs | 11 +++--- 2 files changed, 55 insertions(+), 8 deletions(-) diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index b44147afa3ca..c368a306603c 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/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_mut, 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,7 +734,7 @@ where math_op(left, right, |a, b| a.add_wrapping(b)) } -/// Perform `left + right` operation on two arrays while mutating `left` with operation results. +/// 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 @@ -771,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( + left: PrimitiveArray, + right: &PrimitiveArray, +) -> std::result::Result< + PrimitiveArray, + std::result::Result, 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. /// @@ -3120,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"); + } } diff --git a/arrow/src/compute/kernels/arity.rs b/arrow/src/compute/kernels/arity.rs index 0cedcce64049..152be00e675b 100644 --- a/arrow/src/compute/kernels/arity.rs +++ b/arrow/src/compute/kernels/arity.rs @@ -251,7 +251,7 @@ where .map(|x| len - x.count_set_bits_offset(0, len)) .unwrap_or_default(); - let mut builder = a.into_builder().map_err(|arr| Ok(arr))?; + let mut builder = a.into_builder().map_err(Ok)?; builder .values_slice_mut() @@ -375,7 +375,7 @@ where .map(|x| len - x.count_set_bits_offset(0, len)) .unwrap_or_default(); - let mut builder = a.into_builder().map_err(|arr| Ok(arr))?; + let mut builder = a.into_builder().map_err(Ok)?; let slice = builder.values_slice_mut(); @@ -386,7 +386,7 @@ where }; Ok::<_, ArrowError>(()) }) - .map_err(|err| Err(err))?; + .map_err(Err)?; let array_builder = builder .finish() @@ -437,14 +437,13 @@ where T: ArrowPrimitiveType, F: Fn(T::Native, T::Native) -> Result, { - let mut builder = a.into_builder().map_err(|arr| Ok(arr))?; + 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| Err(err))?; + op(*slice.get_unchecked(idx), b.value_unchecked(idx)).map_err(Err)?; }; } Ok(builder.finish()) From 7191dba516797158d08d4a4b5772a852a4fdd6a8 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 26 Nov 2022 19:32:06 -0800 Subject: [PATCH 4/6] Change result type --- arrow/src/compute/kernels/arithmetic.rs | 10 ++--- arrow/src/compute/kernels/arity.rs | 58 ++++++++++++++++--------- 2 files changed, 42 insertions(+), 26 deletions(-) diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index c368a306603c..9d2d567d9d64 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -746,8 +746,8 @@ pub fn add_mut( left: PrimitiveArray, right: &PrimitiveArray, ) -> std::result::Result< - PrimitiveArray, std::result::Result, ArrowError>, + PrimitiveArray, > where T: ArrowNumericType, @@ -784,8 +784,8 @@ pub fn add_checked_mut( left: PrimitiveArray, right: &PrimitiveArray, ) -> std::result::Result< - PrimitiveArray, std::result::Result, ArrowError>, + PrimitiveArray, > where T: ArrowNumericType, @@ -3149,7 +3149,7 @@ mod tests { 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 c = add_mut(a, &b).unwrap().unwrap(); let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]); assert_eq!(c, expected); } @@ -3159,13 +3159,13 @@ mod tests { let a = Int32Array::from(vec![i32::MAX, i32::MIN]); let b = Int32Array::from(vec![1, 1]); - let wrapped = add_mut(a, &b).unwrap(); + let wrapped = add_mut(a, &b).unwrap().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"); + let _ = overflow.unwrap().expect_err("overflow should be detected"); } } diff --git a/arrow/src/compute/kernels/arity.rs b/arrow/src/compute/kernels/arity.rs index 152be00e675b..452426a678cf 100644 --- a/arrow/src/compute/kernels/arity.rs +++ b/arrow/src/compute/kernels/arity.rs @@ -227,31 +227,34 @@ pub fn binary_mut( b: &PrimitiveArray, op: F, ) -> std::result::Result< - PrimitiveArray, std::result::Result, ArrowError>, + PrimitiveArray, > where T: ArrowPrimitiveType, F: Fn(T::Native, T::Native) -> T::Native, { if a.len() != b.len() { - return Err(Err(ArrowError::ComputeError( + return Ok(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))); + return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty( + &T::DATA_TYPE, + )))); } + let len = a.len(); + 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 mut builder = a.into_builder()?; builder .values_slice_mut() @@ -268,7 +271,7 @@ where .null_count(null_count); let array_data = unsafe { array_builder.build_unchecked() }; - Ok(PrimitiveArray::::from(array_data)) + Ok(Ok(PrimitiveArray::::from(array_data))) } /// Applies the provided fallible binary operation across `a` and `b`, returning any error, @@ -348,22 +351,24 @@ pub fn try_binary_mut( b: &PrimitiveArray, op: F, ) -> std::result::Result< - PrimitiveArray, std::result::Result, ArrowError>, + PrimitiveArray, > where T: ArrowPrimitiveType, F: Fn(T::Native, T::Native) -> Result, { if a.len() != b.len() { - return Err(Err(ArrowError::ComputeError( + return Ok(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))); + return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty( + &T::DATA_TYPE, + )))); } if a.null_count() == 0 && b.null_count() == 0 { @@ -375,18 +380,20 @@ where .map(|x| len - x.count_set_bits_offset(0, len)) .unwrap_or_default(); - let mut builder = a.into_builder().map_err(Ok)?; + let mut builder = a.into_builder()?; let slice = builder.values_slice_mut(); - try_for_each_valid_idx(len, 0, null_count, null_buffer.as_deref(), |idx| { + match 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)?; + }) { + Ok(_) => {} + Err(err) => return Ok(Err(err)), + }; let array_builder = builder .finish() @@ -397,7 +404,7 @@ where .null_count(null_count); let array_data = unsafe { array_builder.build_unchecked() }; - Ok(PrimitiveArray::::from(array_data)) + Ok(Ok(PrimitiveArray::::from(array_data))) } } @@ -430,23 +437,25 @@ fn try_binary_no_nulls_mut( b: &PrimitiveArray, op: F, ) -> std::result::Result< - PrimitiveArray, std::result::Result, ArrowError>, + PrimitiveArray, > where T: ArrowPrimitiveType, F: Fn(T::Native, T::Native) -> Result, { - let mut builder = a.into_builder().map_err(Ok)?; + let mut builder = a.into_builder()?; 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)?; + match op(*slice.get_unchecked(idx), b.value_unchecked(idx)) { + Ok(value) => *slice.get_unchecked_mut(idx) = value, + Err(err) => return Ok(Err(err)), + }; }; } - Ok(builder.finish()) + Ok(Ok(builder.finish())) } #[inline(never)] @@ -589,7 +598,7 @@ mod tests { 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 c = binary_mut(a, &b, |l, r| l + r).unwrap().unwrap(); let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]); assert_eq!(c, expected); @@ -599,11 +608,17 @@ mod tests { 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 c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().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![1, 2, 3, 4, 5]); + let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().unwrap(); + let expected = Int32Array::from(vec![16, 16, 12, 12, 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| { @@ -615,6 +630,7 @@ mod tests { Ok(l + r) } }) + .unwrap() .expect_err("should got error"); } } From 0b672be43905f49c3f3e296bf6c0c42db6c18fb6 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 28 Nov 2022 16:05:22 -0800 Subject: [PATCH 5/6] Remove _mut kernels --- arrow/src/compute/kernels/arithmetic.rs | 58 ++++--------------------- 1 file changed, 9 insertions(+), 49 deletions(-) diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index 9d2d567d9d64..deb24602fecd 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -734,28 +734,6 @@ 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( - left: PrimitiveArray, - right: &PrimitiveArray, -) -> std::result::Result< - std::result::Result, ArrowError>, - PrimitiveArray, -> -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. /// @@ -772,28 +750,6 @@ 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( - left: PrimitiveArray, - right: &PrimitiveArray, -) -> std::result::Result< - std::result::Result, ArrowError>, - PrimitiveArray, -> -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. /// @@ -3145,27 +3101,31 @@ mod tests { } #[test] - fn test_primitive_array_add_mut() { + fn test_primitive_array_add_mut_by_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 = add_mut(a, &b).unwrap().unwrap(); + let c = binary_mut(a, &b, |a, b| a.add_wrapping(b)) + .unwrap() + .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() { + fn test_primitive_add_mut_wrapping_overflow_by_try_binary_mut() { let a = Int32Array::from(vec![i32::MAX, i32::MIN]); let b = Int32Array::from(vec![1, 1]); - let wrapped = add_mut(a, &b).unwrap().unwrap(); + let wrapped = binary_mut(a, &b, |a, b| a.add_wrapping(b)) + .unwrap() + .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 = try_binary_mut(a, &b, |a, b| a.add_checked(b)); let _ = overflow.unwrap().expect_err("overflow should be detected"); } } From fccee61d6770a6802c800a703c4babc206669a93 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 28 Nov 2022 16:14:02 -0800 Subject: [PATCH 6/6] Fix clippy --- arrow/src/compute/kernels/arithmetic.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index d18ac2a39728..c57e27095c23 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -27,8 +27,7 @@ use crate::array::*; use crate::buffer::MutableBuffer; use crate::compute::kernels::arity::unary; use crate::compute::{ - binary, binary_mut, binary_opt, try_binary, try_binary_mut, try_unary, try_unary_dyn, - unary_dyn, + binary, binary_opt, try_binary, try_unary, try_unary_dyn, unary_dyn, }; use crate::datatypes::{ ArrowNativeTypeOp, ArrowNumericType, DataType, Date32Type, Date64Type, @@ -1625,7 +1624,7 @@ where mod tests { use super::*; use crate::array::Int32Array; - use crate::compute::{try_unary_mut, unary_mut}; + use crate::compute::{binary_mut, try_binary_mut, try_unary_mut, unary_mut}; use crate::datatypes::{Date64Type, Int32Type, Int8Type}; use arrow_buffer::i256; use chrono::NaiveDate; @@ -3148,7 +3147,6 @@ mod tests { assert_eq!(expected, wrapped); let a = Int32Array::from(vec![i32::MAX, i32::MIN]); - let b = Int32Array::from(vec![1, 1]); let overflow = try_unary_mut(a, |value| value.add_checked(1)); let _ = overflow.unwrap().expect_err("overflow should be detected"); }