Add RowSelection::skipped_row_count

apache · Sep 20, 2024 · e95c4ed · e95c4ed
1 parent 4683c20
commit e95c4ed
Showing 1 changed file with 34 additions and 0 deletions.
diff --git a/parquet/src/arrow/arrow_reader/selection.rs b/parquet/src/arrow/arrow_reader/selection.rs
@@ -436,6 +436,11 @@ impl RowSelection {
     pub fn row_count(&self) -> usize {
         self.iter().filter(|s| !s.skip).map(|s| s.row_count).sum()
     }
+
+    /// Returns the number of de-selected rows
+    pub fn skipped_row_count(&self) -> usize {
+        self.iter().filter(|s| s.skip).map(|s| s.row_count).sum()
+    }
 }
 
 impl From<Vec<RowSelector>> for RowSelection {
@@ -1345,4 +1350,33 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn test_row_count() {
+        let mut 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 mut 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 mut 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 mut selection = RowSelection::from(vec![]);
+
+        assert_eq!(selection.row_count(), 0);
+        assert_eq!(selection.skipped_row_count(), 0);
+    }
 }