HaoYang670 commented on code in PR #2666:
URL: https://github.com/apache/arrow-rs/pull/2666#discussion_r964390038
##########
arrow/src/compute/kernels/arithmetic.rs:
##########
@@ -211,32 +110,35 @@ where
/// This function errors if:
/// * the arrays have different lengths
/// * there is an element where both left and right values are valid and the
right value is `0`
-fn math_checked_divide_op<T, F>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
+fn math_checked_divide_op<LT, RT, F>(
+ left: &PrimitiveArray<LT>,
+ right: &PrimitiveArray<RT>,
op: F,
-) -> Result<PrimitiveArray<T>>
+) -> Result<PrimitiveArray<LT>>
Review Comment:
Why is the return type same as the left type?
##########
arrow/src/compute/kernels/arity.rs:
##########
@@ -198,6 +259,80 @@ where
}
}
+pub fn binary<A, B, F, O>(
Review Comment:
I guess we need to add docs for `binary` and `try_binary`.
##########
arrow/src/compute/kernels/arity.rs:
##########
@@ -198,6 +259,80 @@ where
}
}
+pub fn binary<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) -> O::Native,
+{
+ assert_eq!(a.len(), b.len());
+ let len = a.len();
+
+ if a.is_empty() {
+ return PrimitiveArray::from(ArrayData::new_empty(&O::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())
+ .unwrap_or_default();
+
+ let values = a.values().iter().zip(b.values()).map(|(l, r)| op(*l, *r));
+ // JUSTIFICATION
+ // Benefit
+ // ~60% speedup
+ // Soundness
+ // `values` is an iterator with a known size from a PrimitiveArray
+ let buffer = unsafe { Buffer::from_trusted_len_iter(values) };
+
+ unsafe { build_primitive_array(len, buffer, null_count, null_buffer) }
+}
+
+pub fn try_binary<A, B, F, O>(
+ a: &PrimitiveArray<A>,
+ b: &PrimitiveArray<B>,
+ op: F,
+) -> Result<PrimitiveArray<O>>
+where
+ A: ArrowPrimitiveType,
+ B: ArrowPrimitiveType,
+ O: ArrowPrimitiveType,
+ F: Fn(A::Native, B::Native) -> Result<O::Native>,
+{
+ assert_eq!(a.len(), b.len());
+ let len = a.len();
+
+ if a.is_empty() {
+ return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
+ }
+
+ let null_buffer = combine_option_bitmap(&[a.data(), b.data()],
len).unwrap();
Review Comment:
```suggestion
let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len)?;
```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
