Rich-T-kid commented on code in PR #12076:
URL: https://github.com/apache/arrow-rs/pull/12076#discussion_r4556827481
##########
arrow-arith/src/numeric.rs:
##########
@@ -1198,6 +1313,118 @@ mod tests {
date.signed_duration_since(EPOCH).num_days() as i32
}
+ fn assert_kernel_result<T: ArrowPrimitiveType>(
+ actual: Result<ArrayRef, ArrowError>,
+ expected: Result<Vec<Option<T::Native>>, ArrowError>,
+ context: &str,
+ ) {
+ match (actual, expected) {
+ (Ok(actual), Ok(expected)) => {
+ let actual: Vec<_> =
actual.as_primitive::<T>().iter().collect();
+ assert_eq!(actual, expected, "{context}");
+ }
+ (Err(actual), Err(expected)) => {
+ assert_eq!(actual.to_string(), expected.to_string(),
"{context}")
+ }
+ (actual, expected) => panic!("{context}: {actual:?} !=
{expected:?}"),
+ }
+ }
+
+ /// Checks `div` and `rem` of `values` by each of `divisors` and by a null
divisor, against
+ /// `div_checked` and `mod_wrapping` on each valid value, as `integer_op`
computes them.
+ /// `values` is repeated to 64 values so that the vectorized loop runs,
and is also checked
+ /// sliced to 50 values, which leaves a remainder for the scalar loop,
with seven in eight
+ /// values null, all null, and empty
+ fn test_div_rem_by_scalar<T: ArrowPrimitiveType>(
+ values: &[Option<T::Native>],
+ divisors: &[T::Native],
+ ) {
+ let values: PrimitiveArray<T> =
values.iter().cycle().take(64).collect();
+ let mostly_null = NullBuffer::from_iter((0..values.len()).map(|i| i %
8 == 0));
+ let arrays = [
+ values.slice(3, 50),
+ PrimitiveArray::new(values.values().clone(), Some(mostly_null)),
+ values.slice(0, 0),
+ PrimitiveArray::new_null(2),
+ values,
+ ];
+ for array in &arrays {
+ for &d in divisors {
+ let divisor = PrimitiveArray::<T>::new_scalar(d);
+ let expected = array
+ .iter()
+ .map(|v| v.map(|v| v.div_checked(d)).transpose())
+ .collect();
+ assert_kernel_result::<T>(div(array, &divisor), expected,
&format!("/ {d:?}"));
+ let expected = array
+ .iter()
+ .map(|v| {
+ v.map(|v| match d.is_zero() {
+ true => Err(ArrowError::DivideByZero),
+ false => Ok(v.mod_wrapping(d)),
+ })
+ .transpose()
+ })
+ .collect();
+ assert_kernel_result::<T>(rem(array, &divisor), expected,
&format!("% {d:?}"));
+ }
+ // A null divisor gives a null result, whatever its slot holds
+ for slot in [0, 7] {
+ let slot = vec![T::Native::usize_as(slot)].into();
+ let divisor = Scalar::new(PrimitiveArray::<T>::new(
+ slot,
+ Some(NullBuffer::new_null(1)),
+ ));
+ assert_eq!(div(array, &divisor).unwrap().null_count(),
array.len());
+ assert_eq!(rem(array, &divisor).unwrap().null_count(),
array.len());
+ }
+ }
+ }
+
+ #[test]
+ fn test_div_rem_by_scalar_integer() {
+ // 16, 13 and 7 reach each variant of the precomputed divisor. 0, and
-1 for `/`, fall back
+ // to `integer_op`, which returns an error for a valid value and `Ok`
otherwise
+ let values = [
+ Some(i32::MIN),
+ Some(i32::MIN + 1),
+ Some(-7),
+ None,
+ Some(0),
+ Some(13),
+ Some(i32::MAX),
+ ];
+ test_div_rem_by_scalar::<Int32Type>(&values, &[16, 13, 7, 0, -1]);
Review Comment:
nit: Could we include a negative optimized divisor such as -7 or -16 here?
The signed fast path currently gets direct kernel coverage only for positive
divisors, while -1 exercises the fallback.
--
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]