Implement UnionArray logical_nulls

Cargo.toml

@@ -74,7 +74,7 @@ include = [
     "Cargo.toml",
 ]
 edition = "2021"
-rust-version = "1.62"
+rust-version = "1.63"
 
 [workspace.dependencies]
 arrow = { version = "52.2.0", path = "./arrow", default-features = false }

arrow-arith/src/boolean.rs

@@ -352,6 +352,9 @@ pub fn is_not_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
 
 #[cfg(test)]
 mod tests {
+    use arrow_buffer::ScalarBuffer;
+    use arrow_schema::{DataType, Field, UnionFields};
+
     use super::*;
     use std::sync::Arc;
 
@@ -911,4 +914,61 @@ mod tests {
         assert_eq!(expected, res);
         assert!(res.nulls().is_none());
     }
+
+    #[test]
+    fn test_dense_union_is_null() {
+        // union of [{A=1}, {A=}, {B=3.2}, {B=}, {C="a"}, {C=}]
+        let int_array = Int32Array::from(vec![Some(1), None]);
+        let float_array = Float64Array::from(vec![Some(3.2), None]);
+        let str_array = StringArray::from(vec![Some("a"), None]);
+        let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
+        let offsets = [0, 1, 0, 1, 0, 1]
+            .into_iter()
+            .collect::<ScalarBuffer<i32>>();
+
+        let children = vec![
+            Arc::new(int_array) as Arc<dyn Array>,
+            Arc::new(float_array),
+            Arc::new(str_array),
+        ];
+
+        let array = UnionArray::try_new(union_fields(), type_ids, Some(offsets), children).unwrap();
+
+        let result = is_null(&array).unwrap();
+
+        let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
+        assert_eq!(expected, &result);
+    }
+
+    #[test]
+    fn test_sparse_union_is_null() {
+        // union of [{A=1}, {A=}, {B=3.2}, {B=}, {C="a"}, {C=}]
+        let int_array = Int32Array::from(vec![Some(1), None, None, None, None, None]);
+        let float_array = Float64Array::from(vec![None, None, Some(3.2), None, None, None]);
+        let str_array = StringArray::from(vec![None, None, None, None, Some("a"), None]);
+        let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
+
+        let children = vec![
+            Arc::new(int_array) as Arc<dyn Array>,
+            Arc::new(float_array),
+            Arc::new(str_array),
+        ];
+
+        let array = UnionArray::try_new(union_fields(), type_ids, None, children).unwrap();
+
+        let result = is_null(&array).unwrap();
+
+        let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
+        assert_eq!(expected, &result);
+    }
+
+    fn union_fields() -> UnionFields {
+        [
+            (0, Arc::new(Field::new("A", DataType::Int32, true))),
+            (1, Arc::new(Field::new("B", DataType::Float64, true))),
+            (2, Arc::new(Field::new("C", DataType::Utf8, true))),
+        ]
+        .into_iter()
+        .collect()
+    }
 }

arrow-array/Cargo.toml

@@ -71,3 +71,7 @@ harness = false
 [[bench]]
 name = "fixed_size_list_array"
 harness = false
+
+[[bench]]
+name = "union_array"
+harness = false

arrow-array/benches/union_array.rs

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::{iter::repeat_n, sync::Arc};
+
+use arrow_array::{Array, ArrayRef, Int32Array, UnionArray};
+use arrow_schema::{DataType, Field, UnionFields};
+use criterion::*;
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let child = Arc::new(Int32Array::new(
+        (0..4096).collect(),
+        Some(repeat_n([true, false], 2048).flatten().collect()),
+    )) as ArrayRef;
+
+    for i in 1..5 {
+        c.bench_function(&format!("union logical nulls 4096 {i} children"), |b| {
+            let fields = UnionFields::new(
+                0..i,
+                (0..i).map(|i| Field::new(format!("f{i}"), DataType::Int32, true)),
+            );
+
+            let array = UnionArray::try_new(
+                fields,
+                (0..i).cycle().take(4096).collect(),
+                None,
+                repeat_n(child.clone(), i as usize).collect(),
+            )
+            .unwrap();
+
+            b.iter(|| black_box(array.logical_nulls()))
+        });
+    }
+
+    c.bench_function("single with nulls 4096", |b| {
+        let fields = UnionFields::new(
+            [1, 3],
+            [
+                Field::new("a", DataType::Int32, true),
+                Field::new("b", DataType::Int32, true),
+            ],
+        );
+
+        let array = UnionArray::try_new(
+            fields,
+            repeat_n([1, 3], 2048).flatten().collect(),
+            None,
+            vec![child.clone(), Arc::new(Int32Array::from_value(-5, 4096))],
+        )
+        .unwrap();
+
+        b.iter(|| black_box(array.logical_nulls()))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);