diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs index dc1b9f07da1..13b446b0382 100644 --- a/arrow-array/src/array/list_array.rs +++ b/arrow-array/src/array/list_array.rs @@ -549,7 +549,7 @@ pub type LargeListArray = GenericListArray; #[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::("a", 1).unwrap(); + builder.append::("b", 2).unwrap(); + builder.append::("b", 3).unwrap(); + builder.append::("a", 4).unwrap(); + builder.append::("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 index 1feef8c5675..3c6da5a7b5c 100644 --- 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::>(); + let offsets = [0, 0, 0].into_iter().collect::>(); + let union_fields = [ + (0, Arc::new(Field::new("A", DataType::Int32, true))), + (1, Arc::new(Field::new("B", DataType::Float64, true))), + ] + .into_iter() + .collect::(); + + let children = vec![Arc::new(int_array) as Arc, Arc::new(float_array)]; + + UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap() + } }