diff --git a/arrow-arith/src/arity.rs b/arrow-arith/src/arity.rs index f15ea955eaaa..42aa5b08891c 100644 --- a/arrow-arith/src/arity.rs +++ b/arrow-arith/src/arity.rs @@ -270,11 +270,17 @@ where } let len = a.len(); - if a.null_count() == 0 && b.null_count() == 0 { + // Physical nulls are not the whole story: a `RunArray` or `DictionaryArray` can have + // logical nulls in its values while its own null buffer is absent. `is_nullable` covers + // those, but is allowed to be conservative, so the union of the logical nulls can still + // be empty. + if !a.is_nullable() && !b.is_nullable() { try_binary_no_nulls(len, a, b, op) } else { - let nulls = - NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()).unwrap(); + let Some(nulls) = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()) + else { + return try_binary_no_nulls(len, a, b, op); + }; let mut buffer = BufferBuilder::::new(len); buffer.append_n_zeroed(len); @@ -324,12 +330,17 @@ where )))); } - if a.null_count() == 0 && b.null_count() == 0 { + // Physical and logical nulls coincide for a `PrimitiveArray`, but gate on `is_nullable` to + // match `try_binary`. That check is allowed to be conservative, so fall back to the no-nulls + // path when the union of the logical nulls turns out to be empty, instead of unwrapping it. + if !a.is_nullable() && !b.is_nullable() { try_binary_no_nulls_mut(len, a, b, op) } else { - let nulls = + let Some(nulls) = create_union_null_buffer(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()) - .unwrap(); + else { + return try_binary_no_nulls_mut(len, a, b, op); + }; let mut builder = a.into_builder()?; @@ -431,6 +442,22 @@ mod tests { ); } + #[test] + fn test_try_binary_run_array_logical_nulls() { + // A `RunArray` has no null buffer of its own, so its logical nulls live in the values. + let run_ends = Int32Array::from(vec![1, 2, 3]); + let values = Int32Array::from(vec![Some(10), None, Some(30)]); + let run = RunArray::::try_new(&run_ends, &values).expect("valid run array"); + assert_eq!(run.null_count(), 0); + assert_eq!(run.logical_null_count(), 1); + + let typed = run.downcast::().expect("Int32 values"); + let other = Int32Array::from(vec![1, 1, 1]); + let result = + try_binary::<_, _, _, Int32Type>(typed, &other, |a, b| Ok(a + b)).expect("no overflow"); + assert_eq!(result, Int32Array::from(vec![Some(11), None, Some(31)])); + } + #[test] fn test_binary_mut() { let a = Int32Array::from(vec![15, 14, 9, 8, 1]); @@ -509,6 +536,18 @@ mod tests { assert_eq!(r1.unwrap(), r2.unwrap()); } + #[test] + fn test_try_binary_mut_all_valid_null_buffers() { + // Both arrays carry a null buffer with no nulls in it: the no-nulls path must still run. + let a = Int32Array::new(vec![1, 2].into(), Some(vec![true, true].into())); + let b = Int32Array::new(vec![10, 20].into(), Some(vec![true, true].into())); + let c = try_binary_mut(a, &b, |a, b| Ok(a + b)) + .expect("not shared") + .expect("no overflow"); + assert_eq!(c, Int32Array::from(vec![11, 22])); + assert_eq!(c.logical_null_count(), 0); + } + #[test] fn test_unary_dict_mut() { let values = Int32Array::from(vec![Some(10), Some(20), None]);