wiedld commented on code in PR #6303:
URL: https://github.com/apache/arrow-rs/pull/6303#discussion_r1775622577
##########
arrow-array/src/array/mod.rs:
##########
@@ -283,6 +284,8 @@ pub trait Array: std::fmt::Debug + Send + Sync {
/// are present. For example a [`DictionaryArray`] with nullable values
will still return true,
/// even if the nulls present in [`DictionaryArray::values`] are not
referenced by any key,
/// and therefore would not appear in [`Array::logical_nulls`].
+ /// Similary, a [`UnionArray`] with any nullable child will always return
true, even if all
+ /// selected values are valid, and therefore would not appear in
[`Array::logical_nulls`].
fn is_nullable(&self) -> bool {
self.null_count() != 0
Review Comment:
Does the added code comment match the implementation?
UnionArrays always [returns a zero
null_count](https://github.com/apache/arrow-rs/blob/0b9a443a7aed67bed47692317496ddb51c726bf0/arrow-array/src/array/union_array.rs#L873-L876).
So this conditional `self.null_count() != 0` always evaluates to false, which
means is never nullable.
##########
arrow-array/src/array/union_array.rs:
##########
@@ -479,20 +739,139 @@ impl Array for UnionArray {
None
}
+ fn logical_nulls(&self) -> Option<NullBuffer> {
+ let fields = match self.data_type() {
+ DataType::Union(fields, _) => fields,
+ _ => unreachable!(),
+ };
+
+ if fields.len() <= 1 {
+ return self
+ .fields
+ .iter()
+ .flatten()
+ .map(Array::logical_nulls)
+ .next()
+ .flatten();
+ }
+
+ let logical_nulls = fields
+ .iter()
+ .filter_map(|(type_id, _)| Some((type_id,
self.child(type_id).logical_nulls()?)))
+ .filter(|(_, nulls)| nulls.null_count() > 0)
+ .collect::<Vec<_>>();
+
+ if logical_nulls.is_empty() {
+ return None;
+ }
+
+ let fully_null_count = logical_nulls
+ .iter()
+ .filter(|(_, nulls)| nulls.null_count() == nulls.len())
+ .count();
+
+ if fully_null_count == fields.len() {
+ if let Some((_, exactly_sized)) = logical_nulls
+ .iter()
+ .find(|(_, nulls)| nulls.len() == self.len())
+ {
+ return Some(exactly_sized.clone());
+ }
+
+ if let Some((_, bigger)) = logical_nulls
+ .iter()
+ .find(|(_, nulls)| nulls.len() > self.len())
+ {
+ return Some(bigger.slice(0, self.len()));
+ }
+
+ return Some(NullBuffer::new_null(self.len()));
+ }
+
+ let boolean_buffer = match &self.offsets {
+ Some(_) => self.gather_nulls(logical_nulls),
+ None => {
+ enum SparseStrategy {
+ Gather,
+ AllNullsSkipOne,
+ MixedSkipWithoutNulls,
+ MixedSkipFullyNull,
+ }
+
+ // This is the threshold where masking becomes slower than
gather
+ // TODO: test on avx512f(feature is still unstable)
+ let gather_relative_cost = if cfg!(target_feature = "avx2") {
+ 10
+ } else if cfg!(target_feature = "sse4.1") {
+ 3
+ } else {
+ 2
+ };
Review Comment:
Do you have any sources (or data) on specifically why these numbers?
There is a clear preference to using the gather approach when we have many
fields (partial null) to iterate thru, but I didn't follow how these exact
numbers were chosen.
##########
arrow-array/src/array/union_array.rs:
##########
@@ -380,6 +392,254 @@ impl UnionArray {
_ => unreachable!(),
}
}
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields without nulls
+ fn mask_sparse_skip_without_nulls(&self, nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ // Example logic for a union with 5 fields, a, b & c with nulls, d & e
without nulls:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c] = masks;
+ // let is_d_or_e = !(is_a | is_b)
+ // let union_chunk_nulls = is_d_or_e | (is_b & b_nulls) | (is_c &
c_nulls)
+ let fold = |(any_nullable_selected, union_nulls), (is_field,
field_nulls)| {
+ (
+ any_nullable_selected | is_field,
+ union_nulls | (is_field & field_nulls),
+ )
+ };
+
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let (any_nullable_selected, union_nulls) = nulls_masks_iters
+ .iter_mut()
+ .map(|(field_type_id, field_nulls)| {
+ let field_nulls = field_nulls.next().unwrap();
+ let is_field = selection_mask(chunk, *field_type_id);
+
+ (is_field, field_nulls)
+ })
+ .fold((0, 0), fold);
+
+ !any_nullable_selected | union_nulls
+ },
+ |remainder, bit_chunks| {
+ let (any_nullable_selected, union_nulls) = bit_chunks
+ .iter()
+ .map(|(field_type_id, field_bit_chunks)| {
+ let field_nulls = field_bit_chunks.remainder_bits();
+ let is_field = selection_mask(remainder,
*field_type_id);
+
+ (is_field, field_nulls)
+ })
+ .fold((0, 0), fold);
+
+ !any_nullable_selected | union_nulls
+ },
+ )
+ }
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields fully null
+ fn mask_sparse_skip_fully_null(&self, mut nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ let fields = match self.data_type() {
+ DataType::Union(fields, _) => fields,
+ _ => unreachable!("Union array's data type is not a union!"),
+ };
+
+ let type_ids = fields.iter().map(|(id, _)| id).collect::<HashSet<_>>();
+ let with_nulls = nulls.iter().map(|(id, _)|
*id).collect::<HashSet<_>>();
+
+ let without_nulls_ids = type_ids
+ .difference(&with_nulls)
+ .copied()
+ .collect::<Vec<_>>();
+
+ nulls.retain(|(_, nulls)| nulls.null_count() < nulls.len());
+
+ // Example logic for a union with 6 fields, a, b & c with nulls, d & e
without nulls, and f fully_null:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c, is_d, is_e] = masks;
+ // let union_chunk_nulls = is_d | is_e | (is_a & a_nulls) | (is_b &
b_nulls) | (is_c & c_nulls)
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let union_nulls = nulls_masks_iters.iter_mut().fold(
+ 0,
+ |union_nulls, (field_type_id, nulls_iter)| {
+ let field_nulls = nulls_iter.next().unwrap();
+
+ if field_nulls == 0 {
+ union_nulls
+ } else {
+ let is_field = selection_mask(chunk,
*field_type_id);
+
+ union_nulls | (is_field & field_nulls)
+ }
+ },
+ );
+
+ union_nulls | without_nulls_selected(chunk, &without_nulls_ids)
+ },
+ |remainder, bit_chunks| {
+ let union_nulls =
+ bit_chunks
+ .iter()
+ .fold(0, |union_nulls, (field_type_id,
field_bit_chunks)| {
+ let is_field = selection_mask(remainder,
*field_type_id);
+ let field_nulls =
field_bit_chunks.remainder_bits();
+
+ union_nulls | is_field & field_nulls
+ });
+
+ union_nulls | without_nulls_selected(remainder,
&without_nulls_ids)
+ },
+ )
+ }
+
+ /// Computes the logical nulls for a sparse union, optimized for when all
fields contains nulls
+ fn mask_sparse_all_with_nulls_skip_one(&self, nulls: Vec<(i8,
NullBuffer)>) -> BooleanBuffer {
+ // Example logic for a union with 3 fields, a, b & c, all containing
nulls:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // We can skip the first field: it's selection mask is the negation of
all others selection mask
+ // let [is_b, is_c] = selection_masks;
+ // let is_a = !(is_b | is_c)
+ // let union_chunk_nulls = (is_a & a_nulls) | (is_b & b_nulls) | (is_c
& c_nulls)
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let (is_not_first, union_nulls) = nulls_masks_iters[1..] //
skip first
+ .iter_mut()
+ .fold(
+ (0, 0),
+ |(is_not_first, union_nulls), (field_type_id,
nulls_iter)| {
+ let field_nulls = nulls_iter.next().unwrap();
+ let is_field = selection_mask(chunk,
*field_type_id);
+
+ (
+ is_not_first | is_field,
+ union_nulls | (is_field & field_nulls),
+ )
+ },
+ );
+
+ let is_first = !is_not_first;
+ let first_nulls = nulls_masks_iters[0].1.next().unwrap();
+
+ (is_first & first_nulls) | union_nulls
+ },
+ |remainder, bit_chunks| {
+ bit_chunks
+ .iter()
+ .fold(0, |union_nulls, (field_type_id, field_bit_chunks)| {
+ let field_nulls = field_bit_chunks.remainder_bits();
+ // The same logic as above, except that since this
runs at most once,
+ // it doesn't make difference to speed-up the first
selection mask
+ let is_field = selection_mask(remainder,
*field_type_id);
+
+ union_nulls | (is_field & field_nulls)
+ })
+ },
+ )
+ }
+
+ /// Maps `nulls` to `BitChunk's` and then to `BitChunkIterator's`, then
divides `self.type_ids` into exact chunks of 64 values,
+ /// calling `mask_chunk` for every exact chunk, and `mask_remainder` for
the remainder, if any, collecting the result in a `BooleanBuffer`
+ fn mask_sparse_helper(
+ &self,
+ nulls: Vec<(i8, NullBuffer)>,
+ mut mask_chunk: impl FnMut(&[i8; 64], &mut [(i8, BitChunkIterator)])
-> u64,
+ mask_remainder: impl FnOnce(&[i8], &[(i8, BitChunks)]) -> u64,
+ ) -> BooleanBuffer {
+ let bit_chunks = nulls
+ .iter()
+ .map(|(type_id, nulls)| (*type_id, nulls.inner().bit_chunks()))
+ .collect::<Vec<_>>();
+
+ let mut nulls_masks_iter = bit_chunks
+ .iter()
+ .map(|(type_id, bit_chunks)| (*type_id, bit_chunks.iter()))
+ .collect::<Vec<_>>();
+
+ let chunks_exact = self.type_ids.chunks_exact(64);
+ let remainder = chunks_exact.remainder();
+
+ let chunks = chunks_exact.map(|chunk| {
+ let chunk_array = <&[i8; 64]>::try_from(chunk).unwrap();
+
+ mask_chunk(chunk_array, &mut nulls_masks_iter)
+ });
Review Comment:
Nit: would you mind renaming the variable `chunk_array` (and where the
mask_chunk closures are defined) as `type_ids_chunk_array` or something similar?
##########
arrow-array/src/array/union_array.rs:
##########
@@ -380,6 +392,254 @@ impl UnionArray {
_ => unreachable!(),
}
}
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields without nulls
+ fn mask_sparse_skip_without_nulls(&self, nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ // Example logic for a union with 5 fields, a, b & c with nulls, d & e
without nulls:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c] = masks;
+ // let is_d_or_e = !(is_a | is_b)
+ // let union_chunk_nulls = is_d_or_e | (is_b & b_nulls) | (is_c &
c_nulls)
+ let fold = |(any_nullable_selected, union_nulls), (is_field,
field_nulls)| {
+ (
+ any_nullable_selected | is_field,
+ union_nulls | (is_field & field_nulls),
+ )
+ };
+
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let (any_nullable_selected, union_nulls) = nulls_masks_iters
+ .iter_mut()
+ .map(|(field_type_id, field_nulls)| {
+ let field_nulls = field_nulls.next().unwrap();
+ let is_field = selection_mask(chunk, *field_type_id);
+
+ (is_field, field_nulls)
+ })
+ .fold((0, 0), fold);
+
+ !any_nullable_selected | union_nulls
Review Comment:
Could we add a comment on this line?
It's basically the inverse of line 480 (see requested comment there). It's
clear once reading the NullBuffer docs, but it's an extra hop for anyone newer
to the codebase. Thank you!
##########
arrow-array/src/array/union_array.rs:
##########
@@ -380,6 +392,254 @@ impl UnionArray {
_ => unreachable!(),
}
}
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields without nulls
+ fn mask_sparse_skip_without_nulls(&self, nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ // Example logic for a union with 5 fields, a, b & c with nulls, d & e
without nulls:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c] = masks;
+ // let is_d_or_e = !(is_a | is_b)
+ // let union_chunk_nulls = is_d_or_e | (is_b & b_nulls) | (is_c &
c_nulls)
+ let fold = |(any_nullable_selected, union_nulls), (is_field,
field_nulls)| {
+ (
+ any_nullable_selected | is_field,
+ union_nulls | (is_field & field_nulls),
+ )
+ };
+
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let (any_nullable_selected, union_nulls) = nulls_masks_iters
+ .iter_mut()
+ .map(|(field_type_id, field_nulls)| {
+ let field_nulls = field_nulls.next().unwrap();
+ let is_field = selection_mask(chunk, *field_type_id);
+
+ (is_field, field_nulls)
+ })
+ .fold((0, 0), fold);
+
+ !any_nullable_selected | union_nulls
+ },
+ |remainder, bit_chunks| {
+ let (any_nullable_selected, union_nulls) = bit_chunks
+ .iter()
+ .map(|(field_type_id, field_bit_chunks)| {
+ let field_nulls = field_bit_chunks.remainder_bits();
+ let is_field = selection_mask(remainder,
*field_type_id);
+
+ (is_field, field_nulls)
+ })
+ .fold((0, 0), fold);
+
+ !any_nullable_selected | union_nulls
+ },
+ )
+ }
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields fully null
+ fn mask_sparse_skip_fully_null(&self, mut nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ let fields = match self.data_type() {
+ DataType::Union(fields, _) => fields,
+ _ => unreachable!("Union array's data type is not a union!"),
+ };
+
+ let type_ids = fields.iter().map(|(id, _)| id).collect::<HashSet<_>>();
+ let with_nulls = nulls.iter().map(|(id, _)|
*id).collect::<HashSet<_>>();
+
+ let without_nulls_ids = type_ids
+ .difference(&with_nulls)
+ .copied()
+ .collect::<Vec<_>>();
+
+ nulls.retain(|(_, nulls)| nulls.null_count() < nulls.len());
+
+ // Example logic for a union with 6 fields, a, b & c with nulls, d & e
without nulls, and f fully_null:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c, is_d, is_e] = masks;
+ // let union_chunk_nulls = is_d | is_e | (is_a & a_nulls) | (is_b &
b_nulls) | (is_c & c_nulls)
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let union_nulls = nulls_masks_iters.iter_mut().fold(
+ 0,
+ |union_nulls, (field_type_id, nulls_iter)| {
+ let field_nulls = nulls_iter.next().unwrap();
+
+ if field_nulls == 0 {
+ union_nulls
+ } else {
+ let is_field = selection_mask(chunk,
*field_type_id);
+
+ union_nulls | (is_field & field_nulls)
+ }
+ },
+ );
+
+ union_nulls | without_nulls_selected(chunk, &without_nulls_ids)
Review Comment:
Could we add a comment on this line?
I believe the without_nulls_selected returns the bitmask for the without
nulls (the `is_d | is_e`), which provides the proper true NullBuffer slot
(which translates to not-the-f-array-null, a.k.a. not one of the fully nulls).
##########
arrow-array/src/array/mod.rs:
##########
@@ -283,6 +284,8 @@ pub trait Array: std::fmt::Debug + Send + Sync {
/// are present. For example a [`DictionaryArray`] with nullable values
will still return true,
/// even if the nulls present in [`DictionaryArray::values`] are not
referenced by any key,
/// and therefore would not appear in [`Array::logical_nulls`].
+ /// Similary, a [`UnionArray`] with any nullable child will always return
true, even if all
+ /// selected values are valid, and therefore would not appear in
[`Array::logical_nulls`].
fn is_nullable(&self) -> bool {
self.null_count() != 0
Review Comment:
I think the proposed/added doc comments is what **_should_** happen -- but
it's not happening now.
Union types have no null validity bitmap ([per
spec](https://arrow.apache.org/docs/format/Columnar.html#union-layout)). I
believe this is why the null_count is always zero and the
[UnionArray::nulls](https://github.com/apache/arrow-rs/blob/0b9a443a7aed67bed47692317496ddb51c726bf0/arrow-array/src/array/union_array.rs#L739)
is None.
In contrast, I believe the is_nullable is based on the child types (as the
docs added here suggest). But what's missing is to define
`UnionArray::is_nullable` on it's implementation, and in a way that examines
the UnionArray's children. (Dictionary arrays [already does
so](https://github.com/apache/arrow-rs/blob/62825b27e98e6719cb66258535c75c7490ddba44/arrow-array/src/array/dictionary_array.rs#L753).)
Do you agree?
If you agree, then maybe add this doc comment & the code change in a follow
up PR?
##########
arrow-array/src/array/union_array.rs:
##########
@@ -380,6 +392,254 @@ impl UnionArray {
_ => unreachable!(),
}
}
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields without nulls
+ fn mask_sparse_skip_without_nulls(&self, nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ // Example logic for a union with 5 fields, a, b & c with nulls, d & e
without nulls:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c] = masks;
+ // let is_d_or_e = !(is_a | is_b)
+ // let union_chunk_nulls = is_d_or_e | (is_b & b_nulls) | (is_c &
c_nulls)
Review Comment:
```suggestion
// let is_d_or_e = !(is_a | is_b | is_c)
// let union_chunk_nulls = is_d_or_e | (is_a & a_nulls) | (is_b &
b_nulls) | (is_c & c_nulls)
```
##########
arrow-array/src/array/union_array.rs:
##########
@@ -479,20 +739,139 @@ impl Array for UnionArray {
None
}
+ fn logical_nulls(&self) -> Option<NullBuffer> {
+ let fields = match self.data_type() {
+ DataType::Union(fields, _) => fields,
+ _ => unreachable!(),
+ };
+
+ if fields.len() <= 1 {
+ return self
+ .fields
+ .iter()
+ .flatten()
+ .map(Array::logical_nulls)
+ .next()
+ .flatten();
+ }
+
+ let logical_nulls = fields
+ .iter()
+ .filter_map(|(type_id, _)| Some((type_id,
self.child(type_id).logical_nulls()?)))
+ .filter(|(_, nulls)| nulls.null_count() > 0)
+ .collect::<Vec<_>>();
+
+ if logical_nulls.is_empty() {
+ return None;
+ }
+
+ let fully_null_count = logical_nulls
+ .iter()
+ .filter(|(_, nulls)| nulls.null_count() == nulls.len())
+ .count();
+
+ if fully_null_count == fields.len() {
+ if let Some((_, exactly_sized)) = logical_nulls
+ .iter()
+ .find(|(_, nulls)| nulls.len() == self.len())
+ {
+ return Some(exactly_sized.clone());
+ }
+
+ if let Some((_, bigger)) = logical_nulls
+ .iter()
+ .find(|(_, nulls)| nulls.len() > self.len())
+ {
+ return Some(bigger.slice(0, self.len()));
+ }
+
+ return Some(NullBuffer::new_null(self.len()));
+ }
+
+ let boolean_buffer = match &self.offsets {
+ Some(_) => self.gather_nulls(logical_nulls),
+ None => {
+ enum SparseStrategy {
+ Gather,
+ AllNullsSkipOne,
+ MixedSkipWithoutNulls,
+ MixedSkipFullyNull,
Review Comment:
Could you move this enum elsewhere and doc comment what the options mean?
Since the names are not self explanatory.
For example, the `AllNullsSkipOne` is actually "all fields have some (not
all) nulls" and the `SkipOne` part refers to how the child arrays are iterated.
Having that enum explained up front helps make the rest of the code read much
faster. 🙏🏼
##########
arrow-array/src/array/union_array.rs:
##########
@@ -380,6 +392,254 @@ impl UnionArray {
_ => unreachable!(),
}
}
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields without nulls
+ fn mask_sparse_skip_without_nulls(&self, nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ // Example logic for a union with 5 fields, a, b & c with nulls, d & e
without nulls:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c] = masks;
+ // let is_d_or_e = !(is_a | is_b)
+ // let union_chunk_nulls = is_d_or_e | (is_b & b_nulls) | (is_c &
c_nulls)
+ let fold = |(any_nullable_selected, union_nulls), (is_field,
field_nulls)| {
+ (
+ any_nullable_selected | is_field,
+ union_nulls | (is_field & field_nulls),
+ )
+ };
+
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let (any_nullable_selected, union_nulls) = nulls_masks_iters
+ .iter_mut()
+ .map(|(field_type_id, field_nulls)| {
+ let field_nulls = field_nulls.next().unwrap();
+ let is_field = selection_mask(chunk, *field_type_id);
+
+ (is_field, field_nulls)
+ })
+ .fold((0, 0), fold);
+
+ !any_nullable_selected | union_nulls
+ },
+ |remainder, bit_chunks| {
+ let (any_nullable_selected, union_nulls) = bit_chunks
+ .iter()
+ .map(|(field_type_id, field_bit_chunks)| {
+ let field_nulls = field_bit_chunks.remainder_bits();
+ let is_field = selection_mask(remainder,
*field_type_id);
+
+ (is_field, field_nulls)
+ })
+ .fold((0, 0), fold);
+
+ !any_nullable_selected | union_nulls
+ },
+ )
+ }
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields fully null
+ fn mask_sparse_skip_fully_null(&self, mut nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
+ let fields = match self.data_type() {
+ DataType::Union(fields, _) => fields,
+ _ => unreachable!("Union array's data type is not a union!"),
+ };
+
+ let type_ids = fields.iter().map(|(id, _)| id).collect::<HashSet<_>>();
+ let with_nulls = nulls.iter().map(|(id, _)|
*id).collect::<HashSet<_>>();
+
+ let without_nulls_ids = type_ids
+ .difference(&with_nulls)
+ .copied()
+ .collect::<Vec<_>>();
+
+ nulls.retain(|(_, nulls)| nulls.null_count() < nulls.len());
+
+ // Example logic for a union with 6 fields, a, b & c with nulls, d & e
without nulls, and f fully_null:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // let [is_a, is_b, is_c, is_d, is_e] = masks;
+ // let union_chunk_nulls = is_d | is_e | (is_a & a_nulls) | (is_b &
b_nulls) | (is_c & c_nulls)
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let union_nulls = nulls_masks_iters.iter_mut().fold(
+ 0,
+ |union_nulls, (field_type_id, nulls_iter)| {
+ let field_nulls = nulls_iter.next().unwrap();
+
+ if field_nulls == 0 {
+ union_nulls
+ } else {
+ let is_field = selection_mask(chunk,
*field_type_id);
+
+ union_nulls | (is_field & field_nulls)
+ }
+ },
+ );
+
+ union_nulls | without_nulls_selected(chunk, &without_nulls_ids)
+ },
+ |remainder, bit_chunks| {
+ let union_nulls =
+ bit_chunks
+ .iter()
+ .fold(0, |union_nulls, (field_type_id,
field_bit_chunks)| {
+ let is_field = selection_mask(remainder,
*field_type_id);
+ let field_nulls =
field_bit_chunks.remainder_bits();
+
+ union_nulls | is_field & field_nulls
+ });
+
+ union_nulls | without_nulls_selected(remainder,
&without_nulls_ids)
+ },
+ )
+ }
+
+ /// Computes the logical nulls for a sparse union, optimized for when all
fields contains nulls
+ fn mask_sparse_all_with_nulls_skip_one(&self, nulls: Vec<(i8,
NullBuffer)>) -> BooleanBuffer {
+ // Example logic for a union with 3 fields, a, b & c, all containing
nulls:
+ // let [a_nulls, b_nulls, c_nulls] = nulls;
+ // We can skip the first field: it's selection mask is the negation of
all others selection mask
+ // let [is_b, is_c] = selection_masks;
+ // let is_a = !(is_b | is_c)
+ // let union_chunk_nulls = (is_a & a_nulls) | (is_b & b_nulls) | (is_c
& c_nulls)
+ self.mask_sparse_helper(
+ nulls,
+ |chunk, nulls_masks_iters| {
+ let (is_not_first, union_nulls) = nulls_masks_iters[1..] //
skip first
+ .iter_mut()
+ .fold(
+ (0, 0),
+ |(is_not_first, union_nulls), (field_type_id,
nulls_iter)| {
+ let field_nulls = nulls_iter.next().unwrap();
+ let is_field = selection_mask(chunk,
*field_type_id);
+
+ (
+ is_not_first | is_field,
+ union_nulls | (is_field & field_nulls),
+ )
+ },
+ );
+
+ let is_first = !is_not_first;
+ let first_nulls = nulls_masks_iters[0].1.next().unwrap();
+
+ (is_first & first_nulls) | union_nulls
+ },
+ |remainder, bit_chunks| {
+ bit_chunks
+ .iter()
+ .fold(0, |union_nulls, (field_type_id, field_bit_chunks)| {
+ let field_nulls = field_bit_chunks.remainder_bits();
+ // The same logic as above, except that since this
runs at most once,
+ // it doesn't make difference to speed-up the first
selection mask
+ let is_field = selection_mask(remainder,
*field_type_id);
+
+ union_nulls | (is_field & field_nulls)
+ })
+ },
+ )
+ }
+
+ /// Maps `nulls` to `BitChunk's` and then to `BitChunkIterator's`, then
divides `self.type_ids` into exact chunks of 64 values,
+ /// calling `mask_chunk` for every exact chunk, and `mask_remainder` for
the remainder, if any, collecting the result in a `BooleanBuffer`
+ fn mask_sparse_helper(
+ &self,
+ nulls: Vec<(i8, NullBuffer)>,
+ mut mask_chunk: impl FnMut(&[i8; 64], &mut [(i8, BitChunkIterator)])
-> u64,
+ mask_remainder: impl FnOnce(&[i8], &[(i8, BitChunks)]) -> u64,
+ ) -> BooleanBuffer {
+ let bit_chunks = nulls
+ .iter()
+ .map(|(type_id, nulls)| (*type_id, nulls.inner().bit_chunks()))
+ .collect::<Vec<_>>();
+
+ let mut nulls_masks_iter = bit_chunks
+ .iter()
+ .map(|(type_id, bit_chunks)| (*type_id, bit_chunks.iter()))
+ .collect::<Vec<_>>();
+
+ let chunks_exact = self.type_ids.chunks_exact(64);
+ let remainder = chunks_exact.remainder();
+
+ let chunks = chunks_exact.map(|chunk| {
+ let chunk_array = <&[i8; 64]>::try_from(chunk).unwrap();
+
+ mask_chunk(chunk_array, &mut nulls_masks_iter)
+ });
+
+ // SAFETY:
+ // chunks is a ChunksExact iterator, which implements TrustedLen, and
correctly reports its length
+ let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks)
};
+
+ if !remainder.is_empty() {
+ buffer.push(mask_remainder(remainder, &bit_chunks));
+ }
+
+ BooleanBuffer::new(buffer.into(), 0, self.type_ids.len())
+ }
+
+ /// Computes the logical nulls for a sparse or dense union, by gathering
individual bits from the null buffer of the selected field
+ fn gather_nulls(&self, nulls: Vec<(i8, NullBuffer)>) -> BooleanBuffer {
+ let one_null = NullBuffer::new_null(1);
+ let one_valid = NullBuffer::new_valid(1);
+
+ // Unsafe code below depend on it:
+ // To remove one branch from the loop, if the a type_id is not
utilized, or it's logical_nulls is None/all set,
+ // we use a null buffer of len 1 and a index_mask of 0, or the true
null buffer and usize::MAX otherwise.
+ // We then unconditionally access the null buffer with index &
index_mask,
+ // which always return 0 for the 1-len buffer, or the true index
unchanged otherwise
+ // We also use a 256 array, so llvm knows that `type_id as u8 as
usize` is always in bounds
+ let mut logical_nulls_array = [(&one_valid, Mask::Zero); 256];
Review Comment:
❤️
##########
arrow-array/src/array/union_array.rs:
##########
@@ -479,20 +675,91 @@ impl Array for UnionArray {
None
}
+ fn logical_nulls(&self) -> Option<NullBuffer> {
+ let DataType::Union(fields, _) = &self.data_type else {
+ unreachable!()
+ };
+
+ if fields.len() <= 1 {
+ return self
+ .fields
+ .iter()
+ .flatten()
+ .map(Array::logical_nulls)
+ .next()
+ .flatten();
+ }
+
+ let with_nulls = fields
+ .iter()
+ .filter_map(|(type_id, _)| Some((type_id,
self.child(type_id).logical_nulls()?)))
+ .filter(|(_, nulls)| nulls.null_count() > 0)
+ .collect::<Vec<_>>();
+
+ if with_nulls.is_empty() {
+ return None;
+ }
+
+ if with_nulls
+ .iter()
+ .all(|(_, nulls)| nulls.null_count() == nulls.len())
+ {
+ if let Some((_, exactly_sized)) = with_nulls
+ .iter()
+ .find(|(_, nulls)| nulls.len() == self.len())
+ {
+ return Some(exactly_sized.clone());
+ }
+
+ if let Some((_, bigger)) = with_nulls
+ .iter()
+ .find(|(_, nulls)| nulls.len() > self.len())
+ {
+ return Some(bigger.slice(0, self.len()));
+ }
+
+ return Some(NullBuffer::new_null(self.len()));
+ }
+
+ let buffer = match &self.offsets {
+ Some(_) => self.gather_nulls(&with_nulls),
+ None => {
+ // masking is only efficient up to a certain point
+ // when all fields contains nulls, one field mask can be
cheaply calculated
+ // by negating the others, so it's threshold is higher
+ if with_nulls.len() <= 10 && with_nulls.len() == fields.len() {
+ self.mask_sparse_all_nulls(&with_nulls)
+ } else if with_nulls.len() <= 9 {
+ self.mask_sparse_mixed_nulls(&with_nulls)
+ } else {
+ self.gather_nulls(&with_nulls)
+ }
+ }
+ };
+
+ let null_buffer = NullBuffer::from(buffer);
+
+ if null_buffer.null_count() > 0 {
+ Some(null_buffer)
+ } else {
+ None
+ }
+ }
+
/// Union types always return non null as there is no validity buffer.
Review Comment:
If we remove these specific implementations, the behavior will not change
but the documented reason (because UnionArray has no validity buffer) will be
lost.
If we remove these -- can we update the documentation for the default
implementations to mention the UnionArray special case?
##########
arrow-array/src/array/union_array.rs:
##########
@@ -380,6 +392,254 @@ impl UnionArray {
_ => unreachable!(),
}
}
+
+ /// Computes the logical nulls for a sparse union, optimized for when
there's a lot of fields without nulls
+ fn mask_sparse_skip_without_nulls(&self, nulls: Vec<(i8, NullBuffer)>) ->
BooleanBuffer {
Review Comment:
This function is not hit by any of the test cases, but does get used by the
benchmark.
There are also several other branch points in these mask methods, also only
used in the benchmark. I don't believe the benchmark tests for correctness (is
`black_box(array.logical_nulls()`). What is our policy here @alamb ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
