Add test for page-level limit pruning

Author: xudong963

datafusion/core/tests/parquet/page_pruning.rs

@@ -17,6 +17,7 @@
 
 use std::sync::Arc;
 
+use crate::parquet::utils::MetricsFinder;
 use crate::parquet::Unit::Page;
 use crate::parquet::{ContextWithParquet, Scenario};
 
@@ -45,6 +46,7 @@ async fn get_parquet_exec(
     state: &SessionState,
     filter: Expr,
     pushdown_filters: bool,
+    limit: Option<usize>,
 ) -> DataSourceExec {
     let object_store_url = ObjectStoreUrl::local_filesystem();
     let store = state.runtime_env().object_store(&object_store_url).unwrap();
@@ -88,6 +90,7 @@ async fn get_parquet_exec(
     );
     let base_config = FileScanConfigBuilder::new(object_store_url, schema, source)
         .with_file(partitioned_file)
+        .with_limit(limit)
         .build();
 
     DataSourceExec::new(Arc::new(base_config))
@@ -98,7 +101,7 @@ async fn get_filter_results(
     filter: Expr,
     pushdown_filters: bool,
 ) -> Vec<RecordBatch> {
-    let parquet_exec = get_parquet_exec(state, filter, pushdown_filters).await;
+    let parquet_exec = get_parquet_exec(state, filter, pushdown_filters, None).await;
     let task_ctx = state.task_ctx();
     let mut results = parquet_exec.execute(0, task_ctx.clone()).unwrap();
     let mut batches = Vec::new();
@@ -108,6 +111,29 @@ async fn get_filter_results(
     batches
 }
 
+#[tokio::test]
+async fn limit_pruning_filter_one_col() {
+    let session_ctx = SessionContext::new();
+    let state = session_ctx.state();
+    let filter = col("month").eq(lit(1_i32));
+
+    let parquet_exec = get_parquet_exec(&state, filter, false, Some(10)).await;
+    let task_ctx = state.task_ctx();
+    let mut results = parquet_exec.execute(0, task_ctx).unwrap();
+    let mut limited_rows = 0;
+    while let Some(Ok(batch)) = results.next().await {
+        limited_rows += batch.num_rows();
+    }
+
+    assert_eq!(limited_rows, 10);
+    let metrics = MetricsFinder::find_metrics(&parquet_exec).unwrap();
+    if let Some(matched_rows) =
+        cast_count_metric(metrics.sum_by_name("limit_pruning_matched_rows").unwrap())
+    {
+        assert_eq!(matched_rows, 10);
+    }
+}
+
 #[tokio::test]
 async fn page_index_filter_one_col() {
     let session_ctx = SessionContext::new();

datafusion/datasource-parquet/src/metrics.rs

@@ -50,6 +50,8 @@ pub struct ParquetFileMetrics {
     pub row_groups_pruned_bloom_filter: Count,
     /// Number of row groups pruned due to limit pruning.
     pub limit_pruned_row_groups: Count,
+    /// Number of rows matched by limit pruning.
+    pub limit_pruning_matched_rows: Count,
     /// Number of row groups whose statistics were checked and matched (not pruned)
     pub row_groups_matched_statistics: Count,
     /// Number of row groups pruned by statistics
@@ -99,6 +101,10 @@ impl ParquetFileMetrics {
             .with_new_label("filename", filename.to_string())
             .counter("limit_pruned_row_groups", partition);
 
+        let limit_pruning_matched_rows = MetricBuilder::new(metrics)
+            .with_new_label("filename", filename.to_string())
+            .counter("limit_pruning_matched_rows", partition);
+
         let row_groups_matched_statistics = MetricBuilder::new(metrics)
             .with_new_label("filename", filename.to_string())
             .counter("row_groups_matched_statistics", partition);
@@ -154,6 +160,7 @@ impl ParquetFileMetrics {
             row_groups_matched_statistics,
             row_groups_pruned_statistics,
             limit_pruned_row_groups,
+            limit_pruning_matched_rows,
             bytes_scanned,
             pushdown_rows_pruned,
             pushdown_rows_matched,

datafusion/datasource-parquet/src/opener.rs

@@ -395,7 +395,6 @@ impl FileOpener for ParquetOpener {
                         );
 
                     access_plan = new_access_plan;
-
                     // Apply page-level limit pruning if limit is specified
                     if let Some(limit) = limit {
                         access_plan = p.prune_by_limit(
@@ -411,7 +410,7 @@ impl FileOpener for ParquetOpener {
 
             let row_group_indexes = access_plan.row_group_indexes();
             if let Some(row_selection) =
-                access_plan.into_overall_row_selection(rg_metadata)?
+                access_plan.into_overall_row_selection(rg_metadata).unwrap()
             {
                 builder = builder.with_row_selection(row_selection);
             }

datafusion/datasource-parquet/src/page_filter.rs

@@ -115,11 +115,6 @@ pub struct PagePruningAccessPlanFilter {
     /// single column predicates (e.g. (`col = 5`) extracted from the overall
     /// predicate. Must all be true for a row to be included in the result.
     predicates: Vec<PruningPredicate>,
-    /// For each row group, tracks which pages are fully matched (all rows satisfy predicates)
-    /// Key: row group index, Value: Vec of booleans (one per page)
-    fully_matched_pages: HashMap<usize, Vec<bool>>,
-    /// For each row group, stores the row counts per page
-    page_row_counts: HashMap<usize, Vec<usize>>,
 }
 
 impl PagePruningAccessPlanFilter {
@@ -154,11 +149,7 @@ impl PagePruningAccessPlanFilter {
                 Some(pp)
             })
             .collect::<Vec<_>>();
-        Self {
-            predicates,
-            fully_matched_pages: HashMap::new(),
-            page_row_counts: HashMap::new(),
-        }
+        Self { predicates }
     }
 
     /// Returns an updated [`ParquetAccessPlan`] by applying predicates to the
@@ -337,7 +328,7 @@ impl PagePruningAccessPlanFilter {
         rg_metadata: &[RowGroupMetaData],
         file_metrics: &ParquetFileMetrics,
     ) -> ParquetAccessPlan {
-        if self.fully_matched_pages.is_empty() {
+        if page_match_infos.is_empty() {
             return access_plan;
         }
 
@@ -388,15 +379,16 @@ impl PagePruningAccessPlanFilter {
                 }
 
                 // Build a RowSelection for this row group that includes only the fully matched pages
-                let page_row_counts = self.page_row_counts.get(&rg_idx).unwrap();
-                let row_selection = build_row_selection_for_pages(
-                    &page_indices,
-                    &page_row_counts,
-                    limit - rows_selected,
-                );
-
-                new_access_plan.scan_selection(rg_idx, row_selection);
-                rows_selected += std::cmp::min(row_count, limit - rows_selected);
+                if let Some(page_match_info) = page_match_infos.get(&rg_idx) {
+                    let row_selection = build_row_selection_for_pages(
+                        &page_indices,
+                        &page_match_info.page_row_counts,
+                        limit - rows_selected,
+                    );
+                    new_access_plan.scan(rg_idx);
+                    new_access_plan.scan_selection(rg_idx, row_selection);
+                    rows_selected += std::cmp::min(row_count, limit - rows_selected);
+                }
             }
 
             let original_row_groups = access_plan.row_group_indexes().len();
@@ -406,7 +398,7 @@ impl PagePruningAccessPlanFilter {
             if pruned_row_groups > 0 {
                 file_metrics.limit_pruned_row_groups.add(pruned_row_groups);
             }
-
+            file_metrics.limit_pruning_matched_rows.add(rows_selected);
             return new_access_plan;
         }
 
@@ -490,7 +482,10 @@ fn build_row_selection_for_pages(
     let mut page_idx_iter = page_indices.iter().peekable();
 
     for (page_num, &page_row_count) in page_row_counts.iter().enumerate() {
-        if let Some(&&next_page_idx) = page_idx_iter.peek() {
+        if rows_selected >= limit {
+            // Once we've reached the limit, skip all remaining pages
+            current_row += page_row_count;
+        } else if let Some(&&next_page_idx) = page_idx_iter.peek() {
             if page_num == next_page_idx {
                 // This page should be selected
                 page_idx_iter.next();
@@ -510,10 +505,6 @@ fn build_row_selection_for_pages(
                 if rows_to_select < page_row_count {
                     current_row = page_row_count - rows_to_select;
                 }
-
-                if rows_selected >= limit {
-                    break;
-                }
             } else {
                 // This page should be skipped
                 current_row += page_row_count;
@@ -876,7 +867,7 @@ mod tests {
             build_row_selection_for_pages(&page_indices, &page_row_counts, limit);
         let result = row_selection_to_vec(&selection);
 
-        assert_eq!(result, vec![(true, 200), (false, 50)]);
+        assert_eq!(result, vec![(true, 200), (false, 250)]);
     }
 
     #[test]
@@ -890,6 +881,6 @@ mod tests {
             build_row_selection_for_pages(&page_indices, &page_row_counts, limit);
         let result = row_selection_to_vec(&selection);
 
-        assert_eq!(result, vec![(false, 50), (true, 150), (false, 30),]);
+        assert_eq!(result, vec![(false, 50), (true, 150), (false, 210),]);
     }
 }