This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new dd5a2294b Improve UnionArray::is_nullable (#6540)
dd5a2294b is described below
commit dd5a2294b8b28f768b991e0e89fe7686b296c4ec
Author: Raphael Taylor-Davies <[email protected]>
AuthorDate: Fri Oct 18 12:14:20 2024 +0100
Improve UnionArray::is_nullable (#6540)
* Improve UnionArray::is_nullable
* Add specific unit test for is_nullable
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
arrow-array/src/array/list_array.rs | 16 ++++++++++++-
arrow-array/src/array/union_array.rs | 44 +++++++++++++++++++++++++++++++++++-
2 files changed, 58 insertions(+), 2 deletions(-)
diff --git a/arrow-array/src/array/list_array.rs
b/arrow-array/src/array/list_array.rs
index dc1b9f07d..13b446b03 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<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
index 1feef8c56..3c6da5a7b 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::<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()
+ }
}
