haohuaijin commented on code in PR #10141:
URL: https://github.com/apache/arrow-rs/pull/10141#discussion_r3610572627
##########
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:
[496d7cb](https://github.com/apache/arrow-rs/pull/10141/commits/496d7cb62a2aff8064ac35c2bcdc5d682e78d291)
##########
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:
solve in
[496d7cb](https://github.com/apache/arrow-rs/pull/10141/commits/496d7cb62a2aff8064ac35c2bcdc5d682e78d291)
--
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]
