alamb commented on code in PR #10141:
URL: https://github.com/apache/arrow-rs/pull/10141#discussion_r3462896701
##########
parquet/src/arrow/arrow_reader/selection.rs:
##########
@@ -135,12 +143,439 @@ impl RowSelector {
/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows
///
/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
-#[derive(Debug, Clone, Default, Eq, PartialEq)]
+#[derive(Default, Clone)]
pub struct RowSelection {
- selectors: Vec<RowSelector>,
+ inner: RowSelectionInner,
+}
+
+/// Internal storage for [`RowSelection`].
+#[derive(Debug, Clone)]
+pub(crate) enum RowSelectionInner {
+ Selectors(Vec<RowSelector>),
+ Mask(BooleanBuffer),
+}
+
+impl Default for RowSelectionInner {
+ fn default() -> Self {
+ Self::Selectors(Vec::new())
+ }
+}
+
+impl std::fmt::Debug for RowSelection {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match &self.inner {
+ RowSelectionInner::Selectors(s) => f
+ .debug_struct("RowSelection")
+ .field("selectors", s)
+ .finish(),
+ RowSelectionInner::Mask(m) => f
+ .debug_struct("RowSelection")
+ .field("mask_len", &m.len())
+ .finish_non_exhaustive(),
+ }
+ }
+}
+
+impl PartialEq for RowSelection {
+ fn eq(&self, other: &Self) -> bool {
+ match (&self.inner, &other.inner) {
+ (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b))
=> a == b,
+ (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a == b,
+ (RowSelectionInner::Mask(mask),
RowSelectionInner::Selectors(selectors))
+ | (RowSelectionInner::Selectors(selectors),
RowSelectionInner::Mask(mask)) => {
+ if selectors
+ .iter()
+ .try_fold(0usize, |acc, selector|
acc.checked_add(selector.row_count))
+ != Some(mask.len())
+ {
+ return false;
+ }
+
+ let mut slices = mask.set_slices().peekable();
+ let mut cursor = 0usize;
+
+ for selector in selectors {
+ let end = cursor + selector.row_count;
+
+ if selector.skip {
+ if slices.peek().is_some_and(|(start, _)| *start <
end) {
+ return false;
+ }
+ } else {
+ match slices.next() {
+ Some((start, slice_end)) if start == cursor &&
slice_end == end => {}
+ _ => return false,
+ }
+ }
+
+ cursor = end;
+ }
+
+ slices.next().is_none()
+ }
+ }
+ }
+}
+
+impl Eq for RowSelection {}
+
+/// Iterator over the [`RowSelector`]s of a [`RowSelection`].
+///
+/// Yields owned [`RowSelector`] values; mask-backed selections stream the
+/// run-length form via [`BitSliceIterator`] without allocation.
+#[derive(Debug)]
+pub enum RowSelectionIter<'a> {
+ Selectors(std::iter::Copied<std::slice::Iter<'a, RowSelector>>),
+ Mask(MaskRunIter<'a>),
+}
+
+impl Iterator for RowSelectionIter<'_> {
+ type Item = RowSelector;
+
+ #[inline]
+ fn next(&mut self) -> Option<RowSelector> {
+ match self {
+ Self::Selectors(it) => it.next(),
+ Self::Mask(it) => it.next(),
+ }
+ }
+}
+
+/// Streaming RLE view of a [`BooleanBuffer`].
+#[derive(Debug)]
+pub struct MaskRunIter<'a> {
+ slices: BitSliceIterator<'a>,
+ cursor: usize,
+ total: usize,
+ pending: Option<RowSelector>,
+ finished: bool,
+}
+
+impl<'a> MaskRunIter<'a> {
+ fn new(mask: &'a BooleanBuffer) -> Self {
+ Self {
+ slices: mask.set_slices(),
+ cursor: 0,
+ total: mask.len(),
+ pending: None,
+ finished: false,
+ }
+ }
+}
+
+impl Iterator for MaskRunIter<'_> {
+ type Item = RowSelector;
+
+ fn next(&mut self) -> Option<RowSelector> {
+ if let Some(p) = self.pending.take() {
+ return Some(p);
+ }
+ if self.finished {
+ return None;
+ }
+ match self.slices.next() {
Review Comment:
Much of this code appears untested
I ran code coverage like
```shell
cargo llvm-cov --html test -p parquet --features arrow
```
<img width="1111" height="942" alt="Image"
src="https://github.com/user-attachments/assets/bbc86094-2da0-4240-bf32-0d94af99a87e";
/>
##########
parquet/src/arrow/arrow_reader/selection.rs:
##########
@@ -135,12 +143,439 @@ impl RowSelector {
/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows
///
/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
-#[derive(Debug, Clone, Default, Eq, PartialEq)]
+#[derive(Default, Clone)]
pub struct RowSelection {
- selectors: Vec<RowSelector>,
+ inner: RowSelectionInner,
+}
+
+/// Internal storage for [`RowSelection`].
+#[derive(Debug, Clone)]
+pub(crate) enum RowSelectionInner {
+ Selectors(Vec<RowSelector>),
+ Mask(BooleanBuffer),
+}
+
+impl Default for RowSelectionInner {
+ fn default() -> Self {
+ Self::Selectors(Vec::new())
+ }
+}
+
+impl std::fmt::Debug for RowSelection {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match &self.inner {
+ RowSelectionInner::Selectors(s) => f
+ .debug_struct("RowSelection")
+ .field("selectors", s)
+ .finish(),
+ RowSelectionInner::Mask(m) => f
+ .debug_struct("RowSelection")
+ .field("mask_len", &m.len())
+ .finish_non_exhaustive(),
+ }
+ }
+}
+
+impl PartialEq for RowSelection {
+ fn eq(&self, other: &Self) -> bool {
+ match (&self.inner, &other.inner) {
+ (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b))
=> a == b,
+ (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a == b,
+ (RowSelectionInner::Mask(mask),
RowSelectionInner::Selectors(selectors))
+ | (RowSelectionInner::Selectors(selectors),
RowSelectionInner::Mask(mask)) => {
+ if selectors
+ .iter()
+ .try_fold(0usize, |acc, selector|
acc.checked_add(selector.row_count))
+ != Some(mask.len())
+ {
+ return false;
+ }
+
+ let mut slices = mask.set_slices().peekable();
+ let mut cursor = 0usize;
+
+ for selector in selectors {
+ let end = cursor + selector.row_count;
+
+ if selector.skip {
+ if slices.peek().is_some_and(|(start, _)| *start <
end) {
+ return false;
+ }
+ } else {
+ match slices.next() {
+ Some((start, slice_end)) if start == cursor &&
slice_end == end => {}
+ _ => return false,
+ }
+ }
+
+ cursor = end;
+ }
+
+ slices.next().is_none()
+ }
+ }
+ }
+}
+
+impl Eq for RowSelection {}
+
+/// Iterator over the [`RowSelector`]s of a [`RowSelection`].
+///
+/// Yields owned [`RowSelector`] values; mask-backed selections stream the
+/// run-length form via [`BitSliceIterator`] without allocation.
+#[derive(Debug)]
+pub enum RowSelectionIter<'a> {
+ Selectors(std::iter::Copied<std::slice::Iter<'a, RowSelector>>),
+ Mask(MaskRunIter<'a>),
+}
+
+impl Iterator for RowSelectionIter<'_> {
+ type Item = RowSelector;
+
+ #[inline]
+ fn next(&mut self) -> Option<RowSelector> {
+ match self {
Review Comment:
In general I think this will add a branch in the main loops - every chunk of
bits will require matching -- it might make sense to figure out how to have
iterators for the two different types, so that code paths that want to iterate
over the selection can match outside
```rust
match iter {
RowSelectionIter::Selectors(iter) => { for i in iter ... }, // special
loop for selector backed
RowSelectionIter::Mask(iter) => { for i in iter ... }, // special loop
for mask backed
}
```
##########
parquet/src/arrow/push_decoder/reader_builder/mod.rs:
##########
@@ -916,7 +916,7 @@ mod tests {
#[test]
// Verify that the size of RowGroupDecoderState does not grow too large
fn test_structure_size() {
- assert_eq!(std::mem::size_of::<RowGroupDecoderState>(), 232);
+ assert_eq!(std::mem::size_of::<RowGroupDecoderState>(), 256);
Review Comment:
I think this is bigger because a `RowSelection` is now bigger.
##########
parquet/src/arrow/arrow_reader/selection.rs:
##########
@@ -412,7 +868,12 @@ impl RowSelection {
///
/// returned: NNNNNNNNYYNYN
pub fn intersection(&self, other: &Self) -> Self {
- intersect_row_selections(&self.selectors, &other.selectors)
+ if let (Some(l), Some(r)) = (self.as_mask(), other.as_mask()) {
+ return Self::from_boolean_buffer(intersect_masks(l, r));
+ }
+ let l = self.materialize_for_combine();
Review Comment:
for example, I think and_then can use some of the fancy BMI instructions
that @devanbenz is investigating in
https://github.com/apache/arrow-rs/pull/10136
##########
parquet/src/arrow/arrow_reader/selection.rs:
##########
@@ -135,12 +143,439 @@ impl RowSelector {
/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows
///
/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
-#[derive(Debug, Clone, Default, Eq, PartialEq)]
+#[derive(Default, Clone)]
pub struct RowSelection {
- selectors: Vec<RowSelector>,
+ inner: RowSelectionInner,
+}
+
+/// Internal storage for [`RowSelection`].
+#[derive(Debug, Clone)]
+pub(crate) enum RowSelectionInner {
+ Selectors(Vec<RowSelector>),
+ Mask(BooleanBuffer),
+}
+
+impl Default for RowSelectionInner {
+ fn default() -> Self {
+ Self::Selectors(Vec::new())
+ }
+}
+
+impl std::fmt::Debug for RowSelection {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match &self.inner {
+ RowSelectionInner::Selectors(s) => f
+ .debug_struct("RowSelection")
+ .field("selectors", s)
+ .finish(),
+ RowSelectionInner::Mask(m) => f
+ .debug_struct("RowSelection")
+ .field("mask_len", &m.len())
+ .finish_non_exhaustive(),
+ }
+ }
+}
+
+impl PartialEq for RowSelection {
+ fn eq(&self, other: &Self) -> bool {
+ match (&self.inner, &other.inner) {
+ (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b))
=> a == b,
+ (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a == b,
+ (RowSelectionInner::Mask(mask),
RowSelectionInner::Selectors(selectors))
+ | (RowSelectionInner::Selectors(selectors),
RowSelectionInner::Mask(mask)) => {
+ if selectors
+ .iter()
+ .try_fold(0usize, |acc, selector|
acc.checked_add(selector.row_count))
+ != Some(mask.len())
+ {
+ return false;
+ }
+
+ let mut slices = mask.set_slices().peekable();
+ let mut cursor = 0usize;
+
+ for selector in selectors {
+ let end = cursor + selector.row_count;
+
+ if selector.skip {
+ if slices.peek().is_some_and(|(start, _)| *start <
end) {
+ return false;
+ }
+ } else {
+ match slices.next() {
+ Some((start, slice_end)) if start == cursor &&
slice_end == end => {}
+ _ => return false,
+ }
+ }
+
+ cursor = end;
+ }
+
+ slices.next().is_none()
+ }
+ }
+ }
+}
+
+impl Eq for RowSelection {}
+
+/// Iterator over the [`RowSelector`]s of a [`RowSelection`].
+///
+/// Yields owned [`RowSelector`] values; mask-backed selections stream the
+/// run-length form via [`BitSliceIterator`] without allocation.
+#[derive(Debug)]
+pub enum RowSelectionIter<'a> {
+ Selectors(std::iter::Copied<std::slice::Iter<'a, RowSelector>>),
+ Mask(MaskRunIter<'a>),
+}
+
+impl Iterator for RowSelectionIter<'_> {
+ type Item = RowSelector;
+
+ #[inline]
+ fn next(&mut self) -> Option<RowSelector> {
+ match self {
+ Self::Selectors(it) => it.next(),
+ Self::Mask(it) => it.next(),
+ }
+ }
+}
+
+/// Streaming RLE view of a [`BooleanBuffer`].
+#[derive(Debug)]
+pub struct MaskRunIter<'a> {
+ slices: BitSliceIterator<'a>,
+ cursor: usize,
+ total: usize,
+ pending: Option<RowSelector>,
+ finished: bool,
+}
+
+impl<'a> MaskRunIter<'a> {
+ fn new(mask: &'a BooleanBuffer) -> Self {
+ Self {
+ slices: mask.set_slices(),
+ cursor: 0,
+ total: mask.len(),
+ pending: None,
+ finished: false,
+ }
+ }
+}
+
+impl Iterator for MaskRunIter<'_> {
+ type Item = RowSelector;
+
+ fn next(&mut self) -> Option<RowSelector> {
Review Comment:
I think this is the main API that is not compatible -- it need to return
`&RowSelector` -- if we could change that we can merge this PR before the next
major release (with breaking API changes)
##########
parquet/src/arrow/arrow_reader/selection.rs:
##########
@@ -412,7 +868,12 @@ impl RowSelection {
///
/// returned: NNNNNNNNYYNYN
pub fn intersection(&self, other: &Self) -> Self {
- intersect_row_selections(&self.selectors, &other.selectors)
+ if let (Some(l), Some(r)) = (self.as_mask(), other.as_mask()) {
+ return Self::from_boolean_buffer(intersect_masks(l, r));
+ }
+ let l = self.materialize_for_combine();
Review Comment:
this is a clever idea -- as a follow on we could potentially implement
specialized versions of each of these implementations
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
