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 80c0f1a48 Add divide dyn kernel which produces null for division by
zero (#2764)
80c0f1a48 is described below
commit 80c0f1a480ee51bc05a2224108575dfad1be8deb
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Wed Sep 21 19:12:24 2022 -0700
Add divide dyn kernel which produces null for division by zero (#2764)
* Add divide_dyn_opt kernel
* Add test
* Fix clippy
* Rename function
---
arrow/src/compute/kernels/arithmetic.rs | 102 ++++++++++++++++++++++++++++++++
arrow/src/compute/kernels/arity.rs | 69 ++++++++++++---------
2 files changed, 143 insertions(+), 28 deletions(-)
diff --git a/arrow/src/compute/kernels/arithmetic.rs
b/arrow/src/compute/kernels/arithmetic.rs
index b44cb8b94..d33827594 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -697,6 +697,39 @@ where
)
}
+#[cfg(feature = "dyn_arith_dict")]
+fn math_divide_safe_op_dict<K, T, F>(
+ left: &DictionaryArray<K>,
+ right: &DictionaryArray<K>,
+ op: F,
+) -> Result<ArrayRef>
+where
+ K: ArrowNumericType,
+ T: ArrowNumericType,
+ T::Native: One + Zero,
+ F: Fn(T::Native, T::Native) -> Option<T::Native>,
+{
+ let left = left.downcast_dict::<PrimitiveArray<T>>().unwrap();
+ let right = right.downcast_dict::<PrimitiveArray<T>>().unwrap();
+ let array: PrimitiveArray<T> = binary_opt::<_, _, _, T>(left, right, op)?;
+ Ok(Arc::new(array) as ArrayRef)
+}
+
+fn math_safe_divide_op<LT, RT, F>(
+ left: &PrimitiveArray<LT>,
+ right: &PrimitiveArray<RT>,
+ op: F,
+) -> Result<ArrayRef>
+where
+ LT: ArrowNumericType,
+ RT: ArrowNumericType,
+ RT::Native: One + Zero,
+ F: Fn(LT::Native, RT::Native) -> Option<LT::Native>,
+{
+ let array: PrimitiveArray<LT> = binary_opt::<_, _, _, LT>(left, right,
op)?;
+ Ok(Arc::new(array) as ArrayRef)
+}
+
/// Perform `left + right` operation on two arrays. If either left or right
value is null
/// then the result is also null.
///
@@ -1406,6 +1439,51 @@ pub fn divide_dyn_checked(left: &dyn Array, right: &dyn
Array) -> Result<ArrayRe
}
}
+/// Perform `left / right` operation on two arrays. If either left or right
value is null
+/// then the result is also null.
+///
+/// If any right hand value is zero, the operation value will be replaced with
null in the
+/// result.
+///
+/// Unlike `divide_dyn` or `divide_dyn_checked`, division by zero will get a
null value instead
+/// returning an `Err`, this also doesn't check overflowing, overflowing will
just wrap
+/// the result around.
+pub fn divide_dyn_opt(left: &dyn Array, right: &dyn Array) -> Result<ArrayRef>
{
+ match left.data_type() {
+ DataType::Dictionary(_, _) => {
+ typed_dict_math_op!(
+ left,
+ right,
+ |a, b| {
+ if b.is_zero() {
+ None
+ } else {
+ Some(a.div_wrapping(b))
+ }
+ },
+ math_divide_safe_op_dict
+ )
+ }
+ _ => {
+ downcast_primitive_array!(
+ (left, right) => {
+ math_safe_divide_op(left, right, |a, b| {
+ if b.is_zero() {
+ None
+ } else {
+ Some(a.div_wrapping(b))
+ }
+ })
+ }
+ _ => Err(ArrowError::CastError(format!(
+ "Unsupported data type {}, {}",
+ left.data_type(), right.data_type()
+ )))
+ )
+ }
+ }
+}
+
/// Perform `left / right` operation on two arrays without checking for
division by zero.
/// For floating point types, the result of dividing by zero follows normal
floating point
/// rules. For other numeric types, dividing by zero will panic,
@@ -2752,4 +2830,28 @@ mod tests {
let overflow = divide_dyn_checked(&a, &b);
overflow.expect_err("overflow should be detected");
}
+
+ #[test]
+ #[cfg(feature = "dyn_arith_dict")]
+ fn test_div_dyn_opt_overflow_division_by_zero() {
+ let a = Int32Array::from(vec![i32::MIN]);
+ let b = Int32Array::from(vec![0]);
+
+ let division_by_zero = divide_dyn_opt(&a, &b);
+ let expected = Arc::new(Int32Array::from(vec![None])) as ArrayRef;
+ assert_eq!(&expected, &division_by_zero.unwrap());
+
+ let mut builder =
+ PrimitiveDictionaryBuilder::<Int8Type,
Int32Type>::with_capacity(1, 1);
+ builder.append(i32::MIN).unwrap();
+ let a = builder.finish();
+
+ let mut builder =
+ PrimitiveDictionaryBuilder::<Int8Type,
Int32Type>::with_capacity(1, 1);
+ builder.append(0).unwrap();
+ let b = builder.finish();
+
+ let division_by_zero = divide_dyn_opt(&a, &b);
+ assert_eq!(&expected, &division_by_zero.unwrap());
+ }
}
diff --git a/arrow/src/compute/kernels/arity.rs
b/arrow/src/compute/kernels/arity.rs
index bf1028968..5f875e6dd 100644
--- a/arrow/src/compute/kernels/arity.rs
+++ b/arrow/src/compute/kernels/arity.rs
@@ -357,6 +357,26 @@ where
Ok(unsafe { build_primitive_array(len, buffer.into(), 0, None) })
}
+#[inline(never)]
+fn try_binary_opt_no_nulls<A: ArrayAccessor, B: ArrayAccessor, F, O>(
+ len: usize,
+ a: A,
+ b: B,
+ op: F,
+) -> Result<PrimitiveArray<O>>
+where
+ O: ArrowPrimitiveType,
+ F: Fn(A::Item, B::Item) -> Option<O::Native>,
+{
+ let mut buffer = Vec::with_capacity(10);
+ for idx in 0..len {
+ unsafe {
+ buffer.push(op(a.value_unchecked(idx), b.value_unchecked(idx)));
+ };
+ }
+ Ok(buffer.iter().collect())
+}
+
/// Applies the provided binary operation across `a` and `b`, collecting the
optional results
/// into a [`PrimitiveArray`]. If any index is null in either `a` or `b`, the
corresponding
/// index in the result will also be null. The binary operation could return
`None` which
@@ -367,16 +387,14 @@ where
/// # Error
///
/// This function gives error if the arrays have different lengths
-pub(crate) fn binary_opt<A, B, F, O>(
- a: &PrimitiveArray<A>,
- b: &PrimitiveArray<B>,
+pub(crate) fn binary_opt<A: ArrayAccessor + Array, B: ArrayAccessor + Array,
F, O>(
+ a: A,
+ b: B,
op: F,
) -> Result<PrimitiveArray<O>>
where
- A: ArrowPrimitiveType,
- B: ArrowPrimitiveType,
O: ArrowPrimitiveType,
- F: Fn(A::Native, B::Native) -> Option<O::Native>,
+ F: Fn(A::Item, B::Item) -> Option<O::Native>,
{
if a.len() != b.len() {
return Err(ArrowError::ComputeError(
@@ -389,29 +407,24 @@ where
}
if a.null_count() == 0 && b.null_count() == 0 {
- Ok(a.values()
- .iter()
- .zip(b.values().iter())
- .map(|(a, b)| op(*a, *b))
- .collect())
- } else {
- let iter_a = ArrayIter::new(a);
- let iter_b = ArrayIter::new(b);
-
- let values =
- iter_a
- .into_iter()
- .zip(iter_b.into_iter())
- .map(|(item_a, item_b)| {
- if let (Some(a), Some(b)) = (item_a, item_b) {
- op(a, b)
- } else {
- None
- }
- });
-
- Ok(values.collect())
+ return try_binary_opt_no_nulls(a.len(), a, b, op);
}
+
+ let iter_a = ArrayIter::new(a);
+ let iter_b = ArrayIter::new(b);
+
+ let values = iter_a
+ .into_iter()
+ .zip(iter_b.into_iter())
+ .map(|(item_a, item_b)| {
+ if let (Some(a), Some(b)) = (item_a, item_b) {
+ op(a, b)
+ } else {
+ None
+ }
+ });
+
+ Ok(values.collect())
}
#[cfg(test)]
