extract some logic into identify_fully_matched_row_groups

Author: xudong963

datafusion/datasource-parquet/src/row_group_filter.rs

@@ -84,6 +84,9 @@ impl RowGroupAccessPlanFilter {
     }
 
     /// Prunes the access plan based on the limit and fully contained row groups.
+    /// See the [description](https://github.com/apache/datafusion/issues/18860#issuecomment-3563442093)
+    /// for how the pruning works and improves performance.
+    /// For more information, see the [paper](https://arxiv.org/pdf/2504.11540)'s "Pruning for LIMIT Queries" part
     pub fn prune_by_limit(
         &mut self,
         limit: usize,
@@ -197,47 +200,15 @@ impl RowGroupAccessPlanFilter {
                     }
                 }
 
-                // Note: this part of code shouldn't be expensive with a limited number of row groups
-                // If we do find it's expensive, we can consider optimizing it further.
-                if !fully_contained_candidates_original_idx.is_empty() {
-                    // Use NotExpr to create the inverted predicate
-                    let inverted_expr =
-                        Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr())));
-                    // Simplify the NOT expression (e.g., NOT(c1 = 0) -> c1 != 0)
-                    // before building the pruning predicate
-                    let mut simplifier = PhysicalExprSimplifier::new(arrow_schema);
-                    let inverted_expr = simplifier.simplify(inverted_expr).unwrap();
-                    if let Ok(inverted_predicate) = PruningPredicate::try_new(
-                        inverted_expr,
-                        Arc::clone(predicate.schema()),
-                    ) {
-                        let inverted_pruning_stats = RowGroupPruningStatistics {
-                            parquet_schema,
-                            row_group_metadatas: fully_contained_candidates_original_idx
-                                .iter()
-                                .map(|&i| &groups[i])
-                                .collect::<Vec<_>>(),
-                            arrow_schema,
-                        };
-
-                        if let Ok(inverted_values) =
-                            inverted_predicate.prune(&inverted_pruning_stats)
-                        {
-                            for (i, &original_row_group_idx) in
-                                fully_contained_candidates_original_idx.iter().enumerate()
-                            {
-                                // If the inverted predicate *also* prunes this row group (meaning inverted_values[i] is false),
-                                // it implies that *all* rows in this group satisfy the original predicate.
-                                if !inverted_values[i] {
-                                    self.is_fully_matched[original_row_group_idx] = true;
-                                    metrics
-                                        .row_groups_pruned_statistics
-                                        .add_fully_matched(1);
-                                }
-                            }
-                        }
-                    }
-                }
+                // Check if any of the matched row groups are fully contained by the predicate
+                self.identify_fully_matched_row_groups(
+                    fully_contained_candidates_original_idx,
+                    arrow_schema,
+                    parquet_schema,
+                    groups,
+                    predicate,
+                    metrics,
+                );
             }
             // stats filter array could not be built, so we can't prune
             Err(e) => {
@@ -247,6 +218,68 @@ impl RowGroupAccessPlanFilter {
         }
     }
 
+    /// Identifies row groups that are fully matched by the predicate.
+    ///
+    /// This optimization checks whether all rows in a row group satisfy the predicate
+    /// by inverting the predicate and checking if it prunes the row group. If the
+    /// inverted predicate prunes a row group, it means no rows match the inverted
+    /// predicate, which implies all rows match the original predicate.
+    ///
+    /// Note: This optimization is relatively inexpensive for a limited number of row groups.
+    fn identify_fully_matched_row_groups(
+        &mut self,
+        candidate_row_group_indices: Vec<usize>,
+        arrow_schema: &Schema,
+        parquet_schema: &SchemaDescriptor,
+        groups: &[RowGroupMetaData],
+        predicate: &PruningPredicate,
+        metrics: &ParquetFileMetrics,
+    ) {
+        if candidate_row_group_indices.is_empty() {
+            return;
+        }
+
+        // Use NotExpr to create the inverted predicate
+        let inverted_expr = Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr())));
+
+        // Simplify the NOT expression (e.g., NOT(c1 = 0) -> c1 != 0)
+        // before building the pruning predicate
+        let mut simplifier = PhysicalExprSimplifier::new(arrow_schema);
+        let Ok(inverted_expr) = simplifier.simplify(inverted_expr) else {
+            return;
+        };
+
+        let Ok(inverted_predicate) =
+            PruningPredicate::try_new(inverted_expr, Arc::clone(predicate.schema()))
+        else {
+            return;
+        };
+
+        let inverted_pruning_stats = RowGroupPruningStatistics {
+            parquet_schema,
+            row_group_metadatas: candidate_row_group_indices
+                .iter()
+                .map(|&i| &groups[i])
+                .collect::<Vec<_>>(),
+            arrow_schema,
+        };
+
+        let Ok(inverted_values) = inverted_predicate.prune(&inverted_pruning_stats)
+        else {
+            return;
+        };
+
+        for (i, &original_row_group_idx) in candidate_row_group_indices.iter().enumerate()
+        {
+            // If the inverted predicate *also* prunes this row group (meaning inverted_values[i] is false),
+            // it implies that *all* rows in this group satisfy the original predicate.
+            if !inverted_values[i] {
+                self.is_fully_matched[original_row_group_idx] = true;
+                metrics.row_groups_pruned_statistics.add_fully_matched(1);
+            }
+        }
+    }
+
     /// Prune remaining row groups using available bloom filters and the
     /// [`PruningPredicate`].
     ///