feat(parquet): RowSelection can be backed by a BooleanBuffer#10141
feat(parquet): RowSelection can be backed by a BooleanBuffer#10141haohuaijin wants to merge 31 commits into
RowSelection can be backed by a BooleanBuffer#10141Conversation
|
Thanks @alamb for the pointers. I looked at #6624, #7454, and the work that landed in #8733. My understanding is:
This does not replace selector-backed selections. Selectors are still better for clustered/page-index-style selections. The case this helps is when the caller already has a row-level bitmap, such as from an external index, FTS index, or bitmap index. Today that bitmap has to go through So I see this as a small representation-layer improvement that complements #8733. #8733 added mask execution; this PR adds a direct bitmap-backed input path for it. It does not try to implement the broader adaptive predicate-pushdown/page-cache design from #7454. |
|
run benchmark arrow_reader arrow_reader_row_filter |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing export-mask (1720537) to 2e035fd (merge-base) diff File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing export-mask (1720537) to 2e035fd (merge-base) diff File an issue against this benchmark runner |
|
run benchmark arrow_reader_clickbench |
|
Merging up to fix CI |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing export-mask (7b102f4) to 6ba533d (merge-base) diff File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark row_selector row_selection_cursor |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you @haohuaijin . I think this is totally the right direction and I really like where this is going
As long as this doesn't cause a performance regression I think it is good to go.
I have actually tried this myself (see #6624) as did @XiangpengHao in #6624.
Some other thoughts
- We need to make sure it doesn't cause a performance regression (I can help here)
- It currently has a public API change, so we can't merge it until main opens for the next major release (likely in mid-late July). If we can figure out how to keep the API the same we could potentially release it sooner
I haven' tmade it through the entire PR yet, but I did fire off some benchmarks
|
|
||
| #[inline] | ||
| fn next(&mut self) -> Option<RowSelector> { | ||
| match self { |
There was a problem hiding this comment.
In general I think this will add a branch in the main loops - every chunk of bits will require matching -- it might make sense to figure out how to have iterators for the two different types, so that code paths that want to iterate over the selection can match outside
match iter {
RowSelectionIter::Selectors(iter) => { for i in iter ... }, // special loop for selector backed
RowSelectionIter::Mask(iter) => { for i in iter ... }, // special loop for mask backed
}| if self.finished { | ||
| return None; | ||
| } | ||
| match self.slices.next() { |
There was a problem hiding this comment.
Much of this code appears untested
I ran code coverage like
cargo llvm-cov --html test -p parquet --features arrow
| impl Iterator for MaskRunIter<'_> { | ||
| type Item = RowSelector; | ||
|
|
||
| fn next(&mut self) -> Option<RowSelector> { |
There was a problem hiding this comment.
I think this is the main API that is not compatible -- it need to return &RowSelector -- if we could change that we can merge this PR before the next major release (with breaking API changes)
There was a problem hiding this comment.
I updated RowSelection::iter() to keep returning Iterator<Item = &RowSelector>.
- For selector-backed selections this still borrows from the internal
Vec<RowSelector>. - For mask-backed selections I added a lazy selector cache, because the returned
&RowSelectorvalues need stable storage insideRowSelection; otherwise they would point into a temporaryVec.
The internal hot paths still avoid this cache and use the BooleanBuffer or MaskRunIter directly.
the struct is below
pub struct RowSelection {
inner: RowSelectionInner,
}
pub(crate) enum RowSelectionInner {
Selectors(Vec<RowSelector>),
Mask(Box<MaskSelection>),
}
pub(crate) struct MaskSelection {
mask: BooleanBuffer,
selectors: OnceLock<Vec<RowSelector>>,
}I am not fully sure this is the best approach, so suggestions and feedback are welcome.
| if let (Some(l), Some(r)) = (self.as_mask(), other.as_mask()) { | ||
| return Self::from_boolean_buffer(intersect_masks(l, r)); | ||
| } | ||
| let l = self.materialize_for_combine(); |
There was a problem hiding this comment.
this is a clever idea -- as a follow on we could potentially implement specialized versions of each of these implementations
There was a problem hiding this comment.
for example, I think and_then can use some of the fancy BMI instructions that @devanbenz is investigating in #10136
| // Verify that the size of RowGroupDecoderState does not grow too large | ||
| fn test_structure_size() { | ||
| assert_eq!(std::mem::size_of::<RowGroupDecoderState>(), 232); | ||
| assert_eq!(std::mem::size_of::<RowGroupDecoderState>(), 256); |
There was a problem hiding this comment.
I think this is bigger because a RowSelection is now bigger.
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing export-mask (7b102f4) to 6ba533d (merge-base) diff File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing export-mask (7b102f4) to 6ba533d (merge-base) diff File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤔 looks like arrow_reader_row_filter is slightly slower. Rerunning to try an reproduce |
|
run benchmark arrow_reader_row_filter |
|
I reran one representative point from For
All five pairs moved in the same direction and Criterion produced no warnings. However, I want to be precise about the interpretation: this benchmark forces Could we rerun the complete L4–L40 curve in
Since this PR adds a second input representation, it would be particularly |
hhhizzz
left a comment
hhhizzz
left a comment
There was a problem hiding this comment.
The direct bitmap path looks valuable, especially for fragmented selections,
but I found several performance issues that seem worth addressing:
- a reproducible regression in existing selector/mask cursor paths; (Addressed in last comment)
- dense-mask trimming enumerating every selected bit;
- repeated full popcounts across row groups;
- cached selectors being materialized a second time;
- loss of slice-iterator optimizations; and
- mixed representations eagerly expanding fragmented masks back to RLE.
The first three look blocking to me; the remaining ones should at least be
resolved or explicitly scoped as follow-ups with benchmark coverage.
| self.iter().copied().collect() | ||
| } | ||
|
|
||
| fn into_selectors_vec(self) -> Vec<RowSelector> { |
There was a problem hiding this comment.
Could we consume the lazy selector cache here instead of always materializing
the mask again?
RowSelection::iter() populates MaskSelection::selectors, but
into_selectors_vec() ignores that cache and calls mask_to_selectors() a
second time. Current DataFusion does exactly this sequence in
ParquetAccessPlan::into_overall_row_selection: it first calls iter() for
validation and then consumes the same selection with selection.into().
For an alternating 3M-row mask, each selector Vec has a 64 MiB capacity. This
path therefore performs a second full scan and can temporarily retain about
128 MiB of selector storage.
Could MaskSelection provide an owning into_selectors() that uses
OnceLock::into_inner() when initialized, and share it with the cursor
conversion path?
| type Item = &'a RowSelector; | ||
|
|
||
| #[inline] | ||
| fn next(&mut self) -> Option<Self::Item> { |
There was a problem hiding this comment.
Could RowSelectionIter preserve the underlying slice iterator's optimized
methods?
The previous iter() implementation delegated directly to slice::Iter, so
its exact size_hint and specialized nth / count implementations were used.
This wrapper only forwards next(), which falls back to (0, None) for
size_hint; consequently, existing selector-backed code such as
selection.iter().copied().collect::<Vec<_>>() now repeatedly grows and copies
the Vec, while count() and nth() become element-by-element walks.
Since both backing types expose a slice iterator after any lazy materialization,
could this delegate size_hint, nth, count, and last, and ideally
implement ExactSizeIterator (plus DoubleEndedIterator / FusedIterator)?
There was a problem hiding this comment.
This also came up with claude code in my review
I think it would basically look like
impl<'a> Iterator for RowSelectionIter<'a> {
type Item = &'a RowSelector;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
impl ExactSizeIterator for RowSelectionIter<'_> {}
// once it returns None, it will continue returning None
impl std::iter::FusedIterator for RowSelectionIter<'_> {}| (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => { | ||
| intersect_row_selections(l, r) | ||
| } | ||
| (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => { |
There was a problem hiding this comment.
Could the mixed-representation path avoid eagerly converting the entire mask
with mask_to_selectors? The same issue also applies to the mixed union arms
below.
A natural target use case is an externally supplied fragmented bitmap composed
with selector-backed page-index pruning. DataFusion's
ParquetAccessPlan::scan_selection performs exactly this kind of
intersection. An alternating 3M-row mask can therefore materialize up to 3M
RowSelectors before reading begins, reintroducing the memory cost this PR is
intended to avoid.
Could this either stream the mixed operation, or choose which side to convert
based on the two representations' shapes and preserve mask backing when that is
cheaper? A mixed bitmap + page-selection benchmark would also protect this
case.
- rewrite boolean_mask_from_selectors to fill bitmap bytes directly, fixing a 3-11% read regression caused by append_n no longer being inlined on this hot path - trim_mask: fast-path a set final bit and reverse-scan bytes for the last set bit instead of walking every selected bit - cache the selected-row count in MaskSelection and propagate it through split_off/trim/offset/limit so the row-group frontier no longer rescans the shrinking bitmap tail per row group
Runs every matrix point with RowSelectionPolicy::default() (read_auto) in addition to the two forced policies, and repeats all modes with a mask-backed RowSelection as input (*_mask_backed groups). Selector-backed groups keep their historical names so existing baselines stay comparable.
|
Thanks for the thorough review @hhhizzz — I verified all six findings and they are accurate. The three blocking items are fixed in 496d7cb; 325d774 extends the 1. Regression — fixedProfiling showed the regression was entirely in the selector→bitmap conversion: after the module split, LLVM stopped inlining L4–L40 matrix, 10141-auto 325d774 vs main 
Both forced paths are now faster than the main at every point, and the Selectors/Mask crossover (between L12 and L16 on this machine) is unchanged by this PR. 2. Dense-mask trimming — fixed
3. Repeated popcounts — fixed
4–6These are valid findings as well. To keep this PR from growing too large I'd like to address them in follow-up PRs, and I'll file issues after this pr merge for them so they don't get lost. While doing that I also think we should split up |
the command for below figure benchmark that 325d774 added the bitmap input representation into row_selection_cursor 
|
|
I reran the forced-Selectors benchmark on an AMD EPYC 9254 using the exact same historical 2-mode benchmark source for both I then split L40 into selection clone, reader creation, read-plan build, and reader drain. Clone ( However, this does not appear to be additional work introduced by the selector cursor. The selector hot path is unchanged, and the normalized My conclusion is that the residual forced-Selectors difference is a default multi-CGU inlining/code-layout artifact rather than a cursor algorithm regression. I would recommend not change |
|
#10288 is about to land and changes Mask execution across page-pruned data. When rebasing this PR, could we preserve its new_mask_from_buffer(mask, loaded_row_ranges)#10288 only has Could we also add a sparse-page regression test starting from a direct, preferably non-byte-aligned, The old |
optimized this part in 3f89ed0 
i will add those test case after #10288 merge. |
I just merged |
Resolves conflicts with apache#10288 (mask filtering across skipped pages): - Port LoadedRowRanges-aware MaskCursor (next_chunk / next_mask_chunk_non_empty) into selection/boolean.rs - Thread loaded_row_ranges into both RowSelectionCursor::new_mask_from_buffer and new_mask_from_selectors - Keep LoadedRowRanges in selection/mod.rs
- mask_backed_plan_respects_loaded_row_ranges: unit test that a Mask cursor built via new_mask_from_buffer from a non-byte-aligned BooleanBuffer::slice stops decoded chunks at loaded-range boundaries - test_mask_from_boolean_buffer_slice_with_sparse_pages: end-to-end test that RowSelection::from_boolean_buffer crossing pruned pages reads correctly under both Mask and Auto policies
Coverage-guided additions (cargo llvm-cov): - selectors_policy_lowers_mask_backed_selection: the Selectors policy x Mask backing arm of build_cursor, the only untested strategy/backing combination - intersect_masks with a longer right operand (tail pass-through) - and_then where the inner selection selects nothing - mixed mask/selector PartialEq inequality paths (length mismatch, set bit inside a skip run, misaligned select run) - zero-length RowSelector filtering and offset(0)/zero-batch-size identity paths
|
Thanks for the updates. I re-ran the The new early-exit run-count logic substantially reduces mask-backed Auto overhead (for example, L4 drops from +24.5% to +2.8%), and the measured crossover remains between L28 and L32, so the current threshold of 32 still looks appropriate. The residual selector variation is codegen/layout-sensitive rather than additional cursor work. The sparse-page integration and added coverage also look good. LGTM from both correctness and performance perspectives. |
|
Thanks for your reviews @hhhizzz |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you @haohuaijin and @hhhizzz -- I found this PR very easy to read and understand; It is also super well tested, and I think is a foundational piece for more efficient filter evaluation in parquet
cc @etseidl for your interest in all things parquet
cc @devanbenz as maybe some of your BMI work is relevant / could be used here
cc @Dandandan as we have discussed this before in the context of pushdown filter work (when the manipulation of RowSelections was expensive)
cc @adriangb as I think you are also interested in this
| #[derive(Default, Clone)] | ||
| pub struct RowSelection { | ||
| selectors: Vec<RowSelector>, | ||
| inner: RowSelectionInner, |
| /// RowSelection::from_consecutive_ranges(ranges.into_iter(), 20); | ||
| /// let actual: Vec<RowSelector> = selection.into(); | ||
| /// assert_eq!(actual, expected); | ||
| /// |
There was a problem hiding this comment.
It might be worth updating the descriptio of this struct (we can do it in a follow on PR too)
The current one talks about skipping/selecting
/// [`RowSelection`] allows selecting or skipping a provided number of rows
/// when scanning the parquet file.
///
/// This is applied prior to reading column data, and can therefore
/// be used to skip IO to fetch data into memory
Maybe we could focus on the selection aspect part
/// [`RowSelection`] represents selecting a subset of rows
/// when scanning a parquet file.
///
/// This is applied prior to reading column data, and can therefore
/// be used to skip IO to fetch data into memory.
There was a problem hiding this comment.
And then we could add a note somewhere like
///
/// Depending on the pattern or rows to be selected, [`RowSelection`] has
/// either a bitmap or an RLE ([`RowSelector`]) based implementation.
| type Item = &'a RowSelector; | ||
|
|
||
| #[inline] | ||
| fn next(&mut self) -> Option<Self::Item> { |
There was a problem hiding this comment.
This also came up with claude code in my review
I think it would basically look like
impl<'a> Iterator for RowSelectionIter<'a> {
type Item = &'a RowSelector;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
impl ExactSizeIterator for RowSelectionIter<'_> {}
// once it returns None, it will continue returning None
impl std::iter::FusedIterator for RowSelectionIter<'_> {}| /// assert_eq!(selection.row_count(), 3); | ||
| /// ``` | ||
| /// | ||
| /// A [`RowSelection`] maintains the following invariants: |
There was a problem hiding this comment.
The note on invariants I think only apply to RLE/Selector backed one -- it would be good to update this documentation
| } | ||
|
|
||
| impl RowSelection { | ||
| fn from_selectors(selectors: Vec<RowSelector>) -> Self { |
There was a problem hiding this comment.
Should this be pub for parity with the other constructors like from_boolean_buffer? I think there is a From<Vec<RowSelector>> impl, but it is strange that the named method is not also pub
There was a problem hiding this comment.
Claude says "no because this does no validation" -- so maybe a small comment saying that would help
| /// Returns the underlying mask if this selection is mask-backed. | ||
| /// | ||
| /// Public so that engines composing selections (e.g. DataFusion's | ||
| /// `ParquetAccessPlan::into_overall_row_selection`) can concatenate |
There was a problem hiding this comment.
if we are going to refer to this, it probably would make sense to make it hyperlink rather than a name
| } else { | ||
| RowSelectionStrategy::Selectors | ||
| } | ||
| selection.auto_selection_strategy(threshold) |
There was a problem hiding this comment.
Well that is certainly nicer - I like how it is now much more encapsulated
| if l.len() == r.len() { | ||
| return l | r; | ||
| } | ||
| let common = l.len().min(r.len()); |
There was a problem hiding this comment.
you might be able to do this more efficiently by copying the longer one and then using &= -- maybe in some follow on
| pub use selection::{ | ||
| MaskRunIter, RowSelection, RowSelectionCursor, RowSelectionPolicy, RowSelector, | ||
| }; |
There was a problem hiding this comment.
Since it is pub I (and claude) suggest also adding RowSelectionIter here
| pub use selection::{ | |
| MaskRunIter, RowSelection, RowSelectionCursor, RowSelectionPolicy, RowSelector, | |
| }; | |
| pub use selection::{ | |
| MaskRunIter, RowSelection, RowSelectionCursor, RowSelectionIter, RowSelectionPolicy, | |
| RowSelector, | |
| }; |
5 participants
Which issue does this PR close?
RowSelectionto be backed by aBooleanBufferto reduce memory usage #10140Rationale for this change
RowSelectioncurrently stores selections asVec<RowSelector>(16 bytes per selector). This is compact for long runs, but expensive for scattered matches. With ~35% isolated single-row hits, it uses about 11.2 bytes per input row. ABooleanBufferuses 1 bit per input row, about 90x less memory.The reader can also choose the
Maskstrategy, which converts selectors back into a bitmap. When the caller already had a bitmap, this conversion round-trip is unnecessary.This PR lets
RowSelectionpreserve a caller-provided bitmap and pass it directly to mask execution.This is not intended to claim broad DataFusion / TPC-DS / ClickBench speedups. Current common DataFusion SQL paths generally do not naturally produce bitmap-backed
RowSelections. The practical benefit is for integrations that already have a row-level bitmap and need Parquet to consume it without materializing a large selector list.What changes are included in this PR?
RowSelectioncan now be backed by eitherVec<RowSelector>orBooleanBuffer. New public construction:Methods that can work directly on the bitmap now do so:
iter()still returns&RowSelector(non-breaking); mask-backed selections lazily materialize a selector cache on first call, while internal hot paths bypass it and use theBooleanBuffer/MaskRunIterdirectlyrow_count/skipped_row_countuse a cached popcountselects_anyusesset_indices().next()trimpreserves mask backing viaBooleanBuffer::sliceintersection/uniononMask+MaskuseBitAnd/BitOrsplit_offon a mask usesBooleanBuffer::slice(O(1), both halves stay mask-backed)limitslices at the selected-row boundary viafind_nth_set_bit_position, staying mask-backedoffsetfinds the first selected row to keep viafind_nth_set_bit_positionand rebuilds only the mask buffer, avoiding selector materializationand_thenapplies the inner selection over the mask's set positions, returning a mask-backed resultFromIterator<RowSelection>concatenatesBooleanBuffers when every input is mask-backedMixed inputs, and existing selector-backed inputs, still use the existing selector helpers. Existing callers keep the same behavior.
The reader (
ReadPlanBuilder::build) passes a mask-backed selection straight toRowSelectionCursor::new_mask_from_buffer, so it skips rebuilding the bitmap from selectors.Autoresolution works directly on the bitmap (early-exit run counting), without converting the backing.Integrates with #10288: both mask cursor constructors carry
LoadedRowRanges, so a caller-provided bitmap stays within loaded pages when page pruning skips pages.Also adds
MaskRunIter+RowSelection::as_maskfor zero-allocation RLE iteration over a mask, therow_selector_boolean_bufferbenchmark, andread_auto/ mask-backed input modes inrow_selection_cursor.Are these changes tested?
Yes. This PR extends the existing
RowSelectionunit tests with coverage for:BooleanBuffer, including empty and all-unset masksFrom<BooleanBuffer>split_off,limit,offset,and_then, and all-maskFromIterator<RowSelection>intersection/union, including uneven-length inputsfrom_filtersselector pathBooleanBuffer::slice(...), under bothMaskandAutopolicies (fix(parquet): support mask filtering across skipped pages #10288 integration seam)LoadedRowRangesboundariesboolean_mask_from_selectorsandtrim_mask(including non-zero offsets)cargo llvm-covAre there any user-facing changes?
No breaking API changes. New public APIs:
RowSelection::from_boolean_buffer,From<BooleanBuffer>,RowSelection::as_mask,MaskRunIter.