diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index 8024831f74ec..ca5e7a8279be 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -2830,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::::with_capacity(1, 1); + builder.append(i32::MIN).unwrap(); + let a = builder.finish(); + + let mut builder = + PrimitiveDictionaryBuilder::::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 ce02180f2515..751eedeb5b4b 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( + len: usize, + a: A, + b: B, + op: F, +) -> Result> +where + O: ArrowPrimitiveType, + F: Fn(A::Item, B::Item) -> Option, +{ + 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 @@ -386,6 +406,10 @@ where return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE))); } + if a.null_count() == 0 && b.null_count() == 0 { + return Ok(try_binary_opt_no_nulls(a.len(), a, b, op)?); + } + let iter_a = ArrayIter::new(a); let iter_b = ArrayIter::new(b);