[GitHub] [arrow-rs] sunchao commented on a diff in pull request #2710: Add divide_opt kernel which produce null values on division by zero error

[GitBox <gi...@apache.org>]

sunchao commented on code in PR #2710:
URL: https://github.com/apache/arrow-rs/pull/2710#discussion_r969032401


##########
arrow/src/compute/kernels/arity.rs:
##########
@@ -257,6 +257,50 @@ where
     Ok(unsafe { build_primitive_array(len, buffer.finish(), null_count, null_buffer) })
 }
 
+/// 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
+/// results in a new null in the collected [`PrimitiveArray`].
+///
+/// The function is only evaluated for non-null indices
+///
+/// # Panic
+///
+/// Panics if the arrays have different lengths
+pub(crate) fn binary_opt<A, B, F, O>(
+    a: &PrimitiveArray<A>,
+    b: &PrimitiveArray<B>,
+    op: F,
+) -> PrimitiveArray<O>
+where
+    A: ArrowPrimitiveType,
+    B: ArrowPrimitiveType,
+    O: ArrowPrimitiveType,
+    F: Fn(A::Native, B::Native) -> Option<O::Native>,
+{
+    assert_eq!(a.len(), b.len());
+
+    if a.is_empty() {
+        return PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE));
+    }
+
+    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) {

Review Comment:
   what if both `a` and `b` are all non-null? can we implement a fast-path for that case?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org