refactor

apache · May 21, 2024 · 45d4813 · 45d4813
1 parent d64ecf6
commit 45d4813
Showing 1 changed file with 37 additions and 25 deletions.
diff --git a/datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs b/datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs
@@ -1282,6 +1282,42 @@ mod tests {
         NonePruned,
     }
 
+    impl ExpectedPruning {
+        /// asserts that the pruned row group match this expectation
+        fn assert(&self, row_groups: &RowGroupSet) {
+            let num_row_groups = row_groups.len();
+            assert!(num_row_groups > 0);
+            let num_pruned = (0..num_row_groups)
+                .filter_map(|i| {
+                    if row_groups.should_scan(i) {
+                        None
+                    } else {
+                        Some(1)
+                    }
+                })
+                .sum::<usize>();
+
+            match self {
+                Self::AllPruned => {
+                    assert_eq!(
+                        num_row_groups, num_pruned,
+                        "Expected all row groups to be pruned, but got {row_groups:?}"
+                    );
+                }
+                ExpectedPruning::NonePruned => {
+                    assert_eq!(
+                        num_pruned, 0,
+                        "Expected no row groups to be pruned, but got {row_groups:?}"
+                    );
+                }
+                ExpectedPruning::SomePruned(expected) => {
+                    let actual = row_groups.indexes();
+                    assert_eq!(expected, &actual, "Unexpected row groups pruned. Expected {expected:?}, got {actual:?}");
+                }
+            }
+        }
+    }
+
     struct BloomFilterTest {
         file_name: String,
         schema: Schema,
@@ -1370,30 +1406,7 @@ mod tests {
             .await
             .unwrap();
 
-            let num_row_groups = pruned_row_groups.len();
-            assert!(num_row_groups > 0);
-            let num_pruned = (0..num_row_groups)
-                .filter_map(|i| {
-                    if pruned_row_groups.should_scan(i) {
-                        None
-                    } else {
-                        Some(1)
-                    }
-                })
-                .sum::<usize>();
-
-            match post_pruning_row_groups {
-                ExpectedPruning::AllPruned => {
-                    assert_eq!(num_row_groups, num_pruned, "Expected all row groups to be pruned, but got {pruned_row_groups:?}");
-                }
-                ExpectedPruning::NonePruned => {
-                    assert_eq!(num_pruned, 0, "Expected no row groups to be pruned, but got {pruned_row_groups:?}");
-                }
-                ExpectedPruning::SomePruned(expected) => {
-                    let actual = pruned_row_groups.indexes();
-                    assert_eq!(expected, actual, "Unexpected row groups pruned. Expected {expected:?}, got {actual:?}");
-                }
-            }
+            post_pruning_row_groups.assert(&pruned_row_groups);
         }
     }
 
@@ -1427,7 +1440,6 @@ mod tests {
         };
         let mut builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap();
 
-        let metadata = builder.metadata().clone();
         let pruned_row_group = prune_row_groups_by_bloom_filters(
             pruning_predicate.schema(),
             &mut builder,