Skip to content

Commit ecd340d

Browse files
Jefffreyahlag
andauthored
[57_maintenance] fix(arrow-array): use the field's type id in dense union logical_nulls (#11810)
- Backport #11643 - For inclusion in #11666 Cherry-pick of 24c53c7 Co-authored-by: Chang Yong Lik <51813538+ahlag@users.noreply.github.com>
1 parent 5b3960a commit ecd340d

1 file changed

Lines changed: 34 additions & 12 deletions

File tree

‎arrow-array/src/array/union_array.rs‎

Lines changed: 34 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -794,18 +794,23 @@ unsafe impl Array for UnionArray {
794794
};
795795

796796
if fields.len() <= 1 {
797-
return self.fields.iter().find_map(|field_opt| {
798-
field_opt
799-
.as_ref()
800-
.and_then(|field| field.logical_nulls())
801-
.map(|logical_nulls| {
802-
if self.is_dense() {
803-
self.gather_nulls(vec![(0, logical_nulls)]).into()
804-
} else {
805-
logical_nulls
806-
}
807-
})
808-
});
797+
return self
798+
.fields
799+
.iter()
800+
.enumerate()
801+
.find_map(|(type_id, field_opt)| {
802+
field_opt
803+
.as_ref()
804+
.and_then(|field| field.logical_nulls())
805+
.map(|logical_nulls| {
806+
if self.is_dense() {
807+
self.gather_nulls(vec![(type_id as i8, logical_nulls)])
808+
.into()
809+
} else {
810+
logical_nulls
811+
}
812+
})
813+
});
809814
}
810815

811816
let logical_nulls = self.fields_logical_nulls();
@@ -2012,6 +2017,23 @@ mod tests {
20122017
assert_eq!(expected, array.gather_nulls(array.fields_logical_nulls()));
20132018
}
20142019

2020+
#[test]
2021+
fn test_dense_union_one_field_nonzero_type_id_logical_nulls() {
2022+
let fields = UnionFields::try_new([5], [Field::new("a", DataType::Int32, true)]).unwrap();
2023+
let child = Arc::new(Int32Array::from(vec![None, Some(1)]));
2024+
let array = UnionArray::try_new(
2025+
fields,
2026+
vec![5, 5, 5].into(),
2027+
Some(vec![1, 0, 1].into()),
2028+
vec![child],
2029+
)
2030+
.unwrap();
2031+
2032+
let expected = NullBuffer::from(vec![true, false, true]);
2033+
assert_eq!(array.logical_nulls(), Some(expected));
2034+
assert_eq!(array.logical_null_count(), 1);
2035+
}
2036+
20152037
#[test]
20162038
fn test_sparse_union_logical_nulls_mask_all_nulls_skip_one() {
20172039
let fields: UnionFields = [

0 commit comments

Comments
 (0)