Jefffrey commented on code in PR #10236:
URL: https://github.com/apache/arrow-rs/pull/10236#discussion_r3610026253
##########
arrow-select/src/filter.rs:
##########
@@ -687,6 +717,21 @@ pub(crate) fn filter_null_mask(
Some((null_count, nulls))
}
+/// Appends the packed bits of the selected `[start, end)` runs of `src` (a bit
+/// buffer starting at `offset`) to `builder`. Shared by the range-based
+/// [`filter_bits`] strategies and [`filter_validity`].
+#[inline]
+fn append_validity_runs(
Review Comment:
i feel this way of reducing duplication by pulling these tiny code snippets
out into separate functions actually hurts readability; i would prefer if such
functions like this were left inlined
##########
arrow-select/src/filter.rs:
##########
@@ -1064,6 +1131,354 @@ fn filter_sparse_union(
})
}
+/// Whether the filter retains most of `len` rows (a "dense" filter)
+fn filter_keeps_most(predicate: &FilterPredicate, len: usize) -> bool {
+ predicate.count as f64 > len as f64 * FILTER_SLICES_SELECTIVITY_THRESHOLD
+}
Review Comment:
```suggestion
fn filter_keeps_most(predicate: &FilterPredicate, len: usize) -> bool {
predicate.count as f64 / len as f64 >
FILTER_SLICES_SELECTIVITY_THRESHOLD
}
```
i find this a bit more readable, especially as its how we do it here too:
https://github.com/apache/arrow-rs/blob/cc83f7eb25fbf0f6e8f94656d8fd703470e32c04/arrow-select/src/filter.rs#L355-L362
##########
arrow-select/src/filter.rs:
##########
@@ -739,6 +785,25 @@ fn filter_boolean(array: &BooleanArray, predicate:
&FilterPredicate) -> BooleanA
BooleanArray::new(values, nulls)
}
+/// Copies the selected `[start, end)` runs of `values` into a fresh `Vec`.
+///
+/// Assumes `runs` `[start, end)` are within `values`.
+#[inline]
+fn copy_runs<T: ArrowNativeType>(
Review Comment:
same here
##########
arrow-select/src/filter.rs:
##########
@@ -905,17 +963,26 @@ fn filter_bytes<T>(array: &GenericByteArray<T>,
predicate: &FilterPredicate) ->
where
T: ByteArrayType,
{
- let mut filter = FilterBytes::new(predicate.count, array);
-
+ // Range strategies reuse the shared range-driven builder (also used by
+ // list-child byte filtering), keeping the offset-rebuild + value-copy in
one
+ // place and streaming the ranges (no `Vec` materialization).
match &predicate.strategy {
IterationStrategy::SlicesIterator => {
-
filter.extend_offsets_slices(SlicesIterator::new(&predicate.filter),
predicate.count);
- filter.extend_slices(SlicesIterator::new(&predicate.filter))
+ return filter_bytes_runs(
Review Comment:
here too, i feel we should leave `filter_bytes()` unchanged for now as this
refactor adds complexity imo
##########
arrow-select/src/filter.rs:
##########
@@ -875,12 +923,22 @@ where
fn extend_offsets_slices(&mut self, iter: impl Iterator<Item = (usize,
usize)>, count: usize) {
self.dst_offsets.reserve_exact(count);
for (start, end) in iter {
- // These can only fail if `array` contains invalid data
- for idx in start..end {
- let (_, _, len) = self.get_value_range(idx);
- self.cur_offset += len;
- self.dst_offsets.push(self.cur_offset);
- }
+ // A retained run is contiguous in the source, so its source
offsets are
+ // monotonic and the run's new offsets are just that slice shifted
by a
+ // constant (`base - src_base`). Emitting them as a `map` over the
+ // contiguous source slice avoids the loop-carried `cur_offset +=
len`
+ // dependency of a per-element accumulation, letting the compiler
+ // vectorize the rebuild (matters at high selectivity, where the
run
+ // covers most elements).
+ let base = self.cur_offset.as_usize();
+ let src_base = self.src_offsets[start].as_usize();
+ self.dst_offsets.extend(
Review Comment:
this is a nice fix
##########
arrow-select/src/filter.rs:
##########
@@ -1064,6 +1131,354 @@ fn filter_sparse_union(
})
}
+/// Whether the filter retains most of `len` rows (a "dense" filter)
+fn filter_keeps_most(predicate: &FilterPredicate, len: usize) -> bool {
+ predicate.count as f64 > len as f64 * FILTER_SLICES_SELECTIVITY_THRESHOLD
+}
+
+/// An iterator of `[start, end)` ranges. Used for both retained parent-row
ranges and
+/// the child element ranges they map to.
+type RangeIter<'a> = Box<dyn Iterator<Item = (usize, usize)> + 'a>;
+
+/// Iterates the retained rows of `predicate` as `[start, end)` ranges
according
+/// to its strategy.
+fn predicate_row_ranges(predicate: &FilterPredicate) -> RangeIter<'_> {
+ match &predicate.strategy {
+ IterationStrategy::Slices(slices) => Box::new(slices.iter().copied()),
+ IterationStrategy::SlicesIterator =>
Box::new(SlicesIterator::new(&predicate.filter)),
+ IterationStrategy::IndexIterator => {
+ Box::new(IndexIterator::new(&predicate.filter,
predicate.count).map(|i| (i, i + 1)))
+ }
+ IterationStrategy::Indices(indices) =>
Box::new(indices.iter().map(|&i| (i, i + 1))),
+ IterationStrategy::All | IterationStrategy::None => unreachable!(),
+ }
+}
+
+/// `filter` for `List`/`LargeList`: returns a new list array containing only
the
+/// rows selected by `predicate`.
+///
+/// # How it works
+///
+/// A `List`/`LargeList` stores child elements contiguously. Because the
offsets are
+/// monotonically non-decreasing, a *run* of consecutive retained rows
+/// `[start, end)` maps to a single contiguous block of child elements
+/// `[offsets[start], offsets[end])`. Filtering a list is therefore two steps:
+///
+/// 1. Rebuild the offset buffer from the retained rows' lengths.
+/// 2. Filter the child `values` down to those contiguous element ranges.
+///
+/// ```text
+/// input: [ [a,b,c] [d,e] [f,g,h] [i,j,k] ] keep rows 0, 2, 3
+/// offsets: 0 3 5 8 11
+/// values: [ a b c d e f g h i j k ]
+/// └──────┘ (drop) └──────────────────┘
+///
+/// retained row runs -> [0, 1) [2, 4)
+/// child element blocks = [offsets[0], [offsets[2],
+/// offsets[1]) offsets[4])
+/// = [0, 3) [5, 11) rows 2 & 3 are
+/// a b c f g h i j k adjacent -> one
copy
+///
+/// output: [ [a,b,c] [f,g,h] [i,j,k] ] new offsets: 0, 3, 6, 9
+/// ```
+///
+fn filter_list<OffsetType: OffsetSizeTrait>(
+ array: &GenericListArray<OffsetType>,
+ predicate: &FilterPredicate,
+) -> Result<GenericListArray<OffsetType>, ArrowError> {
+ let make_row_ranges = || predicate_row_ranges(predicate);
+ filter_list_inner(array, &make_row_ranges, predicate.count)
+}
+
+/// Core of [`filter_list`], driven by a function producing the retained *row*
+/// ranges (`[start, end)`) of `array` rather than a [`FilterPredicate`]. This
+/// lets nested lists recurse with streamed child ranges. The range source is
+/// type-erased (`&dyn Fn` returning a boxed iterator) so the recursion is a
+/// single instantiation per offset type, not an unbounded tower of closure
+/// types.
+fn filter_list_inner<'a, OffsetType: OffsetSizeTrait>(
+ array: &'a GenericListArray<OffsetType>,
+ make_row_ranges: &dyn Fn() -> RangeIter<'a>,
+ count: usize,
+) -> Result<GenericListArray<OffsetType>, ArrowError> {
+ let (offsets, filtered_values) =
+ filter_list_offsets_and_child(array.offsets(), array.values(),
make_row_ranges, count)?;
+
+ let field = match array.data_type() {
+ DataType::List(field) | DataType::LargeList(field) => field.clone(),
+ _ => unreachable!(),
+ };
+ let nulls = filter_validity(array.nulls(), make_row_ranges(), count);
+
+ Ok(GenericListArray::new(
+ field,
+ offsets,
+ filtered_values,
+ nulls,
+ ))
+}
+
+/// Rebuilds the offset buffer for the retained rows and filters the child
`values` by the corresponding contiguous element ranges.
+///
+/// Returns the new, filtered [`OffsetBuffer`] and [`ArrayRef`].
+fn filter_list_offsets_and_child<'a, O: OffsetSizeTrait>(
+ offsets: &'a OffsetBuffer<O>,
+ values: &'a ArrayRef,
+ make_row_ranges: &dyn Fn() -> RangeIter<'a>,
+ count: usize,
+) -> Result<(OffsetBuffer<O>, ArrayRef), ArrowError> {
+ // Rebuild offsets from the retained rows' lengths, reading consecutive
+ // offsets sequentially (mirrors arrow-data's `try_extend_offsets`).
+ let mut new_offsets = Vec::with_capacity(count + 1);
+ new_offsets.push(O::usize_as(0));
+ let mut acc = 0usize; // running length of selected child elements
+ for (start, end) in make_row_ranges() {
+ // Offsets within a retained run are contiguous and monotonic, so the
new
+ // offsets are the source slice shifted by a constant (`acc -
src_base`).
+ // Mapping over the contiguous slice avoids the loop-carried `acc +=`
+ // dependency of the per-window accumulation and vectorizes.
+ let src_base = offsets[start].as_usize();
+ new_offsets.extend(
+ offsets[start + 1..=end]
+ .iter()
+ .map(|&o| O::usize_as(acc + (o.as_usize() - src_base))),
+ );
+ acc += offsets[end].as_usize() - src_base;
+ }
+
+ // Each retained row run maps to one contiguous block of child elements.
+ let make_child_ranges = || -> RangeIter<'a> {
+ Box::new(
+ make_row_ranges()
+ .map(|(start, end)| (offsets[start].as_usize(),
offsets[end].as_usize())),
+ )
+ };
+ let filtered = filter_list_child(values, &make_child_ranges, acc)?;
+
+ Ok((OffsetBuffer::new(ScalarBuffer::from(new_offsets)), filtered))
+}
+
+/// Specialised method to filter a [`FixedSizeListArray`] by a
[`FilterPredicate`].
+///
+/// Each list occupies exactly `size` contiguous child elements (no offsets),
so
+/// a run of retained rows `[start, end)` maps to one contiguous child block
+/// `[start * size, end * size)`. The child is filtered with those blocks via
the
+/// shared [`filter_list_child`] machinery, reusing the specialized per-type
+/// kernels instead of [`MutableArrayData`].
+fn filter_fixed_size_list(
+ array: &FixedSizeListArray,
+ predicate: &FilterPredicate,
+) -> Result<FixedSizeListArray, ArrowError> {
+ let size = array.value_length() as usize;
+ let child_count = predicate.count * size;
+
+ let child_ranges = || -> RangeIter<'_> {
+ Box::new(
+ predicate_row_ranges(predicate).map(move |(start, end)| (start *
size, end * size)),
+ )
+ };
+ let filtered_values = filter_list_child(array.values(), &child_ranges,
child_count)?;
+
+ let field = match array.data_type() {
+ DataType::FixedSizeList(field, _) => field.clone(),
+ _ => unreachable!(),
+ };
+ let nulls = predicate.filter_nulls(array.nulls());
+ Ok(FixedSizeListArray::new(
+ field,
+ size as i32,
+ filtered_values,
+ nulls,
+ ))
+}
+
+/// A map is structurally a `List<Struct<key, value>>`: i32 offsets over a
struct
+/// `entries` child. Filter the `entries` struct with the offsets from the
retained rows (via
+/// shared [`filter_list_child`]).
+fn filter_map(array: &MapArray, predicate: &FilterPredicate) ->
Result<MapArray, ArrowError> {
+ let make_row_ranges = || predicate_row_ranges(predicate);
+ let entries: ArrayRef = Arc::new(array.entries().clone());
+ let (offsets, filtered_entries) = filter_list_offsets_and_child(
+ array.offsets(),
+ &entries,
+ &make_row_ranges,
+ predicate.count,
+ )?;
+
+ let (field, ordered) = match array.data_type() {
+ DataType::Map(field, ordered) => (field.clone(), *ordered),
+ _ => unreachable!(),
+ };
+ let nulls = predicate.filter_nulls(array.nulls());
+ Ok(MapArray::new(
+ field,
+ offsets,
+ filtered_entries.as_struct().clone(),
+ nulls,
+ ordered,
+ ))
+}
+
+/// Filters the `values` child of a list-like array (`List`/`LargeList`/
+/// `FixedSizeList`) given a function producing the retained child-element
+/// ranges. Byte children stream into [`FilterBytes`]; other children reuse the
+/// specialized [`filter_array`] kernels via a `Slices` predicate (which copies
+/// directly from the ranges and never reads `FilterPredicate::filter`, so an
+/// empty filter suffices).
+#[allow(
Review Comment:
we should prefer `expect` to `allow`; that said, when i tried locally it
seems this isnt actually needed anymore?
--
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]
