Trim parquet row selection (#2705)

apache · Sep 24, 2022 · 7c8080c · 7c8080c
1 parent dc62e40
commit 7c8080c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs
@@ -570,7 +570,7 @@ impl ParquetRecordBatchReader {
             batch_size,
             array_reader,
             schema: Arc::new(schema),
-            selection: selection.map(Into::into),
+            selection: selection.map(|s| s.trim().into()),
         }
     }
 }

diff --git a/parquet/src/arrow/arrow_reader/selection.rs b/parquet/src/arrow/arrow_reader/selection.rs
@@ -285,6 +285,14 @@ impl RowSelection {
     pub fn selects_any(&self) -> bool {
         self.selectors.iter().any(|x| !x.skip)
     }
+
+    /// Trims this [`RowSelection`] removing any trailing skips
+    pub(crate) fn trim(mut self) -> Self {
+        while self.selectors.last().map(|x| x.skip).unwrap_or(false) {
+            self.selectors.pop();
+        }
+        self
+    }
 }
 
 impl From<Vec<RowSelector>> for RowSelection {