haohuaijin commented on code in PR #10141:
URL: https://github.com/apache/arrow-rs/pull/10141#discussion_r3643101412
##########
parquet/src/arrow/arrow_reader/selection/mod.rs:
##########
@@ -135,12 +149,297 @@ impl RowSelector {
/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows
///
/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
-#[derive(Debug, Clone, Default, Eq, PartialEq)]
+#[derive(Default, Clone)]
pub struct RowSelection {
- selectors: Vec<RowSelector>,
+ inner: RowSelectionInner,
+}
+
+/// Internal storage for [`RowSelection`].
+#[derive(Debug, Clone)]
+pub(crate) enum RowSelectionInner {
+ Selectors(Vec<RowSelector>),
+ Mask(Box<MaskSelection>),
+}
+
+impl Default for RowSelectionInner {
+ fn default() -> Self {
+ Self::Selectors(Vec::new())
+ }
+}
+
+impl std::fmt::Debug for RowSelection {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match &self.inner {
+ RowSelectionInner::Selectors(s) => f
+ .debug_struct("RowSelection")
+ .field("selectors", s)
+ .finish(),
+ RowSelectionInner::Mask(m) => f
+ .debug_struct("RowSelection")
+ .field("mask_len", &m.mask().len())
+ .finish_non_exhaustive(),
+ }
+ }
+}
+
+impl PartialEq for RowSelection {
+ fn eq(&self, other: &Self) -> bool {
+ match (&self.inner, &other.inner) {
+ (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b))
=> a == b,
+ (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) =>
a.mask() == b.mask(),
+ (RowSelectionInner::Mask(mask),
RowSelectionInner::Selectors(selectors))
+ | (RowSelectionInner::Selectors(selectors),
RowSelectionInner::Mask(mask)) => {
+ if selectors
+ .iter()
+ .try_fold(0usize, |acc, selector|
acc.checked_add(selector.row_count))
+ != Some(mask.mask().len())
+ {
+ return false;
+ }
+
+ let mut slices = mask.mask().set_slices().peekable();
+ let mut cursor = 0usize;
+
+ for selector in selectors {
+ let end = cursor + selector.row_count;
+
+ if selector.skip {
+ if slices.peek().is_some_and(|(start, _)| *start <
end) {
+ return false;
+ }
+ } else {
+ match slices.next() {
+ Some((start, slice_end)) if start == cursor &&
slice_end == end => {}
+ _ => return false,
+ }
+ }
+
+ cursor = end;
+ }
+
+ slices.next().is_none()
+ }
+ }
+ }
+}
+
+impl Eq for RowSelection {}
+
+/// Borrowed iterator over the [`RowSelector`]s of a [`RowSelection`].
+#[derive(Debug)]
+pub struct RowSelectionIter<'a>(std::slice::Iter<'a, RowSelector>);
+
+impl<'a> Iterator for RowSelectionIter<'a> {
+ type Item = &'a RowSelector;
+
+ #[inline]
+ fn next(&mut self) -> Option<Self::Item> {
Review Comment:
apply in
[86610b5](https://github.com/apache/arrow-rs/pull/10141/commits/86610b5bbb3a2a91c7cae070eddbb57fff24e015)
and
[bdf2726](https://github.com/apache/arrow-rs/pull/10141/commits/bdf2726c21930d5264b49c751c9b5c8f0fc569c4)
--
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]
