Improve UnionArray::is_nullable (apache#6540)

* Improve UnionArray::is_nullable * Add specific unit test for is_nullable --------- Co-authored-by: Andrew Lamb <[email protected]>
integer32llc · Oct 18, 2024 · dd5a229 · dd5a229
1 parent 9485897
commit dd5a229
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 2 deletions.
diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs
@@ -549,7 +549,7 @@ pub type LargeListArray = GenericListArray<i64>;
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::builder::{FixedSizeListBuilder, Int32Builder, ListBuilder};
+    use crate::builder::{FixedSizeListBuilder, Int32Builder, ListBuilder, UnionBuilder};
     use crate::cast::AsArray;
     use crate::types::Int32Type;
     use crate::{Int32Array, Int64Array};
@@ -1181,4 +1181,18 @@ mod tests {
             .collect();
         assert_eq!(values, vec![Some(vec![1, 2, 3]), None, Some(vec![4, 5, 6])])
     }
+
+    #[test]
+    fn test_nullable_union() {
+        let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
+        let mut builder = UnionBuilder::new_dense();
+        builder.append::<Int32Type>("a", 1).unwrap();
+        builder.append::<Int32Type>("b", 2).unwrap();
+        builder.append::<Int32Type>("b", 3).unwrap();
+        builder.append::<Int32Type>("a", 4).unwrap();
+        builder.append::<Int32Type>("a", 5).unwrap();
+        let values = builder.build().unwrap();
+        let field = Arc::new(Field::new("element", values.data_type().clone(), false));
+        ListArray::new(field.clone(), offsets, Arc::new(values), None);
+    }
 }
diff --git a/arrow-array/src/array/union_array.rs b/arrow-array/src/array/union_array.rs
@@ -875,7 +875,10 @@ impl Array for UnionArray {
     }
 
     fn is_nullable(&self) -> bool {
-        true
+        self.fields
+            .iter()
+            .flatten()
+            .any(|field| field.is_nullable())
     }
 
     fn get_buffer_memory_size(&self) -> usize {
@@ -2138,4 +2141,43 @@ mod tests {
         .into_iter()
         .collect()
     }
+
+    #[test]
+    fn test_is_nullable() {
+        assert!(!create_union_array(false, false).is_nullable());
+        assert!(create_union_array(true, false).is_nullable());
+        assert!(create_union_array(false, true).is_nullable());
+        assert!(create_union_array(true, true).is_nullable());
+    }
+
+    /// Create a union array with a float and integer field
+    ///
+    /// If the `int_nullable` is true, the integer field will have nulls
+    /// If the `float_nullable` is true, the float field will have nulls
+    ///
+    /// Note the `Field` definitions are always declared to be nullable
+    fn create_union_array(int_nullable: bool, float_nullable: bool) -> UnionArray {
+        let int_array = if int_nullable {
+            Int32Array::from(vec![Some(1), None, Some(3)])
+        } else {
+            Int32Array::from(vec![1, 2, 3])
+        };
+        let float_array = if float_nullable {
+            Float64Array::from(vec![Some(3.2), None, Some(4.2)])
+        } else {
+            Float64Array::from(vec![3.2, 4.2, 5.2])
+        };
+        let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
+        let offsets = [0, 0, 0].into_iter().collect::<ScalarBuffer<i32>>();
+        let union_fields = [
+            (0, Arc::new(Field::new("A", DataType::Int32, true))),
+            (1, Arc::new(Field::new("B", DataType::Float64, true))),
+        ]
+        .into_iter()
+        .collect::<UnionFields>();
+
+        let children = vec![Arc::new(int_array) as Arc<dyn Array>, Arc::new(float_array)];
+
+        UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap()
+    }
 }