haohuaijin commented on code in PR #10141: URL: https://github.com/apache/arrow-rs/pull/10141#discussion_r3643180870
########## parquet/src/arrow/arrow_reader/selection/boolean.rs: ########## @@ -0,0 +1,1308 @@ +// 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::{LoadedRowRanges, RowSelection, RowSelectionInner, RowSelector}; +use crate::errors::ParquetError; +use arrow_array::BooleanArray; +use arrow_buffer::bit_iterator::BitSliceIterator; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, Buffer}; +use std::cmp::Ordering; +use std::sync::{Arc, 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. +/// +/// `count` caches the popcount; `RowSelection::split_off` propagates it to +/// both halves so repeated `row_count()` calls do not rescan the bitmap. +#[derive(Debug)] +pub(crate) struct MaskSelection { + mask: BooleanBuffer, + selectors: OnceLock<Vec<RowSelector>>, + count: OnceLock<usize>, +} + +impl MaskSelection { + pub(super) fn new(mask: BooleanBuffer) -> Self { + Self { + mask, + selectors: OnceLock::new(), + count: OnceLock::new(), + } + } + + /// Create a selection whose selected-row count is already known. + pub(super) fn with_count(mask: BooleanBuffer, count: usize) -> Self { + debug_assert!(count <= mask.len()); + let cell = OnceLock::new(); + let _ = cell.set(count); + Self { + mask, + selectors: OnceLock::new(), + count: cell, + } + } + + pub(crate) fn mask(&self) -> &BooleanBuffer { + &self.mask + } + + pub(crate) fn into_mask(self) -> BooleanBuffer { + let Self { mask, .. } = self; + mask + } + + /// Number of selected rows, computed once and cached. + pub(super) fn count(&self) -> usize { + *self.count.get_or_init(|| self.mask.count_set_bits()) + } + + /// The cached selected-row count, if it has been computed. + pub(super) fn cached_count(&self) -> Option<usize> { + self.count.get().copied() + } + + 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 { + // Drop the selector cache but keep the cheap count cache. + Self { + mask: self.mask.clone(), + selectors: OnceLock::new(), + count: self.count.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. +/// +/// When page pruning produces sparse column data, `loaded_row_ranges` limits +/// each decoded chunk to rows whose pages are loaded for every projected leaf. +/// For example, two projected columns can have different page boundaries: +/// +/// ```text +/// Row ranges: [0, 4) [4, 6) [6, 8) [8, 10) [10, 12) +/// Selection mask: 1000 00 00 00 01 +/// Column A pages: loaded | missing [4, 8) | loaded [8, 12) +/// Column B pages: loaded [0, 6) | missing [6, 10) | loaded +/// LoadedRowRanges: [0, 4) [10, 12) +/// ``` +/// +/// The first chunk decodes `[0, 4)` with mask `1000`. The next chunk skips to +/// row 11 and decodes `[11, 12)` with mask `1`. The loaded ranges are decode +/// boundaries, not output batch boundaries: [`ParquetRecordBatchReader`] +/// accumulates both chunks and applies the combined mask `10001` once. +/// +/// [`ParquetRecordBatchReader`]: crate::arrow::arrow_reader::ParquetRecordBatchReader +#[derive(Debug)] +pub struct MaskCursor { + pub(super) mask: BooleanBuffer, + /// Current absolute offset into the selection + pub(super) position: usize, + /// Row ranges whose backing pages are loaded for every projected column. + pub(super) loaded_row_ranges: Option<Arc<LoadedRowRanges>>, +} + +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> { + if self.is_empty() { + return None; + } + + Some(self.next_mask_chunk_non_empty(batch_size)) + } + + /// Produces the next chunk for a non-empty, trailing-skip-free mask. + fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk { + debug_assert!(!self.is_empty()); + + let (initial_skip, chunk_rows, selected_rows, mask_start, end_position) = { + let mask = &self.mask; + 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; + } + debug_assert!( + cursor < mask.len(), + "ReadPlan must remove trailing skips from Mask selections" + ); + + 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; + + MaskChunk { + initial_skip, + chunk_rows, + selected_rows, + mask_start, + } + } + + /// Returns the next non-empty mask chunk without crossing an unloaded row range. + /// + /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing + /// skips before constructing this cursor. Callers therefore only invoke + /// this method for a non-empty mask that has another selected row. + pub(crate) fn next_chunk(&mut self, batch_size: usize) -> Result<MaskChunk, ParquetError> { + debug_assert!(batch_size > 0); + debug_assert!(!self.is_empty()); + + if self.loaded_row_ranges.is_none() { + return Ok(self.next_mask_chunk_non_empty(batch_size)); + } + + let start_position = self.position; + let mut cursor = start_position; + while cursor < self.mask.len() && !self.mask.value(cursor) { + cursor += 1; + } + + debug_assert!( + cursor < self.mask.len(), + "ReadPlan must remove trailing skips from Mask selections" + ); + + let loaded_range_end = self + .loaded_row_ranges + .as_ref() + .and_then(|ranges| ranges.end_containing(cursor)) + .ok_or_else(|| { + ParquetError::General(format!( + "Internal Error: selected row {cursor} has no loaded page range" + )) + })?; + + let mask_start = cursor; + let mut selected_rows = 0; + while cursor < loaded_range_end && cursor < self.mask.len() && selected_rows < batch_size { + if self.mask.value(cursor) { + selected_rows += 1; + } + cursor += 1; + } + + self.position = cursor; + Ok(MaskChunk { + initial_skip: mask_start - start_position, + chunk_rows: cursor - mask_start, + 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 +} + +/// Returns whether `mask` contains at least `min_runs` alternating set/unset runs. +/// +/// Stops as soon as the requested number of runs is found, avoiding a full scan +/// when callers only need to know whether a boundary has been crossed. +pub(super) fn mask_has_at_least_runs(mask: &BooleanBuffer, min_runs: usize) -> bool { + if min_runs == 0 { + return true; + } + + let total_rows = mask.len(); + if total_rows == 0 { + return false; + } + + let mut run_count = 0; + let mut last_end = 0; + for (start, end) in mask.set_slices() { + run_count += usize::from(start > last_end) + 1; + if run_count >= min_runs { + return true; + } + last_end = end; + } + + run_count + usize::from(last_end < total_rows) >= min_runs +} + +/// 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()); Review Comment: track in https://github.com/apache/arrow-rs/issues/10425 -- 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]
