comphead commented on code in PR #25754:
URL: https://github.com/apache/datafusion/pull/25754#discussion_r4285410058
##########
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() {
Review Comment:
Optional and not measured. `narrow_2` gains 16% where the 16-wide case with
the same NULL density gains 43%, so per-row bookkeeping looks like what
remains. When `list.null_count() == 0` there are no hidden elements, so
`list.is_valid(row)`, `hidden_count` and `keep_mask` do nothing and the loop is
one `set_bits.count_before(end)` per row with `keep` returned as is. Splitting
that case out would also keep the common path short. Worth a quick try on
`narrow_2` before deciding.
##########
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:
`array_compact` builds the same kind of offsets with the checked
`OffsetBuffer::new` and keeps `unsafe` for the UTF-8 byte copy. Here the
checked constructor is one linear scan over `list.len() + 1` offsets. Is the
unchecked version measurably faster on `int32/narrow_2/null_elements25`, where
per-row cost is highest? If it is within noise, I would use `OffsetBuffer::new`
and drop the `unsafe` block and its SAFETY comment.
##########
datafusion/functions-nested/src/utils.rs:
##########
@@ -529,6 +626,79 @@ pub(crate) mod tests {
use arrow::datatypes::Int64Type;
use datafusion_common::utils::SingleRowListArrayBuilder;
+ /// Compare `prepare_list_filter` with a per-row reference on a sliced list
+ /// with the given row lengths and validity, returning the input and output
+ /// bitmaps so callers can check whether the input was reused.
+ fn check_list_filter<O: OffsetSizeTrait>(
Review Comment:
`check_list_filter` is a good oracle but it runs four fixed shapes. The
logic is stateful across rows (word crossings, the first hidden row triggering
the lazy copy, `hidden_count` carried forward), so a seeded loop over random
row lengths (0 to about 130), validity, `keep` density and slice `(offset,
len)` against the same reference, for `i32` and `i64`, would reach orderings
these arrays cannot. `rand` is already a dev-dependency of this crate.
##########
datafusion/functions-nested/src/array_filter.rs:
##########
@@ -408,7 +382,103 @@ mod tests {
lit(true),
)
.unwrap();
- assert_eq!(res.as_list::<i32>(), &list);
+ assert_reuses_input(res.as_list::<i32>(), &list);
+ }
+
+ #[test]
+ fn filter_sliced_null_elements_and_rows() {
Review Comment:
Missing case: a predicate that selects only elements stored under a NULL
row, for example `[[1, 2], NULL (stores 99, 100), [3]]` with `v > 50`. `keep`
then has set bits but `kept_count` is 0, so `filter_list` returns through
`empty_list_values` after `hidden_count` was subtracted. Neither this test nor
`filter_does_not_copy_values_underlying_null_rows` leaves the result empty
while a NULL row holds passing values. A second case where every visible
element passes but a NULL row stores elements would pin that the early
`original_list` return is not taken (`kept_count < flattened_values.len()`).
##########
datafusion/functions-nested/src/lambda_utils.rs:
##########
@@ -152,12 +151,10 @@ impl EvaluatedListLambda {
self.original_list.nulls()
}
- pub(crate) fn row_range(&self, i: usize) -> (usize, usize) {
- (self.row_offsets[i], self.row_offsets[i + 1])
- }
-
- pub(crate) fn adjusted_offsets<O: OffsetSizeTrait>(&self) ->
OffsetBuffer<O> {
- OffsetBuffer::from_lengths(self.row_offsets.windows(2).map(|w| w[1] -
w[0]))
+ /// Row boundaries within `flattened_values`, including the elements that
+ /// NULL rows store. Functions that walk elements row by row compute this
once.
+ pub(crate) fn row_offsets(&self) -> Result<Vec<usize>> {
+ adjusted_row_offsets(&self.original_list)
Review Comment:
`array_first` and `array_any_match` still pay for the two copies this PR
removes from `array_filter`. `adjusted_row_offsets` calls
`adjust_offsets_for_slice`, which builds an adjusted `OffsetBuffer`, and then
`collect()` copies it again into a `Vec<usize>`. Mapping the original offsets
minus the first one straight into the `Vec<usize>` would leave one allocation.
--
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]
