neilconway commented on code in PR #25754:
URL: https://github.com/apache/datafusion/pull/25754#discussion_r4289989273
##########
datafusion/functions-nested/src/utils.rs:
##########
@@ -50,6 +50,103 @@ pub(crate) fn empty_list_values<O: OffsetSizeTrait>(
))
}
+/// Build result row offsets and exclude elements stored under NULL list rows.
+///
+/// `keep` is a bitmap identifying which elements in the list's visible span
+/// should be preserved. The returned bitmap uses the same indexing, with bits
+/// for NULL list rows cleared. Selected elements may themselves be NULL; their
+/// validity is the caller's responsibility when copying values.
+pub(crate) fn prepare_list_filter<O: OffsetSizeTrait>(
+ list: &GenericListArray<O>,
+ keep: &BooleanBuffer,
+) -> (OffsetBuffer<O>, BooleanBuffer) {
+ let list_offsets = list.offsets();
+ let (first_offset, span_len) = offset_span(list_offsets);
+ debug_assert_eq!(keep.len(), span_len);
+ let mut offsets = Vec::with_capacity(list.len() + 1);
+ offsets.push(O::zero());
+ let mut kept_count = 0;
+ // Count the set bits before each row end in one pass over `keep`, which
+ // stays cheap for short rows.
+ let mut set_bits = SetBitCounter::new(keep.bit_chunks().iter_padded());
+ // Set bits stored under NULL rows, which the result does not keep.
+ let mut hidden_count = 0;
+ // NULL list rows are usually empty, so reuse `keep` until a NULL row
stores
+ // elements. Copy the bits of the preceding rows when that first happens.
+ let mut keep_mask: Option<BooleanBufferBuilder> = None;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ let len = end - start;
+ let bit_start = keep.offset() + start;
+ if list.is_valid(row) {
+ kept_count = set_bits.count_before(end) - hidden_count;
+ if let Some(keep_mask) = &mut keep_mask {
+ keep_mask.append_packed_range(bit_start..bit_start + len,
keep.values());
+ }
+ } else if len != 0 {
+ let set_before_row = set_bits.count_before(start);
+ hidden_count += set_bits.count_before(end) - set_before_row;
+ keep_mask
+ .get_or_insert_with(|| {
+ let mut builder = BooleanBufferBuilder::new(keep.len());
+ builder.append_packed_range(keep.offset()..bit_start,
keep.values());
+ builder
+ })
+ .append_n(len, false);
+ }
+ offsets.push(O::usize_as(kept_count));
+ }
+ let keep_mask = keep_mask
+ .map(BooleanBufferBuilder::build)
+ .unwrap_or_else(|| keep.clone());
+ debug_assert_eq!(keep_mask.len(), keep.len());
+ // SAFETY: offsets start at zero and `kept_count` never decreases. The
count
+ // cannot exceed the input's element span, which fits in `O`, so converting
+ // it to `O` cannot overflow.
+ let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
Review Comment:
I measured and found the checked constructor was about 8% slower, so I kept
`new_unchecked`.
##########
datafusion/functions-nested/src/utils.rs:
##########
@@ -50,6 +50,103 @@ pub(crate) fn empty_list_values<O: OffsetSizeTrait>(
))
}
+/// Build result row offsets and exclude elements stored under NULL list rows.
+///
+/// `keep` is a bitmap identifying which elements in the list's visible span
+/// should be preserved. The returned bitmap uses the same indexing, with bits
+/// for NULL list rows cleared. Selected elements may themselves be NULL; their
+/// validity is the caller's responsibility when copying values.
+pub(crate) fn prepare_list_filter<O: OffsetSizeTrait>(
+ list: &GenericListArray<O>,
+ keep: &BooleanBuffer,
+) -> (OffsetBuffer<O>, BooleanBuffer) {
+ let list_offsets = list.offsets();
+ let (first_offset, span_len) = offset_span(list_offsets);
+ debug_assert_eq!(keep.len(), span_len);
+ let mut offsets = Vec::with_capacity(list.len() + 1);
+ offsets.push(O::zero());
+ let mut kept_count = 0;
+ // Count the set bits before each row end in one pass over `keep`, which
+ // stays cheap for short rows.
+ let mut set_bits = SetBitCounter::new(keep.bit_chunks().iter_padded());
+ // Set bits stored under NULL rows, which the result does not keep.
+ let mut hidden_count = 0;
+ // NULL list rows are usually empty, so reuse `keep` until a NULL row
stores
+ // elements. Copy the bits of the preceding rows when that first happens.
+ let mut keep_mask: Option<BooleanBufferBuilder> = None;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ let len = end - start;
+ let bit_start = keep.offset() + start;
+ if list.is_valid(row) {
+ kept_count = set_bits.count_before(end) - hidden_count;
+ if let Some(keep_mask) = &mut keep_mask {
+ keep_mask.append_packed_range(bit_start..bit_start + len,
keep.values());
+ }
+ } else if len != 0 {
+ let set_before_row = set_bits.count_before(start);
+ hidden_count += set_bits.count_before(end) - set_before_row;
+ keep_mask
+ .get_or_insert_with(|| {
+ let mut builder = BooleanBufferBuilder::new(keep.len());
+ builder.append_packed_range(keep.offset()..bit_start,
keep.values());
+ builder
+ })
+ .append_n(len, false);
+ }
+ offsets.push(O::usize_as(kept_count));
+ }
+ let keep_mask = keep_mask
+ .map(BooleanBufferBuilder::build)
+ .unwrap_or_else(|| keep.clone());
+ debug_assert_eq!(keep_mask.len(), keep.len());
+ // SAFETY: offsets start at zero and `kept_count` never decreases. The
count
+ // cannot exceed the input's element span, which fits in `O`, so converting
+ // it to `O` cannot overflow.
+ let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
Review Comment:
I measured and found the checked constructor was about 8% slower (M4 Max),
so I kept `new_unchecked`.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
