From 3b2ed79b0a377c7b7bfd2f22d7e70bb198704d89 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sat, 25 Jul 2026 11:56:08 +0800 Subject: [PATCH 1/4] refactor(parquet): split `arrow_reader/selection` into smaller modules `selection/mod.rs` had grown past 2100 lines and mixed the `RowSelection` type, the set algebra, the page range mapping and the cursor machinery. Split it by concern, with one module per selection backing: * `mod.rs` - `RowSelection` and its public API, dispatching on backing * `selector.rs` - run length backing: `RowSelector` and its primitives * `boolean.rs` - bitmap backing: `MaskSelection` and its primitives * `algebra.rs` - `and_then`, `intersection` and `union` * `ranges.rs` - page byte ranges and batch boundary expansion * `cursor.rs` - `RowSelectionCursor` and the policy / strategy types No functional change: bodies are moved verbatim and the public API is unchanged. --- .../arrow/arrow_reader/selection/algebra.rs | 828 +++++++++ .../arrow/arrow_reader/selection/boolean.rs | 606 +------ .../arrow/arrow_reader/selection/cursor.rs | 419 +++++ .../src/arrow/arrow_reader/selection/mod.rs | 1594 ++--------------- .../arrow/arrow_reader/selection/ranges.rs | 283 +++ .../arrow/arrow_reader/selection/selector.rs | 603 +++++++ 6 files changed, 2295 insertions(+), 2038 deletions(-) create mode 100644 parquet/src/arrow/arrow_reader/selection/algebra.rs create mode 100644 parquet/src/arrow/arrow_reader/selection/cursor.rs create mode 100644 parquet/src/arrow/arrow_reader/selection/ranges.rs create mode 100644 parquet/src/arrow/arrow_reader/selection/selector.rs diff --git a/parquet/src/arrow/arrow_reader/selection/algebra.rs b/parquet/src/arrow/arrow_reader/selection/algebra.rs new file mode 100644 index 000000000000..9c25afbd3b15 --- /dev/null +++ b/parquet/src/arrow/arrow_reader/selection/algebra.rs @@ -0,0 +1,828 @@ +// 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. + +//! Set algebra backing [`RowSelection::and_then`], [`RowSelection::intersection`] +//! and [`RowSelection::union`] +//! +//! Each operation has two implementations, picked by the backing of its +//! operands: a merge of the [`RowSelector`] runs, and a bitwise variant over +//! [`BooleanBuffer`] masks. + +use super::{MaskRunIter, RowSelection, RowSelectionInner, RowSelector}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use std::cmp::Ordering; +use std::iter::Peekable; + +/// Applies `second` to the rows selected by `first`, both selector-backed. +pub(super) fn and_then_row_selections( + first: &[RowSelector], + second: &[RowSelector], +) -> RowSelection { + let mut selectors = vec![]; + let mut first = first.iter().copied().peekable(); + let mut second = second.iter().copied().peekable(); + and_then_iter(&mut selectors, &mut first, &mut second); + RowSelection::from_selectors(selectors) +} + +/// Applies the mask `second` to the rows selected by the selector-backed `first`. +/// +/// The mask is streamed as [`RowSelector`] runs, so it is never materialized. +pub(super) fn and_then_selectors_with_mask( + first: &[RowSelector], + second: &BooleanBuffer, +) -> RowSelection { + let mut selectors = vec![]; + let mut first = first.iter().copied().peekable(); + let mut second = MaskRunIter::new(second).peekable(); + and_then_iter(&mut selectors, &mut first, &mut second); + RowSelection::from_selectors(selectors) +} + +fn and_then_iter( + selectors: &mut Vec, + first: &mut Peekable, + second: &mut Peekable, +) where + I: Iterator, + J: Iterator, +{ + let mut to_skip = 0; + while let Some(b) = second.peek_mut() { + let a = first + .peek_mut() + .expect("selection exceeds the number of selected rows"); + + if b.row_count == 0 { + second.next().unwrap(); + continue; + } + + if a.row_count == 0 { + first.next().unwrap(); + continue; + } + + if a.skip { + // Records were skipped when producing second + to_skip += a.row_count; + first.next().unwrap(); + continue; + } + + let skip = b.skip; + let to_process = a.row_count.min(b.row_count); + + a.row_count -= to_process; + b.row_count -= to_process; + + match skip { + true => to_skip += to_process, + false => { + if to_skip != 0 { + selectors.push(RowSelector::skip(to_skip)); + to_skip = 0; + } + selectors.push(RowSelector::select(to_process)) + } + } + } + + for v in first { + if v.row_count != 0 { + assert!( + v.skip, + "selection contains less than the number of selected rows" + ); + to_skip += v.row_count + } + } + + if to_skip != 0 { + selectors.push(RowSelector::skip(to_skip)); + } +} + +/// Combine two lists of `RowSelection` return the intersection of them +/// For example: +/// self: NNYYYYNNYYNYN +/// other: NYNNNNNNY +/// +/// returned: NNNNNNNNYYNYN +pub(super) fn intersect_row_selections( + left: &[RowSelector], + right: &[RowSelector], +) -> RowSelection { + let mut l_iter = left.iter().copied().peekable(); + let mut r_iter = right.iter().copied().peekable(); + + let iter = std::iter::from_fn(move || { + loop { + let l = l_iter.peek_mut(); + let r = r_iter.peek_mut(); + + match (l, r) { + (Some(a), _) if a.row_count == 0 => { + l_iter.next().unwrap(); + } + (_, Some(b)) if b.row_count == 0 => { + r_iter.next().unwrap(); + } + (Some(l), Some(r)) => { + return match (l.skip, r.skip) { + // Keep both ranges + (false, false) => { + if l.row_count < r.row_count { + r.row_count -= l.row_count; + l_iter.next() + } else { + l.row_count -= r.row_count; + r_iter.next() + } + } + // skip at least one + _ => { + if l.row_count < r.row_count { + let skip = l.row_count; + r.row_count -= l.row_count; + l_iter.next(); + Some(RowSelector::skip(skip)) + } else { + let skip = r.row_count; + l.row_count -= skip; + r_iter.next(); + Some(RowSelector::skip(skip)) + } + } + }; + } + (Some(_), None) => return l_iter.next(), + (None, Some(_)) => return r_iter.next(), + (None, None) => return None, + } + } + }); + + iter.collect() +} + +/// Combine two lists of `RowSelector` return the union of them +/// For example: +/// self: NNYYYYNNYYNYN +/// other: NYNNNNNNY +/// +/// returned: NYYYYYNNYYNYN +/// +/// This can be removed from here once RowSelection::union is in parquet::arrow +pub(super) fn union_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection { + let mut l_iter = left.iter().copied().peekable(); + let mut r_iter = right.iter().copied().peekable(); + + let iter = std::iter::from_fn(move || { + loop { + let l = l_iter.peek_mut(); + let r = r_iter.peek_mut(); + + match (l, r) { + (Some(a), _) if a.row_count == 0 => { + l_iter.next().unwrap(); + } + (_, Some(b)) if b.row_count == 0 => { + r_iter.next().unwrap(); + } + (Some(l), Some(r)) => { + return match (l.skip, r.skip) { + // Skip both ranges + (true, true) => { + if l.row_count < r.row_count { + let skip = l.row_count; + r.row_count -= l.row_count; + l_iter.next(); + Some(RowSelector::skip(skip)) + } else { + let skip = r.row_count; + l.row_count -= skip; + r_iter.next(); + Some(RowSelector::skip(skip)) + } + } + // Keep rows from left + (false, true) => { + if l.row_count < r.row_count { + r.row_count -= l.row_count; + l_iter.next() + } else { + let r_row_count = r.row_count; + l.row_count -= r_row_count; + r_iter.next(); + Some(RowSelector::select(r_row_count)) + } + } + // Keep rows from right + (true, false) => { + if l.row_count < r.row_count { + let l_row_count = l.row_count; + r.row_count -= l_row_count; + l_iter.next(); + Some(RowSelector::select(l_row_count)) + } else { + l.row_count -= r.row_count; + r_iter.next() + } + } + // Keep at least one + _ => { + if l.row_count < r.row_count { + r.row_count -= l.row_count; + l_iter.next() + } else { + l.row_count -= r.row_count; + r_iter.next() + } + } + }; + } + (Some(_), None) => return l_iter.next(), + (None, Some(_)) => return r_iter.next(), + (None, None) => return None, + } + } + }); + + iter.collect() +} + +/// 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(mask: &BooleanBuffer, other: I) -> BooleanBuffer +where + I: IntoIterator, +{ + 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() +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::BooleanArray; + use rand::{Rng, rng}; + + #[test] + fn test_and() { + let mut a = RowSelection::from(vec![ + RowSelector::skip(12), + RowSelector::select(23), + RowSelector::skip(3), + RowSelector::select(5), + ]); + + let b = RowSelection::from(vec![ + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(15), + RowSelector::skip(4), + ]); + + let mut expected = RowSelection::from(vec![ + RowSelector::skip(12), + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(14), + RowSelector::skip(3), + RowSelector::select(1), + RowSelector::skip(4), + ]); + + assert_eq!(a.and_then(&b), expected); + + a.split_off(7); + expected.split_off(7); + assert_eq!(a.and_then(&b), expected); + + let a = RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(3)]); + + let b = RowSelection::from(vec![ + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + ]); + + assert_eq!( + a.and_then(&b).selectors(), + vec![ + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(4) + ] + ); + } + + #[test] + #[should_panic(expected = "selection exceeds the number of selected rows")] + fn test_and_longer() { + let a = RowSelection::from(vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(3), + RowSelector::skip(33), + ]); + let b = RowSelection::from(vec![RowSelector::select(36)]); + a.and_then(&b); + } + + #[test] + #[should_panic(expected = "selection contains less than the number of selected rows")] + fn test_and_shorter() { + let a = RowSelection::from(vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(3), + RowSelector::skip(33), + ]); + let b = RowSelection::from(vec![RowSelector::select(3)]); + a.and_then(&b); + } + + #[test] + fn test_intersect_row_selection_and_combine() { + // a size equal b size + let a = vec![ + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(1), + ]; + let b = vec![ + RowSelector::select(8), + RowSelector::skip(1), + RowSelector::select(1), + ]; + + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![ + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(1), + ], + ); + + // a size larger than b size + let a = vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(3), + RowSelector::skip(33), + ]; + let b = vec![RowSelector::select(36), RowSelector::skip(36)]; + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![RowSelector::select(3), RowSelector::skip(69)] + ); + + // a size less than b size + let a = vec![RowSelector::select(3), RowSelector::skip(7)]; + let b = vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + ]; + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![RowSelector::select(2), RowSelector::skip(8)] + ); + + let a = vec![RowSelector::select(3), RowSelector::skip(7)]; + let b = vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + ]; + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![RowSelector::select(2), RowSelector::skip(8)] + ); + } + + #[test] + fn test_and_fuzz() { + let mut rand = rng(); + for _ in 0..100 { + let a_len = rand.random_range(10..100); + let a_bools: Vec<_> = (0..a_len).map(|_| rand.random_bool(0.2)).collect(); + let a = RowSelection::from_filters(&[BooleanArray::from(a_bools.clone())]); + + let b_len: usize = a_bools.iter().map(|x| *x as usize).sum(); + let b_bools: Vec<_> = (0..b_len).map(|_| rand.random_bool(0.8)).collect(); + let b = RowSelection::from_filters(&[BooleanArray::from(b_bools.clone())]); + + let mut expected_bools = vec![false; a_len]; + + let mut iter_b = b_bools.iter(); + for (idx, b) in a_bools.iter().enumerate() { + if *b && *iter_b.next().unwrap() { + expected_bools[idx] = true; + } + } + + let expected = RowSelection::from_filters(&[BooleanArray::from(expected_bools)]); + + let total_rows: usize = expected.selectors().iter().map(|s| s.row_count).sum(); + assert_eq!(a_len, total_rows); + + assert_eq!(a.and_then(&b), expected); + } + } + + #[test] + fn test_intersection() { + let selection = RowSelection::from(vec![RowSelector::select(1048576)]); + let result = selection.intersection(&selection); + assert_eq!(result, selection); + + let a = RowSelection::from(vec![ + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(20), + ]); + + let b = RowSelection::from(vec![ + RowSelector::skip(20), + RowSelector::select(20), + RowSelector::skip(10), + ]); + + let result = a.intersection(&b); + assert_eq!( + result.selectors(), + vec![ + RowSelector::skip(30), + RowSelector::select(10), + RowSelector::skip(10) + ] + ); + } + + #[test] + fn test_union() { + let selection = RowSelection::from(vec![RowSelector::select(1048576)]); + let result = selection.union(&selection); + assert_eq!(result, selection); + + // NYNYY + let a = RowSelection::from(vec![ + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(20), + ]); + + // NNYYNYN + let b = RowSelection::from(vec![ + RowSelector::skip(20), + RowSelector::select(20), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + ]); + + let result = a.union(&b); + + // NYYYYYN + assert_eq!( + result.iter().copied().collect::>(), + vec![ + RowSelector::skip(10), + RowSelector::select(50), + RowSelector::skip(10), + ] + ); + } + + #[test] + fn test_mask_and_then_preserves_backing() { + let outer_bits = vec![false, true, true, false, true, false, true]; + let inner_bits = vec![true, false, true, false]; + let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone())); + let inner = RowSelection::from_filters(&[BooleanArray::from(inner_bits.clone())]); + + let result = outer_mask.and_then(&inner); + assert!(result.as_mask().is_some()); + + let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]); + let expected = outer_selectors.and_then(&inner); + assert_eq!(result, expected); + + let result_mask = result.as_mask().unwrap(); + let actual_bits: Vec<_> = (0..result_mask.len()) + .map(|i| result_mask.value(i)) + .collect(); + assert_eq!( + actual_bits, + vec![false, true, false, false, true, false, false] + ); + } + + #[test] + fn test_mask_and_then_mask_preserves_backing() { + let outer_bits = vec![false, true, true, false, true, false, true, true]; + let inner_bits = vec![false, true, false, true, false]; + let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone())); + let inner_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(inner_bits)); + + let result = outer_mask.and_then(&inner_mask); + assert!(result.as_mask().is_some()); + + let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]); + let inner_selectors = RowSelection::from_filters(&[BooleanArray::from(vec![ + false, true, false, true, false, + ])]); + assert_eq!(result, outer_selectors.and_then(&inner_selectors)); + + let result_mask = result.as_mask().unwrap(); + let actual_bits: Vec<_> = (0..result_mask.len()) + .map(|i| result_mask.value(i)) + .collect(); + assert_eq!( + actual_bits, + vec![false, false, true, false, false, false, true, false] + ); + } + + #[test] + fn test_selector_and_then_mask() { + let outer = + RowSelection::from_filters(&[BooleanArray::from(vec![false, true, true, false, true])]); + let inner = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])); + + let result = outer.and_then(&inner); + assert!(result.as_mask().is_none()); + assert_eq!( + result, + RowSelection::from_filters(&[BooleanArray::from(vec![ + false, true, false, false, true, + ])]) + ); + } + + #[test] + fn test_mask_and_then_none_selected_returns_all_unset() { + let outer = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, true, true, false, true, + ])); + let inner = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![false, false, false])); + + let result = outer.and_then(&inner); + let mask = result.as_mask().unwrap(); + assert_eq!(mask.len(), 5); + assert_eq!(mask.count_set_bits(), 0); + } + + #[test] + fn test_mask_intersection_uses_bitwise() { + let a_bits = vec![true, true, false, true, false, true]; + let b_bits = vec![true, false, true, true, true, false]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())); + + let r = a.intersection(&b); + assert!(r.as_mask().is_some()); + + let expected: Vec = a_bits.iter().zip(&b_bits).map(|(x, y)| *x && *y).collect(); + let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]); + assert_eq!(r, expected_sel); + } + + #[test] + fn test_mask_union_uses_bitwise() { + let a_bits = vec![true, false, false, true, false, false]; + let b_bits = vec![false, true, false, false, true, false]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())); + + let r = a.union(&b); + assert!(r.as_mask().is_some()); + + let expected: Vec = a_bits.iter().zip(&b_bits).map(|(x, y)| *x || *y).collect(); + let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]); + assert_eq!(r, expected_sel); + } + + #[test] + fn test_mixed_mask_selector_intersection_and_union() { + let mask_bits = vec![true, false, true, false, true, false]; + let selector_bits = vec![false, true, true, false, false, true]; + let mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(mask_bits.clone())); + let selectors = RowSelection::from_filters(&[BooleanArray::from(selector_bits.clone())]); + + let intersection_bits: Vec<_> = mask_bits + .iter() + .zip(&selector_bits) + .map(|(x, y)| *x && *y) + .collect(); + let expected_intersection = + RowSelection::from_filters(&[BooleanArray::from(intersection_bits)]); + assert_eq!(mask.intersection(&selectors), expected_intersection); + assert_eq!(selectors.intersection(&mask), expected_intersection); + + let union_bits: Vec<_> = mask_bits + .iter() + .zip(&selector_bits) + .map(|(x, y)| *x || *y) + .collect(); + let expected_union = RowSelection::from_filters(&[BooleanArray::from(union_bits)]); + assert_eq!(mask.union(&selectors), expected_union); + assert_eq!(selectors.union(&mask), expected_union); + } + + #[test] + fn test_mask_intersection_uneven_passes_tail_through() { + let a_bits = vec![true, true, true, true, true]; + let b_bits = vec![true, false, true]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits)); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits)); + + let r = a.intersection(&b); + let r_mask = r.as_mask().unwrap(); + assert_eq!(r_mask.len(), 5); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, false, true, true, true]); + + // Swapped operands: the right side is longer and its tail passes through. + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + true, true, true, false, true, + ])); + let r = a.intersection(&b); + let r_mask = r.as_mask().unwrap(); + assert_eq!(r_mask.len(), 5); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, false, true, false, true]); + } + + #[test] + fn test_mask_union_uneven_passes_tail_through() { + let a_bits = vec![true, false, true]; + let b_bits = vec![false, true, false, true, false]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits)); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits)); + + let r = a.union(&b); + let r_mask = r.as_mask().unwrap(); + assert_eq!(r_mask.len(), 5); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, true, true, true, false]); + + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, true, false, false, true, + ])); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, false])); + let r = a.union(&b); + let r_mask = r.as_mask().unwrap(); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, true, false, false, true]); + } +} diff --git a/parquet/src/arrow/arrow_reader/selection/boolean.rs b/parquet/src/arrow/arrow_reader/selection/boolean.rs index 2953359f33d2..2ac879eba9dc 100644 --- a/parquet/src/arrow/arrow_reader/selection/boolean.rs +++ b/parquet/src/arrow/arrow_reader/selection/boolean.rs @@ -15,22 +15,31 @@ // specific language governing permissions and limitations // under the License. -use super::{LoadedRowRanges, RowSelection, RowSelectionInner, RowSelector}; -use crate::errors::ParquetError; -use arrow_array::BooleanArray; +//! The bitmap backed representation of a [`RowSelection`] and the primitives +//! operating on it: conversion to and from the run length ([`RowSelector`]) +//! form, and the transforms backing `split_off`, `trim`, `offset` and `limit`. +//! +//! The bitwise set algebra lives in the `algebra` module. +//! +//! [`RowSelection`]: crate::arrow::arrow_reader::RowSelection + +use super::RowSelector; use arrow_buffer::bit_iterator::BitSliceIterator; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, Buffer}; -use std::cmp::Ordering; -use std::sync::{Arc, OnceLock}; +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. +/// `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 +/// `count` caches the popcount; [`RowSelection::split_off`] propagates it to /// both halves so repeated `row_count()` calls do not rescan the bitmap. +/// +/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection +/// [`RowSelection::iter`]: crate::arrow::arrow_reader::RowSelection::iter +/// [`RowSelection::split_off`]: crate::arrow::arrow_reader::RowSelection::split_off #[derive(Debug)] pub(crate) struct MaskSelection { mask: BooleanBuffer, @@ -107,6 +116,9 @@ impl Clone for MaskSelection { /// for run in MaskRunIter::new(mask) { ... } /// } /// ``` +/// +/// [`RowSelection::iter`]: crate::arrow::arrow_reader::RowSelection::iter +/// [`RowSelection::as_mask`]: crate::arrow::arrow_reader::RowSelection::as_mask #[derive(Debug)] pub struct MaskRunIter<'a> { slices: BitSliceIterator<'a>, @@ -166,178 +178,6 @@ impl Iterator for MaskRunIter<'_> { } } -/// 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>, -} - -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 { - 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 { - 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 { - 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 { let total_rows = mask.len(); @@ -386,133 +226,6 @@ pub(super) fn mask_has_at_least_runs(mask: &BooleanBuffer, min_runs: usize) -> b 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()); - 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(mask: &BooleanBuffer, other: I) -> BooleanBuffer -where - I: IntoIterator, -{ - 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( @@ -624,6 +337,7 @@ pub(super) fn boolean_mask_from_selectors(selectors: &[RowSelector]) -> BooleanB #[cfg(test)] mod tests { use super::*; + use crate::arrow::arrow_reader::selection::{RowSelection, RowSelectionInner}; use arrow_array::BooleanArray; use rand::{Rng, rng}; @@ -823,72 +537,6 @@ mod tests { } } - #[test] - fn test_mask_and_then_preserves_backing() { - let outer_bits = vec![false, true, true, false, true, false, true]; - let inner_bits = vec![true, false, true, false]; - let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone())); - let inner = RowSelection::from_filters(&[BooleanArray::from(inner_bits.clone())]); - - let result = outer_mask.and_then(&inner); - assert!(result.as_mask().is_some()); - - let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]); - let expected = outer_selectors.and_then(&inner); - assert_eq!(result, expected); - - let result_mask = result.as_mask().unwrap(); - let actual_bits: Vec<_> = (0..result_mask.len()) - .map(|i| result_mask.value(i)) - .collect(); - assert_eq!( - actual_bits, - vec![false, true, false, false, true, false, false] - ); - } - - #[test] - fn test_mask_and_then_mask_preserves_backing() { - let outer_bits = vec![false, true, true, false, true, false, true, true]; - let inner_bits = vec![false, true, false, true, false]; - let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone())); - let inner_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(inner_bits)); - - let result = outer_mask.and_then(&inner_mask); - assert!(result.as_mask().is_some()); - - let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]); - let inner_selectors = RowSelection::from_filters(&[BooleanArray::from(vec![ - false, true, false, true, false, - ])]); - assert_eq!(result, outer_selectors.and_then(&inner_selectors)); - - let result_mask = result.as_mask().unwrap(); - let actual_bits: Vec<_> = (0..result_mask.len()) - .map(|i| result_mask.value(i)) - .collect(); - assert_eq!( - actual_bits, - vec![false, false, true, false, false, false, true, false] - ); - } - - #[test] - fn test_selector_and_then_mask() { - let outer = - RowSelection::from_filters(&[BooleanArray::from(vec![false, true, true, false, true])]); - let inner = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])); - - let result = outer.and_then(&inner); - assert!(result.as_mask().is_none()); - assert_eq!( - result, - RowSelection::from_filters(&[BooleanArray::from(vec![ - false, true, false, false, true, - ])]) - ); - } - #[test] fn test_mask_offset_past_end_preserves_empty_mask_backing() { let selection = @@ -914,161 +562,6 @@ mod tests { assert_eq!(actual_bits, vec![false, true, false, true]); } - #[test] - fn test_mask_intersection_uses_bitwise() { - let a_bits = vec![true, true, false, true, false, true]; - let b_bits = vec![true, false, true, true, true, false]; - let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())); - let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())); - - let r = a.intersection(&b); - assert!(r.as_mask().is_some()); - - let expected: Vec = a_bits.iter().zip(&b_bits).map(|(x, y)| *x && *y).collect(); - let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]); - assert_eq!(r, expected_sel); - } - - #[test] - fn test_mask_union_uses_bitwise() { - let a_bits = vec![true, false, false, true, false, false]; - let b_bits = vec![false, true, false, false, true, false]; - let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())); - let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())); - - let r = a.union(&b); - assert!(r.as_mask().is_some()); - - let expected: Vec = a_bits.iter().zip(&b_bits).map(|(x, y)| *x || *y).collect(); - let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]); - assert_eq!(r, expected_sel); - } - - #[test] - fn test_mixed_mask_selector_intersection_and_union() { - let mask_bits = vec![true, false, true, false, true, false]; - let selector_bits = vec![false, true, true, false, false, true]; - let mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(mask_bits.clone())); - let selectors = RowSelection::from_filters(&[BooleanArray::from(selector_bits.clone())]); - - let intersection_bits: Vec<_> = mask_bits - .iter() - .zip(&selector_bits) - .map(|(x, y)| *x && *y) - .collect(); - let expected_intersection = - RowSelection::from_filters(&[BooleanArray::from(intersection_bits)]); - assert_eq!(mask.intersection(&selectors), expected_intersection); - assert_eq!(selectors.intersection(&mask), expected_intersection); - - let union_bits: Vec<_> = mask_bits - .iter() - .zip(&selector_bits) - .map(|(x, y)| *x || *y) - .collect(); - let expected_union = RowSelection::from_filters(&[BooleanArray::from(union_bits)]); - assert_eq!(mask.union(&selectors), expected_union); - assert_eq!(selectors.union(&mask), expected_union); - } - - #[test] - fn test_mask_intersection_uneven_passes_tail_through() { - let a_bits = vec![true, true, true, true, true]; - let b_bits = vec![true, false, true]; - let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits)); - let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits)); - - let r = a.intersection(&b); - let r_mask = r.as_mask().unwrap(); - assert_eq!(r_mask.len(), 5); - let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); - assert_eq!(bits, vec![true, false, true, true, true]); - - // Swapped operands: the right side is longer and its tail passes through. - let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])); - let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ - true, true, true, false, true, - ])); - let r = a.intersection(&b); - let r_mask = r.as_mask().unwrap(); - assert_eq!(r_mask.len(), 5); - let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); - assert_eq!(bits, vec![true, false, true, false, true]); - } - - #[test] - fn test_mask_and_then_none_selected_returns_all_unset() { - let outer = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ - false, true, true, false, true, - ])); - let inner = - RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![false, false, false])); - - let result = outer.and_then(&inner); - let mask = result.as_mask().unwrap(); - assert_eq!(mask.len(), 5); - assert_eq!(mask.count_set_bits(), 0); - } - - #[test] - fn test_mixed_backing_equality_mismatches() { - let mask = - RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true, true])); - - // Total row counts differ - let longer = RowSelection::from(vec![ - RowSelector::select(1), - RowSelector::skip(1), - RowSelector::select(2), - RowSelector::skip(1), - ]); - assert_ne!(mask, longer); - assert_ne!(longer, mask); - - // A selected bit falls inside a skip run - let skip_overlap = RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]); - assert_ne!(mask, skip_overlap); - - // Select run boundaries do not line up - let misaligned = RowSelection::from(vec![ - RowSelector::select(2), - RowSelector::skip(1), - RowSelector::select(1), - ]); - assert_ne!(mask, misaligned); - - let equal = RowSelection::from(vec![ - RowSelector::select(1), - RowSelector::skip(1), - RowSelector::select(2), - ]); - assert_eq!(mask, equal); - assert_eq!(equal, mask); - } - - #[test] - fn test_mask_union_uneven_passes_tail_through() { - let a_bits = vec![true, false, true]; - let b_bits = vec![false, true, false, true, false]; - let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits)); - let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits)); - - let r = a.union(&b); - let r_mask = r.as_mask().unwrap(); - assert_eq!(r_mask.len(), 5); - let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); - assert_eq!(bits, vec![true, true, true, true, false]); - - let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ - false, true, false, false, true, - ])); - let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, false])); - let r = a.union(&b); - let r_mask = r.as_mask().unwrap(); - let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); - assert_eq!(bits, vec![true, true, false, false, true]); - } - #[test] fn test_mask_split_off_preserves_backing() { let bits: Vec = (0..40).map(|i| i % 3 == 0).collect(); @@ -1132,61 +625,6 @@ mod tests { assert_eq!(trimmed.as_mask().unwrap().len(), 0); } - #[test] - fn test_from_iter_all_mask_preserves_mask_backing() { - let a_bits = vec![true, false, true, true]; - let b_bits = vec![false, true, false]; - let c_bits = vec![true, true, false, false, true]; - - let parts = vec![ - RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())), - RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())), - RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())), - ]; - let collected: RowSelection = parts.into_iter().collect(); - - let combined = a_bits - .iter() - .chain(b_bits.iter()) - .chain(c_bits.iter()) - .copied() - .collect::>(); - let expected = RowSelection::from_filters(&[BooleanArray::from(combined)]); - - assert!(collected.as_mask().is_some()); - assert_eq!(collected, expected); - } - - #[test] - fn test_from_iter_mixed_backing_falls_back_to_selectors() { - let a_bits = vec![true, false, true]; - let b_selectors = vec![RowSelector::skip(2), RowSelector::select(3)]; - let c_bits = vec![false, true]; - - let parts = vec![ - RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())), - RowSelection::from(b_selectors), - RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())), - ]; - let collected: RowSelection = parts.into_iter().collect(); - - assert!(collected.as_mask().is_none()); - - let combined_bits = vec![ - true, false, true, false, false, true, true, true, false, true, - ]; - let expected = RowSelection::from_filters(&[BooleanArray::from(combined_bits)]); - assert_eq!(collected, expected); - } - - #[test] - fn test_from_iter_empty_yields_empty_selection() { - let collected: RowSelection = std::iter::empty::().collect(); - assert_eq!(collected, RowSelection::default()); - assert!(collected.as_mask().is_some()); - assert_eq!(collected.as_mask().unwrap().len(), 0); - } - #[test] fn test_boolean_mask_from_selectors_fuzz_equivalence() { let mut rand = rng(); diff --git a/parquet/src/arrow/arrow_reader/selection/cursor.rs b/parquet/src/arrow/arrow_reader/selection/cursor.rs new file mode 100644 index 000000000000..dcb490746c9e --- /dev/null +++ b/parquet/src/arrow/arrow_reader/selection/cursor.rs @@ -0,0 +1,419 @@ +// 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. + +//! Execution time iteration over a [`RowSelection`]. +//! +//! A [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) resolves a +//! [`RowSelectionPolicy`] into a [`RowSelectionStrategy`] and builds the +//! matching [`RowSelectionCursor`], which keeps the per-reader position while +//! the selection itself stays immutable. + +use super::boolean::boolean_mask_from_selectors; +use super::{RowSelection, RowSelector}; +use crate::errors::ParquetError; +use arrow_array::BooleanArray; +use arrow_buffer::BooleanBuffer; +use std::collections::VecDeque; +use std::ops::Range; +use std::sync::Arc; + +/// Policy for picking a strategy to materialize [`RowSelection`] during execution. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RowSelectionPolicy { + /// Use a queue of [`RowSelector`] values + Selectors, + /// Use a boolean mask to materialize the selection + Mask, + /// Choose between [`Self::Mask`] and [`Self::Selectors`] based on selector density + Auto { + /// Average selector length below which masks are preferred + threshold: usize, + }, +} + +impl Default for RowSelectionPolicy { + fn default() -> Self { + Self::Auto { threshold: 32 } + } +} + +/// Fully resolved strategy for materializing [`RowSelection`] during execution. +/// +/// This is determined by [`RowSelectionPolicy`], including selector density for +/// [`RowSelectionPolicy::Auto`]. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum RowSelectionStrategy { + /// Use a queue of [`RowSelector`] values + Selectors, + /// Use a boolean mask to materialize the selection + Mask, +} + +/// Cursor for iterating a [`RowSelection`] during execution within a +/// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan). +/// +/// This keeps per-reader state such as the current position and delegates the +/// actual storage strategy to the internal `RowSelectionInner`. +#[derive(Debug)] +pub enum RowSelectionCursor { + /// Reading all rows + All, + /// Use a bitmask to back the selection (dense selections) + Mask(MaskCursor), + /// Use a queue of selectors to back the selection (sparse selections) + Selectors(SelectorsCursor), +} + +impl RowSelectionCursor { + /// Create a [`MaskCursor`] cursor backed by a bitmask, from an existing set of selectors + pub(crate) fn new_mask_from_selectors( + selectors: Vec, + loaded_row_ranges: Option>, + ) -> Self { + debug_assert!( + selectors + .last() + .map(|selector| !selector.skip) + .unwrap_or(true), + "Mask selectors must not end with a skip" + ); + Self::Mask(MaskCursor { + mask: boolean_mask_from_selectors(&selectors), + position: 0, + loaded_row_ranges, + }) + } + + /// Create a [`MaskCursor`] cursor backed by an existing bitmask. + pub(crate) fn new_mask_from_buffer( + mask: BooleanBuffer, + loaded_row_ranges: Option>, + ) -> Self { + debug_assert!( + mask.is_empty() || mask.value(mask.len() - 1), + "Mask selections must not end with a skip" + ); + Self::Mask(MaskCursor { + mask, + position: 0, + loaded_row_ranges, + }) + } + + /// Create a [`RowSelectionCursor::Selectors`] from the provided selectors + pub(crate) fn new_selectors(selectors: Vec) -> Self { + Self::Selectors(SelectorsCursor { + selectors: selectors.into(), + position: 0, + }) + } + + /// Create a cursor that selects all rows + pub(crate) fn new_all() -> Self { + Self::All + } +} + +/// Cursor for iterating a selector-backed [`RowSelection`] +/// +/// This is best for sparse selections where large contiguous +/// blocks of rows are selected or skipped. +#[derive(Debug)] +pub struct SelectorsCursor { + selectors: VecDeque, + /// Current absolute offset into the selection + position: usize, +} + +impl SelectorsCursor { + /// Returns `true` when no further rows remain + pub fn is_empty(&self) -> bool { + self.selectors.is_empty() + } + + pub(crate) fn selectors_mut(&mut self) -> &mut VecDeque { + &mut self.selectors + } + + /// Return the next [`RowSelector`] + pub(crate) fn next_selector(&mut self) -> RowSelector { + let selector = self.selectors.pop_front().unwrap(); + self.position += selector.row_count; + selector + } + + /// Return a selector to the front, rewinding the position + pub(crate) fn return_selector(&mut self, selector: RowSelector) { + self.position = self.position.saturating_sub(selector.row_count); + self.selectors.push_front(selector); + } +} + +/// 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 { + mask: BooleanBuffer, + /// Current absolute offset into the selection + position: usize, + /// Row ranges whose backing pages are loaded for every projected column. + loaded_row_ranges: Option>, +} + +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 { + 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 { + 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 { + 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, +} + +/// Row ranges whose backing pages are loaded for every projected column. +#[derive(Clone, Debug)] +pub(crate) struct LoadedRowRanges(Vec>); + +impl LoadedRowRanges { + pub(crate) fn from_selection(selection: RowSelection) -> Self { + let selectors: Vec = selection.into(); + let mut position = 0; + let ranges = selectors + .into_iter() + .filter_map(|selector| { + let start = position; + position += selector.row_count; + (!selector.skip).then_some(start..position) + }) + .collect(); + Self(ranges) + } + + fn end_containing(&self, row: usize) -> Option { + let idx = self.0.partition_point(|range| range.end <= row); + self.0 + .get(idx) + .filter(|range| range.start <= row) + .map(|range| range.end) + } + + #[cfg(test)] + pub(crate) fn ranges(&self) -> &[Range] { + &self.0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() { + let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges( + std::iter::once(0..5), + 10, + )); + let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors( + vec![RowSelector::select(1)], + Some(loaded.into()), + ) else { + unreachable!() + }; + + let chunk = cursor.next_chunk(10).unwrap(); + assert_eq!(chunk.chunk_rows, 1); + assert!(cursor.is_empty()); + } + + #[test] + fn test_next_mask_chunk_until_cursor_is_empty() { + let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors( + vec![ + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + ], + None, + ) else { + unreachable!() + }; + + let first = cursor.next_mask_chunk(2).unwrap(); + assert_eq!(first.initial_skip, 2); + assert_eq!(first.chunk_rows, 2); + assert_eq!(first.selected_rows, 2); + + let second = cursor.next_mask_chunk(2).unwrap(); + assert_eq!(second.initial_skip, 1); + assert_eq!(second.chunk_rows, 1); + assert_eq!(second.selected_rows, 1); + + assert!(cursor.next_mask_chunk(2).is_none()); + } +} diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index 735340420ed5..8e4948dd00c9 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -15,83 +15,50 @@ // specific language governing permissions and limitations // under the License. +//! Logic for selecting which rows to read: [`RowSelection`] and [`RowSelector`] +//! +//! This module holds [`RowSelection`] and its public API, which dispatches to +//! one of the two backings depending on how the selection is stored: +//! +//! * `selector`: the run length backing, [`RowSelector`] and its primitives +//! * `boolean`: the bitmap backing, `MaskSelection` and its primitives +//! +//! The remaining modules hold the operations that are common to both: +//! +//! * `algebra`: `and_then`, `intersection` and `union` +//! * `ranges`: mapping a [`RowSelection`] onto page and batch ranges +//! * `cursor`: iterating a [`RowSelection`] while reading + use crate::file::page_index::offset_index::PageLocation; use arrow_array::{Array, BooleanArray}; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; use arrow_select::filter::SlicesIterator; -use std::cmp::Ordering; use std::collections::VecDeque; use std::ops::Range; -use std::sync::Arc; +mod algebra; mod boolean; +mod cursor; +mod ranges; +mod selector; + +use algebra::{ + and_then_mask, and_then_row_selections, and_then_selectors_with_mask, intersect_masks, + intersect_row_selections, union_masks, union_row_selections, +}; pub use boolean::MaskRunIter; -pub(crate) use boolean::{MaskCursor, mask_to_selectors}; +pub(crate) use boolean::mask_to_selectors; use boolean::{ - MaskSelection, and_then_mask, boolean_mask_from_selectors, intersect_masks, limit_mask, - mask_has_at_least_runs, offset_mask, split_off_mask, trim_mask, union_masks, + MaskSelection, limit_mask, mask_has_at_least_runs, offset_mask, split_off_mask, trim_mask, +}; +pub(crate) use cursor::{LoadedRowRanges, MaskCursor, RowSelectionStrategy}; +pub use cursor::{RowSelectionCursor, RowSelectionPolicy}; +use ranges::{expand_to_batch_boundaries_from_selectors, scan_ranges_from_selectors}; +pub use selector::{RowSelectionIter, RowSelector}; +use selector::{ + combine_selectors, limit_selectors, offset_selectors, selectors_from_consecutive_ranges, + split_off_selectors, trim_selectors, }; - -/// Policy for picking a strategy to materialize [`RowSelection`] during execution. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum RowSelectionPolicy { - /// Use a queue of [`RowSelector`] values - Selectors, - /// Use a boolean mask to materialize the selection - Mask, - /// Choose between [`Self::Mask`] and [`Self::Selectors`] based on selector density - Auto { - /// Average selector length below which masks are preferred - threshold: usize, - }, -} - -impl Default for RowSelectionPolicy { - fn default() -> Self { - Self::Auto { threshold: 32 } - } -} - -/// Fully resolved strategy for materializing [`RowSelection`] during execution. -/// -/// This is determined by [`RowSelectionPolicy`], including selector density for -/// [`RowSelectionPolicy::Auto`]. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) enum RowSelectionStrategy { - /// Use a queue of [`RowSelector`] values - Selectors, - /// Use a boolean mask to materialize the selection - Mask, -} - -/// [`RowSelection`] is a collection of [`RowSelector`] used to skip rows when -/// scanning a parquet file -#[derive(Debug, Clone, Copy, Eq, PartialEq)] -pub struct RowSelector { - /// The number of rows - pub row_count: usize, - - /// If true, skip `row_count` rows - pub skip: bool, -} - -impl RowSelector { - /// Select `row_count` rows - pub fn select(row_count: usize) -> Self { - Self { - row_count, - skip: false, - } - } - - /// Skip `row_count` rows - pub fn skip(row_count: usize) -> Self { - Self { - row_count, - skip: true, - } - } -} /// [`RowSelection`] represents selecting a subset of rows /// when scanning a parquet file. @@ -226,158 +193,6 @@ impl PartialEq for RowSelection { 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.0.next() - } - - #[inline] - fn size_hint(&self) -> (usize, Option) { - self.0.size_hint() - } - - #[inline] - fn count(self) -> usize { - self.0.count() - } - - #[inline] - fn nth(&mut self, n: usize) -> Option { - self.0.nth(n) - } - - #[inline] - fn last(self) -> Option { - self.0.last() - } - - #[inline] - fn fold(self, init: B, f: F) -> B - where - F: FnMut(B, Self::Item) -> B, - { - self.0.fold(init, f) - } -} - -impl ExactSizeIterator for RowSelectionIter<'_> {} - -// once it returns None, it will continue returning None -impl std::iter::FusedIterator for RowSelectionIter<'_> {} - -#[inline] -fn scan_ranges_from_selectors(selectors: I, page_locations: &[PageLocation]) -> Vec> -where - I: IntoIterator, -{ - let mut ranges: Vec> = 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( - selectors: I, - batch_size: usize, - total_rows: usize, -) -> RowSelection -where - I: IntoIterator, -{ - 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> = 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 { /// Not `pub`: unlike `From>`, this performs no /// validation/normalization of the selectors (e.g. combining adjacent @@ -530,33 +345,7 @@ impl RowSelection { ranges: I, total_rows: usize, ) -> Self { - let mut selectors: Vec = Vec::with_capacity(ranges.size_hint().0); - let mut last_end = 0; - for range in ranges { - let len = range.end - range.start; - if len == 0 { - continue; - } - - match range.start.cmp(&last_end) { - Ordering::Equal => match selectors.last_mut() { - Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(), - None => selectors.push(RowSelector::select(len)), - }, - Ordering::Greater => { - selectors.push(RowSelector::skip(range.start - last_end)); - selectors.push(RowSelector::select(len)) - } - Ordering::Less => panic!("out of order"), - } - last_end = range.end; - } - - if last_end != total_rows { - selectors.push(RowSelector::skip(total_rows - last_end)) - } - - Self::from_selectors(selectors) + Self::from_selectors(selectors_from_consecutive_ranges(ranges, total_rows)) } /// Given an offset index, return the byte ranges for all data pages selected by `self` @@ -634,39 +423,7 @@ impl RowSelection { return Self::from_mask_selection(head); } - let selectors = self.selectors_mut(); - let mut total_count = 0; - - // Find the index where the selector exceeds the row count - let find = selectors.iter().position(|selector| { - total_count += selector.row_count; - total_count > row_count - }); - - let split_idx = match find { - Some(idx) => idx, - None => { - let drained = std::mem::take(selectors); - return Self::from_selectors(drained); - } - }; - - let mut remaining = selectors.split_off(split_idx); - - // Always present as `split_idx < selectors.len` - let next = remaining.first_mut().unwrap(); - let overflow = total_count - row_count; - - if next.row_count != overflow { - selectors.push(RowSelector { - row_count: next.row_count - overflow, - skip: next.skip, - }) - } - next.row_count = overflow; - - std::mem::swap(&mut remaining, selectors); - Self::from_selectors(remaining) + Self::from_selectors(split_off_selectors(self.selectors_mut(), row_count)) } /// returns a [`RowSelection`] representing rows that are selected in both /// input [`RowSelection`]s. @@ -699,11 +456,7 @@ impl RowSelection { and_then_row_selections(first, second) } (RowSelectionInner::Selectors(first), RowSelectionInner::Mask(second)) => { - let mut selectors = vec![]; - let mut first = first.iter().copied().peekable(); - let mut second = MaskRunIter::new(second.mask()).peekable(); - and_then_iter(&mut selectors, &mut first, &mut second); - Self::from_selectors(selectors) + and_then_selectors_with_mask(first, second.mask()) } } } @@ -785,10 +538,7 @@ impl RowSelection { } return self; } - let selectors = self.selectors_mut(); - while selectors.last().map(|x| x.skip).unwrap_or(false) { - selectors.pop(); - } + trim_selectors(self.selectors_mut()); self } @@ -798,80 +548,39 @@ impl RowSelection { return self; } - let mut selectors = match self.inner { + match self.inner { RowSelectionInner::Mask(mask) => { let count = mask.count(); let buffer = offset_mask((*mask).into_mask(), offset, count); - return Self::from_mask_selection(MaskSelection::with_count( + Self::from_mask_selection(MaskSelection::with_count( buffer, count.saturating_sub(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 = 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 => { - selectors.clear(); - return Self::from_selectors(selectors); + RowSelectionInner::Selectors(selectors) => { + Self::from_selectors(offset_selectors(selectors, offset)) } - }; - - 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::from_selectors(new_selectors) + } } /// Limit this [`RowSelection`] to only select `limit` rows - pub(crate) fn limit(self, mut limit: usize) -> Self { - let mut selectors = match self.inner { + pub(crate) fn limit(self, limit: usize) -> Self { + match self.inner { RowSelectionInner::Mask(mask) => { let cached = mask.cached_count(); let buffer = limit_mask((*mask).into_mask(), limit); - return match cached { + match cached { Some(count) => Self::from_mask_selection(MaskSelection::with_count( buffer, count.min(limit), )), None => Self::from_boolean_buffer(buffer), - }; - } - RowSelectionInner::Selectors(selectors) => selectors, - }; - if limit == 0 { - selectors.clear(); - } - - for (idx, selection) in selectors.iter_mut().enumerate() { - if !selection.skip { - if selection.row_count >= limit { - selection.row_count = limit; - selectors.truncate(idx + 1); - break; - } else { - limit -= selection.row_count; } } + RowSelectionInner::Selectors(selectors) => { + Self::from_selectors(limit_selectors(selectors, limit)) + } } - Self::from_selectors(selectors) } /// Returns a borrowed iterator yielding the [`RowSelector`]s for this selection. @@ -884,8 +593,8 @@ impl RowSelection { /// 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()), + RowSelectionInner::Selectors(s) => RowSelectionIter::new(s), + RowSelectionInner::Mask(m) => RowSelectionIter::new(m.selectors()), } } @@ -946,31 +655,7 @@ impl From for RowSelection { impl FromIterator for RowSelection { fn from_iter>(iter: T) -> Self { - let iter = iter.into_iter(); - - // Capacity before filter - let mut selectors = Vec::with_capacity(iter.size_hint().0); - - let mut filtered = iter.filter(|x| x.row_count != 0); - if let Some(x) = filtered.next() { - selectors.push(x); - } - - for s in filtered { - if s.row_count == 0 { - continue; - } - - // Combine consecutive selectors - let last = selectors.last_mut().unwrap(); - if last.skip == s.skip { - last.row_count = last.row_count.checked_add(s.row_count).unwrap(); - } else { - selectors.push(s) - } - } - - Self::from_selectors(selectors) + Self::from_selectors(combine_selectors(iter)) } } @@ -1024,376 +709,9 @@ impl FromIterator for RowSelection { } } -fn and_then_row_selections(first: &[RowSelector], second: &[RowSelector]) -> RowSelection { - let mut selectors = vec![]; - let mut first = first.iter().copied().peekable(); - let mut second = second.iter().copied().peekable(); - and_then_iter(&mut selectors, &mut first, &mut second); - RowSelection::from_selectors(selectors) -} - -fn and_then_iter( - selectors: &mut Vec, - first: &mut std::iter::Peekable, - second: &mut std::iter::Peekable, -) where - I: Iterator, - J: Iterator, -{ - let mut to_skip = 0; - while let Some(b) = second.peek_mut() { - let a = first - .peek_mut() - .expect("selection exceeds the number of selected rows"); - - if b.row_count == 0 { - second.next().unwrap(); - continue; - } - - if a.row_count == 0 { - first.next().unwrap(); - continue; - } - - if a.skip { - // Records were skipped when producing second - to_skip += a.row_count; - first.next().unwrap(); - continue; - } - - let skip = b.skip; - let to_process = a.row_count.min(b.row_count); - - a.row_count -= to_process; - b.row_count -= to_process; - - match skip { - true => to_skip += to_process, - false => { - if to_skip != 0 { - selectors.push(RowSelector::skip(to_skip)); - to_skip = 0; - } - selectors.push(RowSelector::select(to_process)) - } - } - } - - for v in first { - if v.row_count != 0 { - assert!( - v.skip, - "selection contains less than the number of selected rows" - ); - to_skip += v.row_count - } - } - - if to_skip != 0 { - selectors.push(RowSelector::skip(to_skip)); - } -} - -/// Combine two lists of `RowSelection` return the intersection of them -/// For example: -/// self: NNYYYYNNYYNYN -/// other: NYNNNNNNY -/// -/// returned: NNNNNNNNYYNYN -fn intersect_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection { - let mut l_iter = left.iter().copied().peekable(); - let mut r_iter = right.iter().copied().peekable(); - - let iter = std::iter::from_fn(move || { - loop { - let l = l_iter.peek_mut(); - let r = r_iter.peek_mut(); - - match (l, r) { - (Some(a), _) if a.row_count == 0 => { - l_iter.next().unwrap(); - } - (_, Some(b)) if b.row_count == 0 => { - r_iter.next().unwrap(); - } - (Some(l), Some(r)) => { - return match (l.skip, r.skip) { - // Keep both ranges - (false, false) => { - if l.row_count < r.row_count { - r.row_count -= l.row_count; - l_iter.next() - } else { - l.row_count -= r.row_count; - r_iter.next() - } - } - // skip at least one - _ => { - if l.row_count < r.row_count { - let skip = l.row_count; - r.row_count -= l.row_count; - l_iter.next(); - Some(RowSelector::skip(skip)) - } else { - let skip = r.row_count; - l.row_count -= skip; - r_iter.next(); - Some(RowSelector::skip(skip)) - } - } - }; - } - (Some(_), None) => return l_iter.next(), - (None, Some(_)) => return r_iter.next(), - (None, None) => return None, - } - } - }); - - iter.collect() -} - -/// Combine two lists of `RowSelector` return the union of them -/// For example: -/// self: NNYYYYNNYYNYN -/// other: NYNNNNNNY -/// -/// returned: NYYYYYNNYYNYN -/// -/// This can be removed from here once RowSelection::union is in parquet::arrow -fn union_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection { - let mut l_iter = left.iter().copied().peekable(); - let mut r_iter = right.iter().copied().peekable(); - - let iter = std::iter::from_fn(move || { - loop { - let l = l_iter.peek_mut(); - let r = r_iter.peek_mut(); - - match (l, r) { - (Some(a), _) if a.row_count == 0 => { - l_iter.next().unwrap(); - } - (_, Some(b)) if b.row_count == 0 => { - r_iter.next().unwrap(); - } - (Some(l), Some(r)) => { - return match (l.skip, r.skip) { - // Skip both ranges - (true, true) => { - if l.row_count < r.row_count { - let skip = l.row_count; - r.row_count -= l.row_count; - l_iter.next(); - Some(RowSelector::skip(skip)) - } else { - let skip = r.row_count; - l.row_count -= skip; - r_iter.next(); - Some(RowSelector::skip(skip)) - } - } - // Keep rows from left - (false, true) => { - if l.row_count < r.row_count { - r.row_count -= l.row_count; - l_iter.next() - } else { - let r_row_count = r.row_count; - l.row_count -= r_row_count; - r_iter.next(); - Some(RowSelector::select(r_row_count)) - } - } - // Keep rows from right - (true, false) => { - if l.row_count < r.row_count { - let l_row_count = l.row_count; - r.row_count -= l_row_count; - l_iter.next(); - Some(RowSelector::select(l_row_count)) - } else { - l.row_count -= r.row_count; - r_iter.next() - } - } - // Keep at least one - _ => { - if l.row_count < r.row_count { - r.row_count -= l.row_count; - l_iter.next() - } else { - l.row_count -= r.row_count; - r_iter.next() - } - } - }; - } - (Some(_), None) => return l_iter.next(), - (None, Some(_)) => return r_iter.next(), - (None, None) => return None, - } - } - }); - - iter.collect() -} - -/// Cursor for iterating a selector-backed [`RowSelection`] -/// -/// This is best for sparse selections where large contiguous -/// blocks of rows are selected or skipped. -#[derive(Debug)] -pub struct SelectorsCursor { - selectors: VecDeque, - /// Current absolute offset into the selection - position: usize, -} - -impl SelectorsCursor { - /// Returns `true` when no further rows remain - pub fn is_empty(&self) -> bool { - self.selectors.is_empty() - } - - pub(crate) fn selectors_mut(&mut self) -> &mut VecDeque { - &mut self.selectors - } - - /// Return the next [`RowSelector`] - pub(crate) fn next_selector(&mut self) -> RowSelector { - let selector = self.selectors.pop_front().unwrap(); - self.position += selector.row_count; - selector - } - - /// Return a selector to the front, rewinding the position - pub(crate) fn return_selector(&mut self, selector: RowSelector) { - self.position = self.position.saturating_sub(selector.row_count); - self.selectors.push_front(selector); - } -} - -/// Row ranges whose backing pages are loaded for every projected column. -#[derive(Clone, Debug)] -pub(crate) struct LoadedRowRanges(Vec>); - -impl LoadedRowRanges { - pub(crate) fn from_selection(selection: RowSelection) -> Self { - let selectors: Vec = selection.into(); - let mut position = 0; - let ranges = selectors - .into_iter() - .filter_map(|selector| { - let start = position; - position += selector.row_count; - (!selector.skip).then_some(start..position) - }) - .collect(); - Self(ranges) - } - - fn end_containing(&self, row: usize) -> Option { - let idx = self.0.partition_point(|range| range.end <= row); - self.0 - .get(idx) - .filter(|range| range.start <= row) - .map(|range| range.end) - } - - #[cfg(test)] - pub(crate) fn ranges(&self) -> &[Range] { - &self.0 - } -} - -/// Cursor for iterating a [`RowSelection`] during execution within a -/// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan). -/// -/// This keeps per-reader state such as the current position and delegates the -/// actual storage strategy to the internal `RowSelectionInner`. -#[derive(Debug)] -pub enum RowSelectionCursor { - /// Reading all rows - All, - /// Use a bitmask to back the selection (dense selections) - Mask(MaskCursor), - /// Use a queue of selectors to back the selection (sparse selections) - Selectors(SelectorsCursor), -} - -impl RowSelectionCursor { - /// Create a [`MaskCursor`] cursor backed by a bitmask, from an existing set of selectors - pub(crate) fn new_mask_from_selectors( - selectors: Vec, - loaded_row_ranges: Option>, - ) -> Self { - debug_assert!( - selectors - .last() - .map(|selector| !selector.skip) - .unwrap_or(true), - "Mask selectors must not end with a skip" - ); - Self::Mask(MaskCursor { - mask: boolean_mask_from_selectors(&selectors), - position: 0, - loaded_row_ranges, - }) - } - - /// Create a [`MaskCursor`] cursor backed by an existing bitmask. - pub(crate) fn new_mask_from_buffer( - mask: BooleanBuffer, - loaded_row_ranges: Option>, - ) -> Self { - debug_assert!( - mask.is_empty() || mask.value(mask.len() - 1), - "Mask selections must not end with a skip" - ); - Self::Mask(MaskCursor { - mask, - position: 0, - loaded_row_ranges, - }) - } - - /// Create a [`RowSelectionCursor::Selectors`] from the provided selectors - pub(crate) fn new_selectors(selectors: Vec) -> Self { - Self::Selectors(SelectorsCursor { - selectors: selectors.into(), - position: 0, - }) - } - - /// Create a cursor that selects all rows - pub(crate) fn new_all() -> Self { - Self::All - } -} - #[cfg(test)] mod tests { use super::*; - use rand::{Rng, rng}; - - #[test] - fn test_from_selectors_skips_empty_selectors() { - let selection = RowSelection::from(vec![ - RowSelector::select(0), - RowSelector::skip(0), - RowSelector::select(2), - RowSelector::select(0), - RowSelector::skip(1), - ]); - assert_eq!( - selection.selectors(), - vec![RowSelector::select(2), RowSelector::skip(1)] - ); - } #[test] fn test_offset_zero_and_zero_batch_expand_are_identity() { @@ -1450,770 +768,138 @@ mod tests { } #[test] - fn test_split_off() { - let mut selection = RowSelection::from(vec![ + fn test_iter() { + // use the iter() API to show it does what is expected and + // avoid accidental deletion + let selectors = vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(4), + ]; + + let round_tripped: Vec = RowSelection::from(selectors.clone()) + .iter() + .copied() + .collect(); + assert_eq!(selectors, round_tripped); + } + + #[test] + fn test_row_count() { + let selection = RowSelection::from(vec![ RowSelector::skip(34), RowSelector::select(12), RowSelector::skip(3), RowSelector::select(35), ]); - let split = selection.split_off(34); - assert_eq!(split.selectors(), vec![RowSelector::skip(34)]); - assert_eq!( - selection.selectors(), - vec![ - RowSelector::select(12), - RowSelector::skip(3), - RowSelector::select(35) - ] - ); + assert_eq!(selection.row_count(), 12 + 35); + assert_eq!(selection.skipped_row_count(), 34 + 3); - let split = selection.split_off(5); - assert_eq!(split.selectors(), vec![RowSelector::select(5)]); - assert_eq!( - selection.selectors(), - vec![ - RowSelector::select(7), - RowSelector::skip(3), - RowSelector::select(35) - ] - ); - - let split = selection.split_off(8); - assert_eq!( - split.selectors(), - vec![RowSelector::select(7), RowSelector::skip(1)] - ); - assert_eq!( - selection.selectors(), - vec![RowSelector::skip(2), RowSelector::select(35)] - ); - - let split = selection.split_off(200); - assert_eq!( - split.selectors(), - vec![RowSelector::skip(2), RowSelector::select(35)] - ); - assert!(selection.selectors().is_empty()); - } - - #[test] - fn test_offset() { - let selection = RowSelection::from(vec![ - RowSelector::select(5), - RowSelector::skip(23), - RowSelector::select(7), - RowSelector::skip(33), - RowSelector::select(6), - ]); + let selection = RowSelection::from(vec![RowSelector::select(12), RowSelector::select(35)]); - let selection = selection.offset(2); - assert_eq!( - selection.selectors(), - vec![ - RowSelector::skip(2), - RowSelector::select(3), - RowSelector::skip(23), - RowSelector::select(7), - RowSelector::skip(33), - RowSelector::select(6), - ] - ); + assert_eq!(selection.row_count(), 12 + 35); + assert_eq!(selection.skipped_row_count(), 0); - let selection = selection.offset(5); - assert_eq!( - selection.selectors(), - vec![ - RowSelector::skip(30), - RowSelector::select(5), - RowSelector::skip(33), - RowSelector::select(6), - ] - ); + let selection = RowSelection::from(vec![RowSelector::skip(34), RowSelector::skip(3)]); - let selection = selection.offset(3); - assert_eq!( - selection.selectors(), - vec![ - RowSelector::skip(33), - RowSelector::select(2), - RowSelector::skip(33), - RowSelector::select(6), - ] - ); + assert_eq!(selection.row_count(), 0); + assert_eq!(selection.skipped_row_count(), 34 + 3); - let selection = selection.offset(2); - assert_eq!( - selection.selectors(), - vec![RowSelector::skip(68), RowSelector::select(6),] - ); + let selection = RowSelection::from(vec![]); - let selection = selection.offset(3); - assert_eq!( - selection.selectors(), - vec![RowSelector::skip(71), RowSelector::select(3),] - ); + assert_eq!(selection.row_count(), 0); + assert_eq!(selection.skipped_row_count(), 0); } #[test] - fn test_and() { - let mut a = RowSelection::from(vec![ - RowSelector::skip(12), - RowSelector::select(23), - RowSelector::skip(3), - RowSelector::select(5), - ]); - - let b = RowSelection::from(vec![ - RowSelector::select(5), - RowSelector::skip(4), - RowSelector::select(15), - RowSelector::skip(4), - ]); + fn test_mixed_backing_equality_mismatches() { + let mask = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true, true])); - let mut expected = RowSelection::from(vec![ - RowSelector::skip(12), - RowSelector::select(5), - RowSelector::skip(4), - RowSelector::select(14), - RowSelector::skip(3), + // Total row counts differ + let longer = RowSelection::from(vec![ RowSelector::select(1), - RowSelector::skip(4), - ]); - - assert_eq!(a.and_then(&b), expected); - - a.split_off(7); - expected.split_off(7); - assert_eq!(a.and_then(&b), expected); - - let a = RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(3)]); - - let b = RowSelection::from(vec![ - RowSelector::select(2), RowSelector::skip(1), - RowSelector::select(1), + RowSelector::select(2), RowSelector::skip(1), ]); + assert_ne!(mask, longer); + assert_ne!(longer, mask); - assert_eq!( - a.and_then(&b).selectors(), - vec![ - RowSelector::select(2), - RowSelector::skip(1), - RowSelector::select(1), - RowSelector::skip(4) - ] - ); - } - - #[test] - fn test_combine() { - let a = vec![ - RowSelector::skip(3), - RowSelector::skip(3), - RowSelector::select(10), - RowSelector::skip(4), - ]; - - let b = vec![ - RowSelector::skip(3), - RowSelector::skip(3), - RowSelector::select(10), - RowSelector::skip(4), - RowSelector::skip(0), - ]; + // A selected bit falls inside a skip run + let skip_overlap = RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]); + assert_ne!(mask, skip_overlap); - let c = vec![ - RowSelector::skip(2), - RowSelector::skip(4), - RowSelector::select(3), - RowSelector::select(3), - RowSelector::select(4), - RowSelector::skip(3), + // Select run boundaries do not line up + let misaligned = RowSelection::from(vec![ + RowSelector::select(2), RowSelector::skip(1), - RowSelector::skip(0), - ]; - - let expected = RowSelection::from(vec![ - RowSelector::skip(6), - RowSelector::select(10), - RowSelector::skip(4), - ]); - - assert_eq!(RowSelection::from_iter(a), expected); - assert_eq!(RowSelection::from_iter(b), expected); - assert_eq!(RowSelection::from_iter(c), expected); - } - - #[test] - fn test_combine_2elements() { - let a = vec![RowSelector::select(10), RowSelector::select(5)]; - let a_expect = vec![RowSelector::select(15)]; - assert_eq!(RowSelection::from_iter(a).selectors(), a_expect); - - let b = vec![RowSelector::select(10), RowSelector::skip(5)]; - let b_expect = vec![RowSelector::select(10), RowSelector::skip(5)]; - assert_eq!(RowSelection::from_iter(b).selectors(), b_expect); - - let c = vec![RowSelector::skip(10), RowSelector::select(5)]; - let c_expect = vec![RowSelector::skip(10), RowSelector::select(5)]; - assert_eq!(RowSelection::from_iter(c).selectors(), c_expect); - - let d = vec![RowSelector::skip(10), RowSelector::skip(5)]; - let d_expect = vec![RowSelector::skip(15)]; - assert_eq!(RowSelection::from_iter(d).selectors(), d_expect); - } - - #[test] - fn test_from_one_and_empty() { - let a = vec![RowSelector::select(10)]; - let selection1 = RowSelection::from(a.clone()); - assert_eq!(selection1.selectors(), a); - - let b = vec![]; - let selection1 = RowSelection::from(b.clone()); - assert_eq!(selection1.selectors(), b) - } - - #[test] - #[should_panic(expected = "selection exceeds the number of selected rows")] - fn test_and_longer() { - let a = RowSelection::from(vec![ - RowSelector::select(3), - RowSelector::skip(33), - RowSelector::select(3), - RowSelector::skip(33), - ]); - let b = RowSelection::from(vec![RowSelector::select(36)]); - a.and_then(&b); - } - - #[test] - #[should_panic(expected = "selection contains less than the number of selected rows")] - fn test_and_shorter() { - let a = RowSelection::from(vec![ - RowSelector::select(3), - RowSelector::skip(33), - RowSelector::select(3), - RowSelector::skip(33), + RowSelector::select(1), ]); - let b = RowSelection::from(vec![RowSelector::select(3)]); - a.and_then(&b); - } + assert_ne!(mask, misaligned); - #[test] - fn test_intersect_row_selection_and_combine() { - // a size equal b size - let a = vec![ - RowSelector::select(5), - RowSelector::skip(4), + let equal = RowSelection::from(vec![ RowSelector::select(1), - ]; - let b = vec![ - RowSelector::select(8), RowSelector::skip(1), - RowSelector::select(1), - ]; - - let res = intersect_row_selections(&a, &b); - assert_eq!( - res.selectors(), - vec![ - RowSelector::select(5), - RowSelector::skip(4), - RowSelector::select(1), - ], - ); - - // a size larger than b size - let a = vec![ - RowSelector::select(3), - RowSelector::skip(33), - RowSelector::select(3), - RowSelector::skip(33), - ]; - let b = vec![RowSelector::select(36), RowSelector::skip(36)]; - let res = intersect_row_selections(&a, &b); - assert_eq!( - res.selectors(), - vec![RowSelector::select(3), RowSelector::skip(69)] - ); - - // a size less than b size - let a = vec![RowSelector::select(3), RowSelector::skip(7)]; - let b = vec![ - RowSelector::select(2), - RowSelector::skip(2), - RowSelector::select(2), - RowSelector::skip(2), - RowSelector::select(2), - ]; - let res = intersect_row_selections(&a, &b); - assert_eq!( - res.selectors(), - vec![RowSelector::select(2), RowSelector::skip(8)] - ); - - let a = vec![RowSelector::select(3), RowSelector::skip(7)]; - let b = vec![ - RowSelector::select(2), - RowSelector::skip(2), RowSelector::select(2), - RowSelector::skip(2), - RowSelector::select(2), - ]; - let res = intersect_row_selections(&a, &b); - assert_eq!( - res.selectors(), - vec![RowSelector::select(2), RowSelector::skip(8)] - ); - } - - #[test] - fn test_and_fuzz() { - let mut rand = rng(); - for _ in 0..100 { - let a_len = rand.random_range(10..100); - let a_bools: Vec<_> = (0..a_len).map(|_| rand.random_bool(0.2)).collect(); - let a = RowSelection::from_filters(&[BooleanArray::from(a_bools.clone())]); - - let b_len: usize = a_bools.iter().map(|x| *x as usize).sum(); - let b_bools: Vec<_> = (0..b_len).map(|_| rand.random_bool(0.8)).collect(); - let b = RowSelection::from_filters(&[BooleanArray::from(b_bools.clone())]); - - let mut expected_bools = vec![false; a_len]; - - let mut iter_b = b_bools.iter(); - for (idx, b) in a_bools.iter().enumerate() { - if *b && *iter_b.next().unwrap() { - expected_bools[idx] = true; - } - } - - let expected = RowSelection::from_filters(&[BooleanArray::from(expected_bools)]); - - let total_rows: usize = expected.selectors().iter().map(|s| s.row_count).sum(); - assert_eq!(a_len, total_rows); - - assert_eq!(a.and_then(&b), expected); - } + ]); + assert_eq!(mask, equal); + assert_eq!(equal, mask); } #[test] - fn test_iter() { - // use the iter() API to show it does what is expected and - // avoid accidental deletion - let selectors = vec![ - RowSelector::select(3), - RowSelector::skip(33), - RowSelector::select(4), + fn test_from_iter_all_mask_preserves_mask_backing() { + let a_bits = vec![true, false, true, true]; + let b_bits = vec![false, true, false]; + let c_bits = vec![true, true, false, false, true]; + + let parts = vec![ + RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())), + RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())), + RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())), ]; + let collected: RowSelection = parts.into_iter().collect(); - let round_tripped: Vec = RowSelection::from(selectors.clone()) + let combined = a_bits .iter() + .chain(b_bits.iter()) + .chain(c_bits.iter()) .copied() - .collect(); - assert_eq!(selectors, round_tripped); - } - - #[test] - fn test_limit() { - // Limit to existing limit should no-op - let selection = RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(90)]); - let limited = selection.limit(10); - assert_eq!(RowSelection::from(vec![RowSelector::select(10)]), limited); + .collect::>(); + let expected = RowSelection::from_filters(&[BooleanArray::from(combined)]); - let selection = RowSelection::from(vec![ - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(10), - ]); - - let limited = selection.clone().limit(5); - let expected = vec![RowSelector::select(5)]; - assert_eq!(limited.selectors(), expected); - - let limited = selection.clone().limit(15); - let expected = vec![ - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(5), - ]; - assert_eq!(limited.selectors(), expected); - - let limited = selection.clone().limit(0); - let expected = vec![]; - assert_eq!(limited.selectors(), expected); - - let limited = selection.clone().limit(30); - let expected = vec![ - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(10), - ]; - assert_eq!(limited.selectors(), expected); - - let limited = selection.limit(100); - let expected = vec![ - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(10), - ]; - assert_eq!(limited.selectors(), expected); + assert!(collected.as_mask().is_some()); + assert_eq!(collected, expected); } #[test] - fn test_scan_ranges() { - let index = vec![ - PageLocation { - offset: 0, - compressed_page_size: 10, - first_row_index: 0, - }, - PageLocation { - offset: 10, - compressed_page_size: 10, - first_row_index: 10, - }, - PageLocation { - offset: 20, - compressed_page_size: 10, - first_row_index: 20, - }, - PageLocation { - offset: 30, - compressed_page_size: 10, - first_row_index: 30, - }, - PageLocation { - offset: 40, - compressed_page_size: 10, - first_row_index: 40, - }, - PageLocation { - offset: 50, - compressed_page_size: 10, - first_row_index: 50, - }, - PageLocation { - offset: 60, - compressed_page_size: 10, - first_row_index: 60, - }, + fn test_from_iter_mixed_backing_falls_back_to_selectors() { + let a_bits = vec![true, false, true]; + let b_selectors = vec![RowSelector::skip(2), RowSelector::select(3)]; + let c_bits = vec![false, true]; + + let parts = vec![ + RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())), + RowSelection::from(b_selectors), + RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())), ]; + let collected: RowSelection = parts.into_iter().collect(); - let selection = RowSelection::from(vec![ - // Skip first page - RowSelector::skip(10), - // Multiple selects in same page - RowSelector::select(3), - RowSelector::skip(3), - RowSelector::select(4), - // Select to page boundary - RowSelector::skip(5), - RowSelector::select(5), - // Skip full page past page boundary - RowSelector::skip(12), - // Select across page boundaries - RowSelector::select(12), - // Skip final page - RowSelector::skip(12), - ]); - - let ranges = selection.scan_ranges(&index); + assert!(collected.as_mask().is_none()); - // assert_eq!(mask, vec![false, true, true, false, true, true, false]); - assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60]); - assert_eq!( - selection.row_ranges_for_selected_pages(&index, 70), - vec![10..20, 20..30, 40..50, 50..60] - ); - - let selection = RowSelection::from(vec![ - // Skip first page - RowSelector::skip(10), - // Multiple selects in same page - RowSelector::select(3), - RowSelector::skip(3), - RowSelector::select(4), - // Select to page boundary - RowSelector::skip(5), - RowSelector::select(5), - // Skip full page past page boundary - RowSelector::skip(12), - // Select across page boundaries - RowSelector::select(12), - RowSelector::skip(1), - // Select across page boundaries including final page - RowSelector::select(8), - ]); - - let ranges = selection.scan_ranges(&index); - - // assert_eq!(mask, vec![false, true, true, false, true, true, true]); - assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]); - - let selection = RowSelection::from(vec![ - // Skip first page - RowSelector::skip(10), - // Multiple selects in same page - RowSelector::select(3), - RowSelector::skip(3), - RowSelector::select(4), - // Select to page boundary - RowSelector::skip(5), - RowSelector::select(5), - // Skip full page past page boundary - RowSelector::skip(12), - // Select to final page boundary - RowSelector::select(12), - RowSelector::skip(1), - // Skip across final page boundary - RowSelector::skip(8), - // Select from final page - RowSelector::select(4), - ]); - - let ranges = selection.scan_ranges(&index); - - // assert_eq!(mask, vec![false, true, true, false, true, true, true]); - assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]); - - let selection = RowSelection::from(vec![ - // Skip first page - RowSelector::skip(10), - // Multiple selects in same page - RowSelector::select(3), - RowSelector::skip(3), - RowSelector::select(4), - // Select to remaining in page and first row of next page - RowSelector::skip(5), - RowSelector::select(6), - // Skip remaining - RowSelector::skip(50), - ]); - - let ranges = selection.scan_ranges(&index); - - // assert_eq!(mask, vec![false, true, true, false, true, true, true]); - assert_eq!(ranges, vec![10..20, 20..30, 30..40]); - } - - #[test] - fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() { - let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges( - std::iter::once(0..5), - 10, - )); - let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors( - vec![RowSelector::select(1)], - Some(loaded.into()), - ) else { - unreachable!() - }; - - let chunk = cursor.next_chunk(10).unwrap(); - assert_eq!(chunk.chunk_rows, 1); - assert!(cursor.is_empty()); - } - - #[test] - fn test_next_mask_chunk_until_cursor_is_empty() { - let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors( - vec![ - RowSelector::skip(2), - RowSelector::select(2), - RowSelector::skip(1), - RowSelector::select(1), - ], - None, - ) else { - unreachable!() - }; - - let first = cursor.next_mask_chunk(2).unwrap(); - assert_eq!(first.initial_skip, 2); - assert_eq!(first.chunk_rows, 2); - assert_eq!(first.selected_rows, 2); - - let second = cursor.next_mask_chunk(2).unwrap(); - assert_eq!(second.initial_skip, 1); - assert_eq!(second.chunk_rows, 1); - assert_eq!(second.selected_rows, 1); - - assert!(cursor.next_mask_chunk(2).is_none()); - } - - #[test] - fn test_from_ranges() { - let ranges = [1..3, 4..6, 6..6, 8..8, 9..10]; - let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(), 10); - assert_eq!( - selection.selectors(), - vec![ - RowSelector::skip(1), - RowSelector::select(2), - RowSelector::skip(1), - RowSelector::select(2), - RowSelector::skip(3), - RowSelector::select(1) - ] - ); - - let out_of_order_ranges = [1..3, 8..10, 4..7]; - let result = std::panic::catch_unwind(|| { - RowSelection::from_consecutive_ranges(out_of_order_ranges.into_iter(), 10) - }); - assert!(result.is_err()); - } - - #[test] - fn test_empty_selector() { - let selection = RowSelection::from(vec![ - RowSelector::skip(0), - RowSelector::select(2), - RowSelector::skip(0), - RowSelector::select(2), - ]); - assert_eq!(selection.selectors(), vec![RowSelector::select(4)]); - - let selection = RowSelection::from(vec![ - RowSelector::select(0), - RowSelector::skip(2), - RowSelector::select(0), - RowSelector::skip(2), - ]); - assert_eq!(selection.selectors(), vec![RowSelector::skip(4)]); - } - - #[test] - fn test_intersection() { - let selection = RowSelection::from(vec![RowSelector::select(1048576)]); - let result = selection.intersection(&selection); - assert_eq!(result, selection); - - let a = RowSelection::from(vec![ - RowSelector::skip(10), - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(20), - ]); - - let b = RowSelection::from(vec![ - RowSelector::skip(20), - RowSelector::select(20), - RowSelector::skip(10), - ]); - - let result = a.intersection(&b); - assert_eq!( - result.selectors(), - vec![ - RowSelector::skip(30), - RowSelector::select(10), - RowSelector::skip(10) - ] - ); - } - - #[test] - fn test_union() { - let selection = RowSelection::from(vec![RowSelector::select(1048576)]); - let result = selection.union(&selection); - assert_eq!(result, selection); - - // NYNYY - let a = RowSelection::from(vec![ - RowSelector::skip(10), - RowSelector::select(10), - RowSelector::skip(10), - RowSelector::select(20), - ]); - - // NNYYNYN - let b = RowSelection::from(vec![ - RowSelector::skip(20), - RowSelector::select(20), - RowSelector::skip(10), - RowSelector::select(10), - RowSelector::skip(10), - ]); - - let result = a.union(&b); - - // NYYYYYN - assert_eq!( - result.iter().copied().collect::>(), - vec![ - RowSelector::skip(10), - RowSelector::select(50), - RowSelector::skip(10), - ] - ); - } - - #[test] - fn test_row_count() { - let selection = RowSelection::from(vec![ - RowSelector::skip(34), - RowSelector::select(12), - RowSelector::skip(3), - RowSelector::select(35), - ]); - - assert_eq!(selection.row_count(), 12 + 35); - assert_eq!(selection.skipped_row_count(), 34 + 3); - - let selection = RowSelection::from(vec![RowSelector::select(12), RowSelector::select(35)]); - - assert_eq!(selection.row_count(), 12 + 35); - assert_eq!(selection.skipped_row_count(), 0); - - let selection = RowSelection::from(vec![RowSelector::skip(34), RowSelector::skip(3)]); - - assert_eq!(selection.row_count(), 0); - assert_eq!(selection.skipped_row_count(), 34 + 3); - - let selection = RowSelection::from(vec![]); - - assert_eq!(selection.row_count(), 0); - assert_eq!(selection.skipped_row_count(), 0); + let combined_bits = vec![ + true, false, true, false, false, true, true, true, false, true, + ]; + let expected = RowSelection::from_filters(&[BooleanArray::from(combined_bits)]); + assert_eq!(collected, expected); } #[test] - fn test_trim() { - let selection = RowSelection::from(vec![ - RowSelector::skip(34), - RowSelector::select(12), - RowSelector::skip(3), - RowSelector::select(35), - ]); - - let expected = vec![ - RowSelector::skip(34), - RowSelector::select(12), - RowSelector::skip(3), - RowSelector::select(35), - ]; - - assert_eq!(selection.trim().selectors(), expected); - - let selection = RowSelection::from(vec![ - RowSelector::skip(34), - RowSelector::select(12), - RowSelector::skip(3), - ]); - - let expected = vec![RowSelector::skip(34), RowSelector::select(12)]; - - assert_eq!(selection.trim().selectors(), expected); + fn test_from_iter_empty_yields_empty_selection() { + let collected: RowSelection = std::iter::empty::().collect(); + assert_eq!(collected, RowSelection::default()); + assert!(collected.as_mask().is_some()); + assert_eq!(collected.as_mask().unwrap().len(), 0); } } diff --git a/parquet/src/arrow/arrow_reader/selection/ranges.rs b/parquet/src/arrow/arrow_reader/selection/ranges.rs new file mode 100644 index 000000000000..4c7860d095fa --- /dev/null +++ b/parquet/src/arrow/arrow_reader/selection/ranges.rs @@ -0,0 +1,283 @@ +// 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. + +//! Mapping the [`RowSelector`] runs of a selection onto ranges: the byte ranges +//! of the data pages that must be fetched ([`RowSelection::scan_ranges`]) and +//! the expansion of a selection to batch boundaries. +//! +//! Both are shared by the selector and mask backings, which stream their runs +//! from a slice and a [`MaskRunIter`] respectively. +//! +//! [`MaskRunIter`]: crate::arrow::arrow_reader::MaskRunIter + +use super::{RowSelection, RowSelector}; +use crate::file::page_index::offset_index::PageLocation; +use std::ops::Range; + +/// Byte ranges of the data pages containing at least one selected row. +#[inline] +pub(super) fn scan_ranges_from_selectors( + selectors: I, + page_locations: &[PageLocation], +) -> Vec> +where + I: IntoIterator, +{ + let mut ranges: Vec> = 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 +} + +/// Grows each selected run to the batch boundaries containing it, merging the +/// runs that overlap as a result. +#[inline] +pub(super) fn expand_to_batch_boundaries_from_selectors( + selectors: I, + batch_size: usize, + total_rows: usize, +) -> RowSelection +where + I: IntoIterator, +{ + 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> = 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) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_scan_ranges() { + let index = vec![ + PageLocation { + offset: 0, + compressed_page_size: 10, + first_row_index: 0, + }, + PageLocation { + offset: 10, + compressed_page_size: 10, + first_row_index: 10, + }, + PageLocation { + offset: 20, + compressed_page_size: 10, + first_row_index: 20, + }, + PageLocation { + offset: 30, + compressed_page_size: 10, + first_row_index: 30, + }, + PageLocation { + offset: 40, + compressed_page_size: 10, + first_row_index: 40, + }, + PageLocation { + offset: 50, + compressed_page_size: 10, + first_row_index: 50, + }, + PageLocation { + offset: 60, + compressed_page_size: 10, + first_row_index: 60, + }, + ]; + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to page boundary + RowSelector::skip(5), + RowSelector::select(5), + // Skip full page past page boundary + RowSelector::skip(12), + // Select across page boundaries + RowSelector::select(12), + // Skip final page + RowSelector::skip(12), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, false]); + assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60]); + assert_eq!( + selection.row_ranges_for_selected_pages(&index, 70), + vec![10..20, 20..30, 40..50, 50..60] + ); + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to page boundary + RowSelector::skip(5), + RowSelector::select(5), + // Skip full page past page boundary + RowSelector::skip(12), + // Select across page boundaries + RowSelector::select(12), + RowSelector::skip(1), + // Select across page boundaries including final page + RowSelector::select(8), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, true]); + assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]); + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to page boundary + RowSelector::skip(5), + RowSelector::select(5), + // Skip full page past page boundary + RowSelector::skip(12), + // Select to final page boundary + RowSelector::select(12), + RowSelector::skip(1), + // Skip across final page boundary + RowSelector::skip(8), + // Select from final page + RowSelector::select(4), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, true]); + assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]); + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to remaining in page and first row of next page + RowSelector::skip(5), + RowSelector::select(6), + // Skip remaining + RowSelector::skip(50), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, true]); + assert_eq!(ranges, vec![10..20, 20..30, 30..40]); + } +} diff --git a/parquet/src/arrow/arrow_reader/selection/selector.rs b/parquet/src/arrow/arrow_reader/selection/selector.rs new file mode 100644 index 000000000000..43a15b1892dd --- /dev/null +++ b/parquet/src/arrow/arrow_reader/selection/selector.rs @@ -0,0 +1,603 @@ +// 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. + +//! The run length backed representation of a [`RowSelection`]: [`RowSelector`] +//! and the primitives operating on a sequence of them. +//! +//! This is the counterpart of the bitmap backing in the `boolean` module, and +//! provides the same set of transforms (`split_off`, `trim`, `offset`, +//! `limit`) over `Vec` instead of a `BooleanBuffer`. +//! +//! [`RowSelection`]: crate::arrow::arrow_reader::RowSelection + +use std::cmp::Ordering; +use std::ops::Range; + +/// [`RowSelection`] is a collection of [`RowSelector`] used to skip rows when +/// scanning a parquet file +/// +/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub struct RowSelector { + /// The number of rows + pub row_count: usize, + + /// If true, skip `row_count` rows + pub skip: bool, +} + +impl RowSelector { + /// Select `row_count` rows + pub fn select(row_count: usize) -> Self { + Self { + row_count, + skip: false, + } + } + + /// Skip `row_count` rows + pub fn skip(row_count: usize) -> Self { + Self { + row_count, + skip: true, + } + } +} + +/// Borrowed iterator over the [`RowSelector`]s of a +/// [`RowSelection`](crate::arrow::arrow_reader::RowSelection). +#[derive(Debug)] +pub struct RowSelectionIter<'a>(std::slice::Iter<'a, RowSelector>); + +impl<'a> RowSelectionIter<'a> { + pub(super) fn new(selectors: &'a [RowSelector]) -> Self { + Self(selectors.iter()) + } +} + +impl<'a> Iterator for RowSelectionIter<'a> { + type Item = &'a RowSelector; + + #[inline] + fn next(&mut self) -> Option { + self.0.next() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.0.size_hint() + } + + #[inline] + fn count(self) -> usize { + self.0.count() + } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + self.0.nth(n) + } + + #[inline] + fn last(self) -> Option { + self.0.last() + } + + #[inline] + fn fold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + self.0.fold(init, f) + } +} + +impl ExactSizeIterator for RowSelectionIter<'_> {} + +// once it returns None, it will continue returning None +impl std::iter::FusedIterator for RowSelectionIter<'_> {} + +/// Normalizes a sequence of selectors: drops the empty ones and combines +/// consecutive selectors that both skip or both select. +pub(super) fn combine_selectors(iter: I) -> Vec +where + I: IntoIterator, +{ + let iter = iter.into_iter(); + + // Capacity before filter + let mut selectors = Vec::with_capacity(iter.size_hint().0); + + let mut filtered = iter.filter(|x| x.row_count != 0); + if let Some(x) = filtered.next() { + selectors.push(x); + } + + for s in filtered { + if s.row_count == 0 { + continue; + } + + // Combine consecutive selectors + let last = selectors.last_mut().unwrap(); + if last.skip == s.skip { + last.row_count = last.row_count.checked_add(s.row_count).unwrap(); + } else { + selectors.push(s) + } + } + + selectors +} + +/// Builds the selectors keeping `ranges` out of `total_rows` rows. +/// +/// # Panics +/// +/// Panics if `ranges` are not in ascending order. +pub(super) fn selectors_from_consecutive_ranges(ranges: I, total_rows: usize) -> Vec +where + I: Iterator>, +{ + let mut selectors: Vec = Vec::with_capacity(ranges.size_hint().0); + let mut last_end = 0; + for range in ranges { + let len = range.end - range.start; + if len == 0 { + continue; + } + + match range.start.cmp(&last_end) { + Ordering::Equal => match selectors.last_mut() { + Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(), + None => selectors.push(RowSelector::select(len)), + }, + Ordering::Greater => { + selectors.push(RowSelector::skip(range.start - last_end)); + selectors.push(RowSelector::select(len)) + } + Ordering::Less => panic!("out of order"), + } + last_end = range.end; + } + + if last_end != total_rows { + selectors.push(RowSelector::skip(total_rows - last_end)) + } + + selectors +} + +/// Splits the first `row_count` rows off `selectors`, returning them and +/// leaving the remainder in place. +pub(super) fn split_off_selectors( + selectors: &mut Vec, + row_count: usize, +) -> Vec { + let mut total_count = 0; + + // Find the index where the selector exceeds the row count + let find = selectors.iter().position(|selector| { + total_count += selector.row_count; + total_count > row_count + }); + + let split_idx = match find { + Some(idx) => idx, + None => return std::mem::take(selectors), + }; + + let mut remaining = selectors.split_off(split_idx); + + // Always present as `split_idx < selectors.len` + let next = remaining.first_mut().unwrap(); + let overflow = total_count - row_count; + + if next.row_count != overflow { + selectors.push(RowSelector { + row_count: next.row_count - overflow, + skip: next.skip, + }) + } + next.row_count = overflow; + + std::mem::swap(&mut remaining, selectors); + remaining +} + +/// Removes any trailing skips from `selectors`. +pub(super) fn trim_selectors(selectors: &mut Vec) { + while selectors.last().map(|x| x.skip).unwrap_or(false) { + selectors.pop(); + } +} + +/// Skips the first `offset` selected rows of `selectors`. +pub(super) fn offset_selectors(mut selectors: Vec, offset: usize) -> Vec { + let mut selected_count = 0; + let mut skipped_count = 0; + + // Find the index where the selector exceeds the row count + 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 => { + selectors.clear(); + return selectors; + } + }; + + 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..]); + + new_selectors +} + +/// Keeps only the first `limit` selected rows of `selectors`. +pub(super) fn limit_selectors( + mut selectors: Vec, + mut limit: usize, +) -> Vec { + if limit == 0 { + selectors.clear(); + } + + for (idx, selection) in selectors.iter_mut().enumerate() { + if !selection.skip { + if selection.row_count >= limit { + selection.row_count = limit; + selectors.truncate(idx + 1); + break; + } else { + limit -= selection.row_count; + } + } + } + selectors +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::arrow_reader::selection::RowSelection; + + #[test] + fn test_from_selectors_skips_empty_selectors() { + let selection = RowSelection::from(vec![ + RowSelector::select(0), + RowSelector::skip(0), + RowSelector::select(2), + RowSelector::select(0), + RowSelector::skip(1), + ]); + assert_eq!( + selection.selectors(), + vec![RowSelector::select(2), RowSelector::skip(1)] + ); + } + + #[test] + fn test_split_off() { + let mut selection = RowSelection::from(vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35), + ]); + + let split = selection.split_off(34); + assert_eq!(split.selectors(), vec![RowSelector::skip(34)]); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35) + ] + ); + + let split = selection.split_off(5); + assert_eq!(split.selectors(), vec![RowSelector::select(5)]); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::select(7), + RowSelector::skip(3), + RowSelector::select(35) + ] + ); + + let split = selection.split_off(8); + assert_eq!( + split.selectors(), + vec![RowSelector::select(7), RowSelector::skip(1)] + ); + assert_eq!( + selection.selectors(), + vec![RowSelector::skip(2), RowSelector::select(35)] + ); + + let split = selection.split_off(200); + assert_eq!( + split.selectors(), + vec![RowSelector::skip(2), RowSelector::select(35)] + ); + assert!(selection.selectors().is_empty()); + } + + #[test] + fn test_offset() { + let selection = RowSelection::from(vec![ + RowSelector::select(5), + RowSelector::skip(23), + RowSelector::select(7), + RowSelector::skip(33), + RowSelector::select(6), + ]); + + let selection = selection.offset(2); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(23), + RowSelector::select(7), + RowSelector::skip(33), + RowSelector::select(6), + ] + ); + + let selection = selection.offset(5); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(30), + RowSelector::select(5), + RowSelector::skip(33), + RowSelector::select(6), + ] + ); + + let selection = selection.offset(3); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(33), + RowSelector::select(2), + RowSelector::skip(33), + RowSelector::select(6), + ] + ); + + let selection = selection.offset(2); + assert_eq!( + selection.selectors(), + vec![RowSelector::skip(68), RowSelector::select(6),] + ); + + let selection = selection.offset(3); + assert_eq!( + selection.selectors(), + vec![RowSelector::skip(71), RowSelector::select(3),] + ); + } + + #[test] + fn test_combine() { + let a = vec![ + RowSelector::skip(3), + RowSelector::skip(3), + RowSelector::select(10), + RowSelector::skip(4), + ]; + + let b = vec![ + RowSelector::skip(3), + RowSelector::skip(3), + RowSelector::select(10), + RowSelector::skip(4), + RowSelector::skip(0), + ]; + + let c = vec![ + RowSelector::skip(2), + RowSelector::skip(4), + RowSelector::select(3), + RowSelector::select(3), + RowSelector::select(4), + RowSelector::skip(3), + RowSelector::skip(1), + RowSelector::skip(0), + ]; + + let expected = RowSelection::from(vec![ + RowSelector::skip(6), + RowSelector::select(10), + RowSelector::skip(4), + ]); + + assert_eq!(RowSelection::from_iter(a), expected); + assert_eq!(RowSelection::from_iter(b), expected); + assert_eq!(RowSelection::from_iter(c), expected); + } + + #[test] + fn test_combine_2elements() { + let a = vec![RowSelector::select(10), RowSelector::select(5)]; + let a_expect = vec![RowSelector::select(15)]; + assert_eq!(RowSelection::from_iter(a).selectors(), a_expect); + + let b = vec![RowSelector::select(10), RowSelector::skip(5)]; + let b_expect = vec![RowSelector::select(10), RowSelector::skip(5)]; + assert_eq!(RowSelection::from_iter(b).selectors(), b_expect); + + let c = vec![RowSelector::skip(10), RowSelector::select(5)]; + let c_expect = vec![RowSelector::skip(10), RowSelector::select(5)]; + assert_eq!(RowSelection::from_iter(c).selectors(), c_expect); + + let d = vec![RowSelector::skip(10), RowSelector::skip(5)]; + let d_expect = vec![RowSelector::skip(15)]; + assert_eq!(RowSelection::from_iter(d).selectors(), d_expect); + } + + #[test] + fn test_from_one_and_empty() { + let a = vec![RowSelector::select(10)]; + let selection1 = RowSelection::from(a.clone()); + assert_eq!(selection1.selectors(), a); + + let b = vec![]; + let selection1 = RowSelection::from(b.clone()); + assert_eq!(selection1.selectors(), b) + } + + #[test] + fn test_limit() { + // Limit to existing limit should no-op + let selection = RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(90)]); + let limited = selection.limit(10); + assert_eq!(RowSelection::from(vec![RowSelector::select(10)]), limited); + + let selection = RowSelection::from(vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + ]); + + let limited = selection.clone().limit(5); + let expected = vec![RowSelector::select(5)]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.clone().limit(15); + let expected = vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(5), + ]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.clone().limit(0); + let expected = vec![]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.clone().limit(30); + let expected = vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + ]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.limit(100); + let expected = vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + ]; + assert_eq!(limited.selectors(), expected); + } + + #[test] + fn test_from_ranges() { + let ranges = [1..3, 4..6, 6..6, 8..8, 9..10]; + let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(), 10); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(3), + RowSelector::select(1) + ] + ); + + let out_of_order_ranges = [1..3, 8..10, 4..7]; + let result = std::panic::catch_unwind(|| { + RowSelection::from_consecutive_ranges(out_of_order_ranges.into_iter(), 10) + }); + assert!(result.is_err()); + } + + #[test] + fn test_empty_selector() { + let selection = RowSelection::from(vec![ + RowSelector::skip(0), + RowSelector::select(2), + RowSelector::skip(0), + RowSelector::select(2), + ]); + assert_eq!(selection.selectors(), vec![RowSelector::select(4)]); + + let selection = RowSelection::from(vec![ + RowSelector::select(0), + RowSelector::skip(2), + RowSelector::select(0), + RowSelector::skip(2), + ]); + assert_eq!(selection.selectors(), vec![RowSelector::skip(4)]); + } + + #[test] + fn test_trim() { + let selection = RowSelection::from(vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35), + ]); + + let expected = vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35), + ]; + + assert_eq!(selection.trim().selectors(), expected); + + let selection = RowSelection::from(vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + ]); + + let expected = vec![RowSelector::skip(34), RowSelector::select(12)]; + + assert_eq!(selection.trim().selectors(), expected); + } +} From b3efe2d29a7f648d748c44063fc838d1e580cb83 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sat, 25 Jul 2026 15:39:13 +0800 Subject: [PATCH 2/4] refactor(parquet): dispatch on `RowSelectionInner` by `match` in selection Follow up cleanup on the module split, in `selection/mod.rs`: * `trim` and `split_off` dispatched on the backing with `if let` / `matches!` plus an early return, and had to re-destructure the taken `RowSelectionInner` behind `unreachable!()`. Both now `match` on `self.inner` directly, which drops all three `unreachable!()` arms. * `RowSelection::selectors_mut` is removed. Its only two callers were `trim` and `split_off`, and both returned early on the mask backing, so the `Mask` -> `Selectors` promotion inside it was unreachable. * `split_off_selectors` now takes the `Vec` by value and returns `(head, tail)`, mirroring `split_off_mask`. The `mem::take` and the trailing `mem::swap` it needed to write the tail back through `&mut` are gone. * `trim_selectors` is inlined back into `trim`: a four line loop only extracted to mirror `trim_mask`. No behavior change; the set of test names is still identical. --- .../src/arrow/arrow_reader/selection/mod.rs | 110 +++++++++--------- .../arrow/arrow_reader/selection/selector.rs | 25 ++-- 2 files changed, 63 insertions(+), 72 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index 8e4948dd00c9..a9a1231d3d6b 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -57,7 +57,7 @@ use ranges::{expand_to_batch_boundaries_from_selectors, scan_ranges_from_selecto pub use selector::{RowSelectionIter, RowSelector}; use selector::{ combine_selectors, limit_selectors, offset_selectors, selectors_from_consecutive_ranges, - split_off_selectors, trim_selectors, + split_off_selectors, }; /// [`RowSelection`] represents selecting a subset of rows @@ -306,21 +306,6 @@ impl RowSelection { } } - /// Promote a mask-backed selection to selector backing in place. - fn selectors_mut(&mut self) -> &mut Vec { - if let RowSelectionInner::Mask(_) = &self.inner { - let mask = match std::mem::take(&mut self.inner) { - RowSelectionInner::Mask(m) => m, - RowSelectionInner::Selectors(_) => unreachable!(), - }; - self.inner = RowSelectionInner::Selectors(mask_to_selectors(mask.mask())); - } - match &mut self.inner { - RowSelectionInner::Selectors(s) => s, - RowSelectionInner::Mask(_) => unreachable!(), - } - } - /// Creates a [`RowSelection`] from a slice of [`BooleanArray`] /// /// # Panic @@ -396,35 +381,37 @@ impl RowSelection { /// Splits off the first `row_count` from this [`RowSelection`] pub fn split_off(&mut self, row_count: usize) -> Self { - if matches!(&self.inner, RowSelectionInner::Mask(_)) { - let mask = match std::mem::take(&mut self.inner) { - RowSelectionInner::Mask(m) => m, - RowSelectionInner::Selectors(_) => unreachable!(), - }; - let total = mask.cached_count(); - let (head, tail) = split_off_mask((*mask).into_mask(), row_count); - // Popcount only the head and derive the tail by subtraction, so - // repeated splits stay O(bitmap) overall. - let (head, tail) = match total { - Some(total) => { - let head_count = if tail.is_empty() { - total - } else { - head.count_set_bits() - }; - ( - MaskSelection::with_count(head, head_count), - MaskSelection::with_count(tail, total - head_count), - ) - } - None => (MaskSelection::new(head), MaskSelection::new(tail)), - }; - self.inner = RowSelectionInner::Mask(Box::new(tail)); - return Self::from_mask_selection(head); + match std::mem::take(&mut self.inner) { + RowSelectionInner::Mask(mask) => { + let total = mask.cached_count(); + let (head, tail) = split_off_mask((*mask).into_mask(), row_count); + // Popcount only the head and derive the tail by subtraction, so + // repeated splits stay O(bitmap) overall. + let (head, tail) = match total { + Some(total) => { + let head_count = if tail.is_empty() { + total + } else { + head.count_set_bits() + }; + ( + MaskSelection::with_count(head, head_count), + MaskSelection::with_count(tail, total - head_count), + ) + } + None => (MaskSelection::new(head), MaskSelection::new(tail)), + }; + self.inner = RowSelectionInner::Mask(Box::new(tail)); + Self::from_mask_selection(head) + } + RowSelectionInner::Selectors(selectors) => { + let (head, tail) = split_off_selectors(selectors, row_count); + self.inner = RowSelectionInner::Selectors(tail); + Self::from_selectors(head) + } } - - Self::from_selectors(split_off_selectors(self.selectors_mut(), row_count)) } + /// returns a [`RowSelection`] representing rows that are selected in both /// input [`RowSelection`]s. /// @@ -525,21 +512,32 @@ impl RowSelection { } /// Trims this [`RowSelection`] removing any trailing skips - pub(crate) fn trim(mut self) -> Self { - if let RowSelectionInner::Mask(m) = &self.inner { - if let Some(mask) = trim_mask(m.mask()) { - // Trimming only drops trailing unset bits; the count is unchanged. - return match m.cached_count() { - Some(count) => { - Self::from_mask_selection(MaskSelection::with_count(mask, count)) - } - None => Self::from_boolean_buffer(mask), - }; + pub(crate) fn trim(self) -> Self { + match self.inner { + RowSelectionInner::Mask(m) => { + let trimmed = trim_mask(m.mask()); + let cached_count = m.cached_count(); + match trimmed { + // Trimming only drops trailing unset bits; the count is unchanged. + Some(mask) => match cached_count { + Some(count) => { + Self::from_mask_selection(MaskSelection::with_count(mask, count)) + } + None => Self::from_boolean_buffer(mask), + }, + // Nothing to trim, hand the existing box back untouched. + None => Self { + inner: RowSelectionInner::Mask(m), + }, + } + } + RowSelectionInner::Selectors(mut selectors) => { + while selectors.last().map(|x| x.skip).unwrap_or(false) { + selectors.pop(); + } + Self::from_selectors(selectors) } - return self; } - trim_selectors(self.selectors_mut()); - self } /// Applies an offset to this [`RowSelection`], skipping the first `offset` selected rows diff --git a/parquet/src/arrow/arrow_reader/selection/selector.rs b/parquet/src/arrow/arrow_reader/selection/selector.rs index 43a15b1892dd..63a7e3b7ddac 100644 --- a/parquet/src/arrow/arrow_reader/selection/selector.rs +++ b/parquet/src/arrow/arrow_reader/selection/selector.rs @@ -182,12 +182,11 @@ where selectors } -/// Splits the first `row_count` rows off `selectors`, returning them and -/// leaving the remainder in place. +/// Splits `selectors` at the first `row_count` rows, returning `(head, tail)`. pub(super) fn split_off_selectors( - selectors: &mut Vec, + mut selectors: Vec, row_count: usize, -) -> Vec { +) -> (Vec, Vec) { let mut total_count = 0; // Find the index where the selector exceeds the row count @@ -198,13 +197,15 @@ pub(super) fn split_off_selectors( let split_idx = match find { Some(idx) => idx, - None => return std::mem::take(selectors), + None => return (selectors, Vec::new()), }; - let mut remaining = selectors.split_off(split_idx); + // `selectors` keeps the head, `tail` takes the rest. The selector straddling + // the boundary is split between the two. + let mut tail = selectors.split_off(split_idx); // Always present as `split_idx < selectors.len` - let next = remaining.first_mut().unwrap(); + let next = tail.first_mut().unwrap(); let overflow = total_count - row_count; if next.row_count != overflow { @@ -215,15 +216,7 @@ pub(super) fn split_off_selectors( } next.row_count = overflow; - std::mem::swap(&mut remaining, selectors); - remaining -} - -/// Removes any trailing skips from `selectors`. -pub(super) fn trim_selectors(selectors: &mut Vec) { - while selectors.last().map(|x| x.skip).unwrap_or(false) { - selectors.pop(); - } + (selectors, tail) } /// Skips the first `offset` selected rows of `selectors`. From 5a5e2d1e97a6a000aa81b72c82aa5cba965bc591 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sat, 25 Jul 2026 18:31:17 +0800 Subject: [PATCH 3/4] update --- .../src/arrow/arrow_reader/selection/mod.rs | 56 ++++++++++++-- .../arrow/arrow_reader/selection/selector.rs | 74 ------------------- 2 files changed, 50 insertions(+), 80 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index a9a1231d3d6b..39dbb94eba2a 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -33,6 +33,7 @@ use crate::file::page_index::offset_index::PageLocation; use arrow_array::{Array, BooleanArray}; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; use arrow_select::filter::SlicesIterator; +use std::cmp::Ordering; use std::collections::VecDeque; use std::ops::Range; @@ -55,10 +56,7 @@ pub(crate) use cursor::{LoadedRowRanges, MaskCursor, RowSelectionStrategy}; pub use cursor::{RowSelectionCursor, RowSelectionPolicy}; use ranges::{expand_to_batch_boundaries_from_selectors, scan_ranges_from_selectors}; pub use selector::{RowSelectionIter, RowSelector}; -use selector::{ - combine_selectors, limit_selectors, offset_selectors, selectors_from_consecutive_ranges, - split_off_selectors, -}; +use selector::{limit_selectors, offset_selectors, split_off_selectors}; /// [`RowSelection`] represents selecting a subset of rows /// when scanning a parquet file. @@ -330,7 +328,33 @@ impl RowSelection { ranges: I, total_rows: usize, ) -> Self { - Self::from_selectors(selectors_from_consecutive_ranges(ranges, total_rows)) + let mut selectors: Vec = Vec::with_capacity(ranges.size_hint().0); + let mut last_end = 0; + for range in ranges { + let len = range.end - range.start; + if len == 0 { + continue; + } + + match range.start.cmp(&last_end) { + Ordering::Equal => match selectors.last_mut() { + Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(), + None => selectors.push(RowSelector::select(len)), + }, + Ordering::Greater => { + selectors.push(RowSelector::skip(range.start - last_end)); + selectors.push(RowSelector::select(len)) + } + Ordering::Less => panic!("out of order"), + } + last_end = range.end; + } + + if last_end != total_rows { + selectors.push(RowSelector::skip(total_rows - last_end)) + } + + Self::from_selectors(selectors) } /// Given an offset index, return the byte ranges for all data pages selected by `self` @@ -653,7 +677,27 @@ impl From for RowSelection { impl FromIterator for RowSelection { fn from_iter>(iter: T) -> Self { - Self::from_selectors(combine_selectors(iter)) + let iter = iter.into_iter(); + + // Capacity before filter + let mut selectors = Vec::with_capacity(iter.size_hint().0); + + let mut filtered = iter.filter(|x| x.row_count != 0); + if let Some(x) = filtered.next() { + selectors.push(x); + } + + for s in filtered { + // Combine consecutive selectors + let last = selectors.last_mut().unwrap(); + if last.skip == s.skip { + last.row_count = last.row_count.checked_add(s.row_count).unwrap(); + } else { + selectors.push(s) + } + } + + Self::from_selectors(selectors) } } diff --git a/parquet/src/arrow/arrow_reader/selection/selector.rs b/parquet/src/arrow/arrow_reader/selection/selector.rs index 63a7e3b7ddac..15261a5c7a89 100644 --- a/parquet/src/arrow/arrow_reader/selection/selector.rs +++ b/parquet/src/arrow/arrow_reader/selection/selector.rs @@ -24,9 +24,6 @@ //! //! [`RowSelection`]: crate::arrow::arrow_reader::RowSelection -use std::cmp::Ordering; -use std::ops::Range; - /// [`RowSelection`] is a collection of [`RowSelector`] used to skip rows when /// scanning a parquet file /// @@ -111,77 +108,6 @@ impl ExactSizeIterator for RowSelectionIter<'_> {} // once it returns None, it will continue returning None impl std::iter::FusedIterator for RowSelectionIter<'_> {} -/// Normalizes a sequence of selectors: drops the empty ones and combines -/// consecutive selectors that both skip or both select. -pub(super) fn combine_selectors(iter: I) -> Vec -where - I: IntoIterator, -{ - let iter = iter.into_iter(); - - // Capacity before filter - let mut selectors = Vec::with_capacity(iter.size_hint().0); - - let mut filtered = iter.filter(|x| x.row_count != 0); - if let Some(x) = filtered.next() { - selectors.push(x); - } - - for s in filtered { - if s.row_count == 0 { - continue; - } - - // Combine consecutive selectors - let last = selectors.last_mut().unwrap(); - if last.skip == s.skip { - last.row_count = last.row_count.checked_add(s.row_count).unwrap(); - } else { - selectors.push(s) - } - } - - selectors -} - -/// Builds the selectors keeping `ranges` out of `total_rows` rows. -/// -/// # Panics -/// -/// Panics if `ranges` are not in ascending order. -pub(super) fn selectors_from_consecutive_ranges(ranges: I, total_rows: usize) -> Vec -where - I: Iterator>, -{ - let mut selectors: Vec = Vec::with_capacity(ranges.size_hint().0); - let mut last_end = 0; - for range in ranges { - let len = range.end - range.start; - if len == 0 { - continue; - } - - match range.start.cmp(&last_end) { - Ordering::Equal => match selectors.last_mut() { - Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(), - None => selectors.push(RowSelector::select(len)), - }, - Ordering::Greater => { - selectors.push(RowSelector::skip(range.start - last_end)); - selectors.push(RowSelector::select(len)) - } - Ordering::Less => panic!("out of order"), - } - last_end = range.end; - } - - if last_end != total_rows { - selectors.push(RowSelector::skip(total_rows - last_end)) - } - - selectors -} - /// Splits `selectors` at the first `row_count` rows, returning `(head, tail)`. pub(super) fn split_off_selectors( mut selectors: Vec, From bea04f0975c746d936ce035d7f9f334f04582a87 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sat, 25 Jul 2026 20:31:12 +0800 Subject: [PATCH 4/4] update --- parquet/src/arrow/arrow_reader/selection/mod.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index 39dbb94eba2a..b8a5bc0f2a74 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -688,6 +688,10 @@ impl FromIterator for RowSelection { } for s in filtered { + if s.row_count == 0 { + continue; + } + // Combine consecutive selectors let last = selectors.last_mut().unwrap(); if last.skip == s.skip {