fix: try_binary dropped the logical nulls of run and dictionary arrays

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::<O::Native>::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::<Int32Type>::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::<Int32Array>().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]);