hhhizzz commented on code in PR #10141: URL: https://github.com/apache/arrow-rs/pull/10141#discussion_r3608539782
########## parquet/src/arrow/arrow_reader/selection/boolean.rs: ########## @@ -0,0 +1,962 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::{RowSelection, RowSelectionInner, RowSelector}; +use crate::errors::ParquetError; +use arrow_array::BooleanArray; +use arrow_buffer::bit_iterator::BitSliceIterator; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use std::cmp::Ordering; +use std::sync::OnceLock; + +/// Mask-backed [`RowSelection`] storage. +/// +/// `selectors` is only populated if callers use the borrowed [`RowSelection::iter`] +/// compatibility API. Internal paths that can stream or consume the bitmap avoid +/// this cache. +#[derive(Debug)] +pub(crate) struct MaskSelection { + mask: BooleanBuffer, + selectors: OnceLock<Vec<RowSelector>>, +} + +impl MaskSelection { + pub(super) fn new(mask: BooleanBuffer) -> Self { + Self { + mask, + selectors: OnceLock::new(), + } + } + + pub(crate) fn mask(&self) -> &BooleanBuffer { + &self.mask + } + + pub(crate) fn into_mask(self) -> BooleanBuffer { + let Self { mask, .. } = self; + mask + } + + pub(super) fn selectors(&self) -> &[RowSelector] { + self.selectors + .get_or_init(|| mask_to_selectors(&self.mask)) + .as_slice() + } +} + +impl Clone for MaskSelection { + fn clone(&self) -> Self { + Self::new(self.mask.clone()) + } +} + +/// Streaming RLE view of a [`BooleanBuffer`], yielding owned [`RowSelector`]s +/// without allocation. +/// +/// Useful as a zero-cost alternative to [`RowSelection::iter`] for mask-backed +/// selections, via [`RowSelection::as_mask`]: +/// +/// ```ignore +/// if let Some(mask) = selection.as_mask() { +/// for run in MaskRunIter::new(mask) { ... } +/// } +/// ``` +#[derive(Debug)] +pub struct MaskRunIter<'a> { + slices: BitSliceIterator<'a>, + cursor: usize, + total: usize, + pending: Option<RowSelector>, + finished: bool, +} + +impl<'a> MaskRunIter<'a> { + /// Create a streaming RLE iterator over a [`BooleanBuffer`]. + pub fn new(mask: &'a BooleanBuffer) -> Self { + Self { + slices: mask.set_slices(), + cursor: 0, + total: mask.len(), + pending: None, + finished: false, + } + } +} + +impl Iterator for MaskRunIter<'_> { + type Item = RowSelector; + + fn next(&mut self) -> Option<RowSelector> { + if let Some(p) = self.pending.take() { + return Some(p); + } + if self.finished { + return None; + } + match self.slices.next() { + Some((start, end)) => { + let select = RowSelector::select(end - start); + if start > self.cursor { + let skip = RowSelector::skip(start - self.cursor); + self.pending = Some(select); + self.cursor = end; + Some(skip) + } else { + self.cursor = end; + Some(select) + } + } + None => { + self.finished = true; + if self.cursor < self.total { + let skip = RowSelector::skip(self.total - self.cursor); + self.cursor = self.total; + Some(skip) + } else { + None + } + } + } + } +} + +/// Cursor for iterating a mask-backed [`RowSelection`] +/// +/// This is best for dense selections where there are many small skips +/// or selections. For example, selecting every other row. +#[derive(Debug)] +pub struct MaskCursor { + pub(super) mask: BooleanBuffer, + /// Current absolute offset into the selection + pub(super) position: usize, +} + +impl MaskCursor { + /// Returns `true` when no further rows remain + pub fn is_empty(&self) -> bool { + self.position >= self.mask.len() + } + + /// Advance through the mask representation, producing the next chunk summary + pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option<MaskChunk> { + let (initial_skip, chunk_rows, selected_rows, mask_start, end_position) = { + let mask = &self.mask; + + if self.position >= mask.len() { + return None; + } + + let start_position = self.position; + let mut cursor = start_position; + let mut initial_skip = 0; + + while cursor < mask.len() && !mask.value(cursor) { + initial_skip += 1; + cursor += 1; + } + + let mask_start = cursor; + let mut chunk_rows = 0; + let mut selected_rows = 0; + + // Advance until enough rows have been selected to satisfy the batch size, + // or until the mask is exhausted. This mirrors the behaviour of the legacy + // `RowSelector` queue-based iteration. + while cursor < mask.len() && selected_rows < batch_size { + chunk_rows += 1; + if mask.value(cursor) { + selected_rows += 1; + } + cursor += 1; + } + + (initial_skip, chunk_rows, selected_rows, mask_start, cursor) + }; + + self.position = end_position; + + Some(MaskChunk { + initial_skip, + chunk_rows, + selected_rows, + mask_start, + }) + } + + /// Materialise the boolean values for a mask-backed chunk + pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result<BooleanArray, ParquetError> { + if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() { + return Err(ParquetError::General( + "Internal Error: MaskChunk exceeds mask length".to_string(), + )); + } + Ok(BooleanArray::from( + self.mask.slice(chunk.mask_start, chunk.chunk_rows), + )) + } +} + +/// Result of computing the next chunk to read when using a [`MaskCursor`] +#[derive(Debug)] +pub struct MaskChunk { + /// Number of leading rows to skip before reaching selected rows + pub initial_skip: usize, + /// Total rows covered by this chunk (selected + skipped) + pub chunk_rows: usize, + /// Rows actually selected within the chunk + pub selected_rows: usize, + /// Starting offset within the mask where the chunk begins + pub mask_start: usize, +} + +/// Materialize a [`BooleanBuffer`] into its RLE form. +pub(crate) fn mask_to_selectors(mask: &BooleanBuffer) -> Vec<RowSelector> { + let total_rows = mask.len(); + if total_rows == 0 { + return Vec::new(); + } + let mut selectors: Vec<RowSelector> = Vec::new(); + let mut last_end = 0; + for (start, end) in mask.set_slices() { + if start > last_end { + selectors.push(RowSelector::skip(start - last_end)); + } + selectors.push(RowSelector::select(end - start)); + last_end = end; + } + if last_end != total_rows { + selectors.push(RowSelector::skip(total_rows - last_end)); + } + selectors +} + +pub(super) fn mask_run_count(mask: &BooleanBuffer) -> usize { + let total_rows = mask.len(); + if total_rows == 0 { + return 0; + } + + let mut run_count = 0; + let mut last_end = 0; + for (start, end) in mask.set_slices() { + if start > last_end { + run_count += 1; + } + if end > start { + run_count += 1; + } + last_end = end; + } + if last_end < total_rows { + run_count += 1; + } + + run_count +} + +/// Bitwise AND of two mask-backed selections. Longer side's tail passes through. +pub(super) fn intersect_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer { + if l.len() == r.len() { + return l & r; + } + let common = l.len().min(r.len()); + let head = &l.slice(0, common) & &r.slice(0, common); + let (longer, longer_len) = if l.len() > r.len() { + (l, l.len()) + } else { + (r, r.len()) + }; + let tail = longer.slice(common, longer_len - common); + let mut builder = BooleanBufferBuilder::new(longer_len); + builder.append_buffer(&head); + builder.append_buffer(&tail); + builder.finish() +} + +/// Bitwise OR of two mask-backed selections. Longer side's tail passes through. +pub(super) fn union_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer { + if l.len() == r.len() { + return l | r; + } + let common = l.len().min(r.len()); + let head = &l.slice(0, common) | &r.slice(0, common); + let (longer, longer_len) = if l.len() > r.len() { + (l, l.len()) + } else { + (r, r.len()) + }; + let tail = longer.slice(common, longer_len - common); + let mut builder = BooleanBufferBuilder::new(longer_len); + builder.append_buffer(&head); + builder.append_buffer(&tail); + builder.finish() +} + +/// Applies `other` to the selected rows of `mask`, preserving the original row domain. +pub(super) fn and_then_mask(mask: &BooleanBuffer, other: &RowSelection) -> BooleanBuffer { + match &other.inner { + RowSelectionInner::Mask(other_mask) => and_then_masks(mask, other_mask.mask()), + RowSelectionInner::Selectors(selectors) => { + and_then_mask_from_selectors(mask, selectors.iter().copied()) + } + } +} + +fn and_then_mask_from_selectors<I>(mask: &BooleanBuffer, other: I) -> BooleanBuffer +where + I: IntoIterator<Item = RowSelector>, +{ + let mut builder = BooleanBufferBuilder::new(mask.len()); + let mut other_iter = other.into_iter(); + let mut current = other_iter.next(); + let mut cursor = 0usize; + + // Iterate only over the set positions in `mask`; the gaps of unset bits + // are filled in bulk with `append_n` instead of bit-by-bit. + for set_idx in mask.set_indices() { + if set_idx > cursor { + builder.append_n(set_idx - cursor, false); + } + cursor = set_idx + 1; + + while current.as_ref().is_some_and(|s| s.row_count == 0) { + current = other_iter.next(); + } + let selector = current + .as_mut() + .expect("selection contains less than the number of selected rows"); + let selected = !selector.skip; + selector.row_count -= 1; + builder.append(selected); + } + if cursor < mask.len() { + builder.append_n(mask.len() - cursor, false); + } + + if current.is_some_and(|s| s.row_count != 0) || other_iter.any(|s| s.row_count != 0) { + panic!("selection exceeds the number of selected rows"); + } + + builder.finish() +} + +fn and_then_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer { + let selected_count = mask.count_set_bits(); + match other.len().cmp(&selected_count) { + Ordering::Less => panic!("selection contains less than the number of selected rows"), + Ordering::Greater => panic!("selection exceeds the number of selected rows"), + Ordering::Equal => {} + } + + let other_true_count = other.count_set_bits(); + if other_true_count == 0 { + return BooleanBuffer::new_unset(mask.len()); + } + if other_true_count == selected_count { + return mask.clone(); + } + + let mut builder = BooleanBufferBuilder::new(mask.len()); + let mut outer_set_indices = mask.set_indices(); + let mut next_selected_ordinal = 0usize; + let mut cursor = 0usize; + + for selected_ordinal in other.set_indices() { + let skip = selected_ordinal - next_selected_ordinal; + let set_idx = outer_set_indices + .nth(skip) + .expect("validated other length matches selected row count"); + if set_idx > cursor { + builder.append_n(set_idx - cursor, false); + } + builder.append(true); + cursor = set_idx + 1; + next_selected_ordinal = selected_ordinal + 1; + } + + if cursor < mask.len() { + builder.append_n(mask.len() - cursor, false); + } + + builder.finish() +} + +/// Split a mask into `(head, tail)` at `row_count`, preserving an empty mask tail +/// when the split point is past the end. +pub(super) fn split_off_mask( + mask: BooleanBuffer, + row_count: usize, +) -> (BooleanBuffer, BooleanBuffer) { + let total = mask.len(); + if row_count >= total { + return (mask, BooleanBuffer::new_unset(0)); + } + + let head = mask.slice(0, row_count); + let tail = mask.slice(row_count, total - row_count); + (head, tail) +} + +/// Trims trailing unset bits from a mask-backed selection. +pub(super) fn trim_mask(mask: &BooleanBuffer) -> Option<BooleanBuffer> { Review Comment: Could we avoid walking every selected bit here? `trim_mask` first computes a full popcount and then calls `find_nth_set_bit_position(0, popcount)`. That helper is implemented using `set_indices().nth(n - 1)`, and `BitIndexIterator` does not override `nth`. Therefore a dense mask enumerates every selected row just to find its final set bit. `ReadPlanBuilder::build` calls `trim()` for every selection, so an all-set or near-dense 1B-row bitmap can execute up to 1B iterator steps before decoding starts, even when there is nothing to trim. Could we at least fast-path a set final bit, and use a reverse word scan for the remaining case? ########## parquet/src/arrow/arrow_reader/selection/mod.rs: ########## @@ -422,96 +662,151 @@ impl RowSelection { /// /// returned: NYYYYYNNYYNYN pub fn union(&self, other: &Self) -> Self { - union_row_selections(&self.selectors, &other.selectors) + match &self.inner { + RowSelectionInner::Mask(l) => match &other.inner { + RowSelectionInner::Mask(r) => { + Self::from_boolean_buffer(union_masks(l.mask(), r.mask())) + } + RowSelectionInner::Selectors(r) => { + let l = mask_to_selectors(l.mask()); + union_row_selections(&l, r) + } + }, + RowSelectionInner::Selectors(l) => match &other.inner { + RowSelectionInner::Mask(r) => { + let r = mask_to_selectors(r.mask()); + union_row_selections(l, &r) + } + RowSelectionInner::Selectors(r) => union_row_selections(l, r), + }, + } } /// Returns `true` if this [`RowSelection`] selects any rows pub fn selects_any(&self) -> bool { - self.selectors.iter().any(|x| !x.skip) + match &self.inner { + RowSelectionInner::Selectors(s) => s.iter().any(|x| !x.skip), + RowSelectionInner::Mask(m) => m.mask().set_indices().next().is_some(), + } } /// Trims this [`RowSelection`] removing any trailing skips pub(crate) fn trim(mut self) -> Self { - while self.selectors.last().map(|x| x.skip).unwrap_or(false) { - self.selectors.pop(); + if let RowSelectionInner::Mask(m) = &self.inner { + if let Some(mask) = trim_mask(m.mask()) { + return Self::from_boolean_buffer(mask); + } + return self; + } + let selectors = self.selectors_mut(); + while selectors.last().map(|x| x.skip).unwrap_or(false) { + selectors.pop(); } self } /// Applies an offset to this [`RowSelection`], skipping the first `offset` selected rows - pub(crate) fn offset(mut self, offset: usize) -> Self { + pub(crate) fn offset(self, offset: usize) -> Self { if offset == 0 { return self; } + let mut selectors = match self.inner { + RowSelectionInner::Mask(mask) => { + return Self::from_boolean_buffer(offset_mask((*mask).into_mask(), offset)); + } + RowSelectionInner::Selectors(selectors) => selectors, + }; let mut selected_count = 0; let mut skipped_count = 0; // Find the index where the selector exceeds the row count - let find = self - .selectors - .iter() - .position(|selector| match selector.skip { - true => { - skipped_count += selector.row_count; - false - } - false => { - selected_count += selector.row_count; - selected_count > offset - } - }); + let find = selectors.iter().position(|selector| match selector.skip { + true => { + skipped_count += selector.row_count; + false + } + false => { + selected_count += selector.row_count; + selected_count > offset + } + }); let split_idx = match find { Some(idx) => idx, None => { - self.selectors.clear(); - return self; + selectors.clear(); + return Self::from_selectors(selectors); } }; - let mut selectors = Vec::with_capacity(self.selectors.len() - split_idx + 1); - selectors.push(RowSelector::skip(skipped_count + offset)); - selectors.push(RowSelector::select(selected_count - offset)); - selectors.extend_from_slice(&self.selectors[split_idx + 1..]); + let mut new_selectors = Vec::with_capacity(selectors.len() - split_idx + 1); + new_selectors.push(RowSelector::skip(skipped_count + offset)); + new_selectors.push(RowSelector::select(selected_count - offset)); + new_selectors.extend_from_slice(&selectors[split_idx + 1..]); - Self { selectors } + Self::from_selectors(new_selectors) } /// Limit this [`RowSelection`] to only select `limit` rows - pub(crate) fn limit(mut self, mut limit: usize) -> Self { + pub(crate) fn limit(self, mut limit: usize) -> Self { + let mut selectors = match self.inner { + RowSelectionInner::Mask(mask) => { + return Self::from_boolean_buffer(limit_mask((*mask).into_mask(), limit)); + } + RowSelectionInner::Selectors(selectors) => selectors, + }; if limit == 0 { - self.selectors.clear(); + selectors.clear(); } - for (idx, selection) in self.selectors.iter_mut().enumerate() { + for (idx, selection) in selectors.iter_mut().enumerate() { if !selection.skip { if selection.row_count >= limit { selection.row_count = limit; - self.selectors.truncate(idx + 1); + selectors.truncate(idx + 1); break; } else { limit -= selection.row_count; } } } - self + Self::from_selectors(selectors) } - /// Returns an iterator over the [`RowSelector`]s for this - /// [`RowSelection`]. - pub fn iter(&self) -> impl Iterator<Item = &RowSelector> { - self.selectors.iter() + /// Returns a borrowed iterator yielding the [`RowSelector`]s for this selection. + /// + /// Mask-backed selections materialize a `Vec<RowSelector>` cache on first + /// call (one allocation, `O(set_slices)` work) so the iterator can hand out + /// `&RowSelector`; the cache is not copied on clone. For single-pass walks + /// over mask-backed selections, prefer streaming directly via + /// [`Self::as_mask`] + [`MaskRunIter::new`] — that path is allocation-free + /// and avoids populating the cache. + pub fn iter(&self) -> RowSelectionIter<'_> { + match &self.inner { + RowSelectionInner::Selectors(s) => RowSelectionIter(s.iter()), + RowSelectionInner::Mask(m) => RowSelectionIter(m.selectors().iter()), + } } /// Returns the number of selected rows pub fn row_count(&self) -> usize { Review Comment: `row_count()` is a full bitmap popcount for mask-backed selections. The row-group frontier calls it on the entire shrinking tail before every row group, then again on the split head. A caller using `peek_next_row_group` repeats the same traversal on a clone. This makes the mask path O(number_of_row_groups * bitmap_size). For example, a 1B-row bitmap is about 125 MB; with 100 equal row groups, the shrinking-tail checks alone rescan about 6.3 GB of bitmap data. Could `MaskSelection` lazily cache the selected count and propagate head/tail counts through `split_off`, or could the frontier explicitly track the remaining selected count? ########## parquet/src/arrow/arrow_reader/selection/mod.rs: ########## @@ -412,7 +637,22 @@ impl RowSelection { /// /// returned: NNNNNNNNYYNYN pub fn intersection(&self, other: &Self) -> Self { - intersect_row_selections(&self.selectors, &other.selectors) + match (&self.inner, &other.inner) { + (RowSelectionInner::Mask(l), RowSelectionInner::Mask(r)) => { + Self::from_boolean_buffer(intersect_masks(l.mask(), r.mask())) + } + (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => { + intersect_row_selections(l, r) + } + (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => { Review 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 `RowSelector`s 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. ########## parquet/src/arrow/arrow_reader/selection/mod.rs: ########## @@ -135,12 +149,297 @@ impl RowSelector { /// * Consecutive [`RowSelector`]s alternate skipping or selecting rows /// /// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData -#[derive(Debug, Clone, Default, Eq, PartialEq)] +#[derive(Default, Clone)] pub struct RowSelection { - selectors: Vec<RowSelector>, + inner: RowSelectionInner, +} + +/// Internal storage for [`RowSelection`]. +#[derive(Debug, Clone)] +pub(crate) enum RowSelectionInner { + Selectors(Vec<RowSelector>), + Mask(Box<MaskSelection>), +} + +impl Default for RowSelectionInner { + fn default() -> Self { + Self::Selectors(Vec::new()) + } +} + +impl std::fmt::Debug for RowSelection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.inner { + RowSelectionInner::Selectors(s) => f + .debug_struct("RowSelection") + .field("selectors", s) + .finish(), + RowSelectionInner::Mask(m) => f + .debug_struct("RowSelection") + .field("mask_len", &m.mask().len()) + .finish_non_exhaustive(), + } + } +} + +impl PartialEq for RowSelection { + fn eq(&self, other: &Self) -> bool { + match (&self.inner, &other.inner) { + (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b)) => a == b, + (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a.mask() == b.mask(), + (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors)) + | (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask)) => { + if selectors + .iter() + .try_fold(0usize, |acc, selector| acc.checked_add(selector.row_count)) + != Some(mask.mask().len()) + { + return false; + } + + let mut slices = mask.mask().set_slices().peekable(); + let mut cursor = 0usize; + + for selector in selectors { + let end = cursor + selector.row_count; + + if selector.skip { + if slices.peek().is_some_and(|(start, _)| *start < end) { + return false; + } + } else { + match slices.next() { + Some((start, slice_end)) if start == cursor && slice_end == end => {} + _ => return false, + } + } + + cursor = end; + } + + slices.next().is_none() + } + } + } +} + +impl Eq for RowSelection {} + +/// Borrowed iterator over the [`RowSelector`]s of a [`RowSelection`]. +#[derive(Debug)] +pub struct RowSelectionIter<'a>(std::slice::Iter<'a, RowSelector>); + +impl<'a> Iterator for RowSelectionIter<'a> { + type Item = &'a RowSelector; + + #[inline] + fn next(&mut self) -> Option<Self::Item> { + self.0.next() + } +} + +#[inline] +fn scan_ranges_from_selectors<I>(selectors: I, page_locations: &[PageLocation]) -> Vec<Range<u64>> +where + I: IntoIterator<Item = RowSelector>, +{ + let mut ranges: Vec<Range<u64>> = vec![]; + let mut row_offset = 0; + + let mut pages = page_locations.iter().peekable(); + let mut selectors = selectors.into_iter(); + let mut current_selector = selectors.next(); + let mut current_page = pages.next(); + + let mut current_page_included = false; + + while let Some((selector, page)) = current_selector.as_mut().zip(current_page) { + if !(selector.skip || current_page_included) { + let start = page.offset as u64; + let end = start + page.compressed_page_size as u64; + ranges.push(start..end); + current_page_included = true; + } + + if let Some(next_page) = pages.peek() { + if row_offset + selector.row_count > next_page.first_row_index as usize { + let remaining_in_page = next_page.first_row_index as usize - row_offset; + selector.row_count -= remaining_in_page; + row_offset += remaining_in_page; + current_page = pages.next(); + current_page_included = false; + + continue; + } else { + if row_offset + selector.row_count == next_page.first_row_index as usize { + current_page = pages.next(); + current_page_included = false; + } + row_offset += selector.row_count; + current_selector = selectors.next(); + } + } else { + if !(selector.skip || current_page_included) { + let start = page.offset as u64; + let end = start + page.compressed_page_size as u64; + ranges.push(start..end); + } + current_selector = selectors.next() + } + } + + ranges +} + +#[inline] +fn expand_to_batch_boundaries_from_selectors<I>( + selectors: I, + batch_size: usize, + total_rows: usize, +) -> RowSelection +where + I: IntoIterator<Item = RowSelector>, +{ + let mut expanded_ranges = Vec::new(); + let mut row_offset = 0; + + for selector in selectors { + if selector.skip { + row_offset += selector.row_count; + } else { + let start = row_offset; + let end = row_offset + selector.row_count; + + // Expand start to batch boundary + let expanded_start = (start / batch_size) * batch_size; + // Expand end to batch boundary + let expanded_end = end.div_ceil(batch_size) * batch_size; + let expanded_end = expanded_end.min(total_rows); + + expanded_ranges.push(expanded_start..expanded_end); + row_offset += selector.row_count; + } + } + + // Sort ranges by start position + expanded_ranges.sort_by_key(|range| range.start); + + // Merge overlapping or consecutive ranges + let mut merged_ranges: Vec<Range<usize>> = Vec::new(); + for range in expanded_ranges { + if let Some(last) = merged_ranges.last_mut() { + if range.start <= last.end { + // Overlapping or consecutive - merge them + last.end = last.end.max(range.end); + } else { + // No overlap - add new range + merged_ranges.push(range); + } + } else { + // First range + merged_ranges.push(range); + } + } + + RowSelection::from_consecutive_ranges(merged_ranges.into_iter(), total_rows) } impl RowSelection { + fn from_selectors(selectors: Vec<RowSelector>) -> Self { + Self { + inner: RowSelectionInner::Selectors(selectors), + } + } + + /// Create a [`RowSelection`] from a packed [`BooleanBuffer`]. + /// + /// Each set bit selects a row, each unset bit skips one. Unlike + /// [`Self::from_filters`], the bitmap is kept as-is rather than + /// eagerly run-length-encoded. [`Self::iter`] materializes and caches the + /// RLE form on first use; use [`MaskRunIter`] to stream the RLE form + /// directly from the bitmap. + pub fn from_boolean_buffer(mask: BooleanBuffer) -> Self { + Self { + inner: RowSelectionInner::Mask(Box::new(MaskSelection::new(mask))), + } + } + + /// 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 + /// mask-backed selections without materialising the RLE form. + pub fn as_mask(&self) -> Option<&BooleanBuffer> { + match &self.inner { + RowSelectionInner::Mask(m) => Some(m.mask()), + _ => None, + } + } + + /// Consume the selection and return its internal storage. + pub(crate) fn into_inner(self) -> RowSelectionInner { + self.inner + } + + /// Choose the automatic materialisation strategy without converting between + /// selector and mask backing. + #[inline] + pub(crate) fn auto_selection_strategy(&self, threshold: usize) -> RowSelectionStrategy { + let (total_rows, effective_count) = match &self.inner { + RowSelectionInner::Selectors(selectors) => { + selectors.iter().fold((0usize, 0usize), |(rows, count), s| { + if s.row_count > 0 { + (rows + s.row_count, count + 1) + } else { + (rows, count) + } + }) + } + RowSelectionInner::Mask(mask) => { + let mask = mask.mask(); + let total_rows = mask.len(); + (total_rows, mask_run_count(mask)) + } + }; + + if effective_count == 0 { + return RowSelectionStrategy::Mask; + } + + if total_rows < effective_count.saturating_mul(threshold) { + RowSelectionStrategy::Mask + } else { + RowSelectionStrategy::Selectors + } + } + + #[cfg(test)] + fn selectors(&self) -> Vec<RowSelector> { + self.iter().copied().collect() + } + + fn into_selectors_vec(self) -> Vec<RowSelector> { Review 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? ########## parquet/src/arrow/arrow_reader/selection/mod.rs: ########## @@ -135,12 +149,297 @@ impl RowSelector { /// * Consecutive [`RowSelector`]s alternate skipping or selecting rows /// /// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData -#[derive(Debug, Clone, Default, Eq, PartialEq)] +#[derive(Default, Clone)] pub struct RowSelection { - selectors: Vec<RowSelector>, + inner: RowSelectionInner, +} + +/// Internal storage for [`RowSelection`]. +#[derive(Debug, Clone)] +pub(crate) enum RowSelectionInner { + Selectors(Vec<RowSelector>), + Mask(Box<MaskSelection>), +} + +impl Default for RowSelectionInner { + fn default() -> Self { + Self::Selectors(Vec::new()) + } +} + +impl std::fmt::Debug for RowSelection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.inner { + RowSelectionInner::Selectors(s) => f + .debug_struct("RowSelection") + .field("selectors", s) + .finish(), + RowSelectionInner::Mask(m) => f + .debug_struct("RowSelection") + .field("mask_len", &m.mask().len()) + .finish_non_exhaustive(), + } + } +} + +impl PartialEq for RowSelection { + fn eq(&self, other: &Self) -> bool { + match (&self.inner, &other.inner) { + (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b)) => a == b, + (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a.mask() == b.mask(), + (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors)) + | (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask)) => { + if selectors + .iter() + .try_fold(0usize, |acc, selector| acc.checked_add(selector.row_count)) + != Some(mask.mask().len()) + { + return false; + } + + let mut slices = mask.mask().set_slices().peekable(); + let mut cursor = 0usize; + + for selector in selectors { + let end = cursor + selector.row_count; + + if selector.skip { + if slices.peek().is_some_and(|(start, _)| *start < end) { + return false; + } + } else { + match slices.next() { + Some((start, slice_end)) if start == cursor && slice_end == end => {} + _ => return false, + } + } + + cursor = end; + } + + slices.next().is_none() + } + } + } +} + +impl Eq for RowSelection {} + +/// Borrowed iterator over the [`RowSelector`]s of a [`RowSelection`]. +#[derive(Debug)] +pub struct RowSelectionIter<'a>(std::slice::Iter<'a, RowSelector>); + +impl<'a> Iterator for RowSelectionIter<'a> { + type Item = &'a RowSelector; + + #[inline] + fn next(&mut self) -> Option<Self::Item> { Review 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`)? -- 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]
