Add fetch_order_sensitive during limit pushdown to decide if use limit pruning

Author: xudong963

datafusion/catalog-listing/src/table.rs

@@ -506,6 +506,7 @@ impl TableProvider for ListingTable {
                     .with_statistics(statistics)
                     .with_projection_indices(projection)
                     .with_limit(limit)
+                    .with_limit_order_sensitive(args.limit_order_sensitive())
                     .with_output_ordering(output_ordering)
                     .with_expr_adapter(self.expr_adapter_factory.clone())
                     .build(),

datafusion/catalog/src/table.rs

@@ -336,6 +336,7 @@ pub struct ScanArgs<'a> {
     filters: Option<&'a [Expr]>,
     projection: Option<&'a [usize]>,
     limit: Option<usize>,
+    limit_order_sensitive: bool,
 }
 
 impl<'a> ScanArgs<'a> {
@@ -397,6 +398,25 @@ impl<'a> ScanArgs<'a> {
     pub fn limit(&self) -> Option<usize> {
         self.limit
     }
+
+    /// Set whether the scan's limit should be order-sensitive.
+    ///
+    /// If specified, the scan should return the limited rows in a specific order.
+    /// Or we can leverage limit pruning to optimize the scan.
+    ///
+    /// # Arguments
+    /// * `order_sensitive` - Whether the scan's limit should be order-sensitive
+    pub fn with_limit_order_sensitive(mut self, order_sensitive: bool) -> Self {
+        self.limit_order_sensitive = order_sensitive;
+        self
+    }
+
+    /// Get whether the scan's limit should be order-sensitive.
+    ///
+    /// Returns `true` if the scan's limit should be order-sensitive, or `false` if not.
+    pub fn limit_order_sensitive(&self) -> bool {
+        self.limit_order_sensitive
+    }
 }
 
 /// Result of a table scan operation from [`TableProvider::scan_with_args`].

datafusion/core/src/physical_planner.rs

@@ -457,6 +457,7 @@ impl DefaultPhysicalPlanner {
                 projection,
                 filters,
                 fetch,
+                fetch_order_sensitive,
                 ..
             }) => {
                 let source = source_as_provider(source)?;
@@ -468,7 +469,8 @@ impl DefaultPhysicalPlanner {
                 let opts = ScanArgs::default()
                     .with_projection(projection.as_deref())
                     .with_filters(Some(&filters_vec))
-                    .with_limit(*fetch);
+                    .with_limit(*fetch)
+                    .with_limit_order_sensitive(*fetch_order_sensitive);
                 let res = source.scan_with_args(session_state, opts).await?;
                 Arc::clone(res.plan())
             }

datafusion/core/tests/parquet/mod.rs

@@ -126,7 +126,7 @@ struct TestOutput {
 impl TestOutput {
     /// retrieve the value of the named metric, if any
     fn metric_value(&self, metric_name: &str) -> Option<usize> {
-        if let Some((pruned, _matched)) = self.pruning_metric(metric_name) {
+        if let Some((pruned, _matched, _fully)) = self.pruning_metric(metric_name) {
             return Some(pruned);
         }
 
@@ -140,9 +140,10 @@ impl TestOutput {
             })
     }
 
-    fn pruning_metric(&self, metric_name: &str) -> Option<(usize, usize)> {
+    fn pruning_metric(&self, metric_name: &str) -> Option<(usize, usize, usize)> {
         let mut total_pruned = 0;
         let mut total_matched = 0;
+        let mut total_fully_matched = 0;
         let mut found = false;
 
         for metric in self.parquet_metrics.iter() {
@@ -154,13 +155,15 @@ impl TestOutput {
                 {
                     total_pruned += pruning_metrics.pruned();
                     total_matched += pruning_metrics.matched();
+                    total_fully_matched += pruning_metrics.fully_matched();
+
                     found = true;
                 }
             }
         }
 
         if found {
-            Some((total_pruned, total_matched))
+            Some((total_pruned, total_matched, total_fully_matched))
         } else {
             None
         }
@@ -172,32 +175,33 @@ impl TestOutput {
     }
 
     /// The number of row_groups pruned / matched by bloom filter
-    fn row_groups_bloom_filter(&self) -> Option<(usize, usize)> {
+    fn row_groups_bloom_filter(&self) -> Option<(usize, usize, usize)> {
         self.pruning_metric("row_groups_pruned_bloom_filter")
     }
 
     /// The number of row_groups matched by statistics
     fn row_groups_matched_statistics(&self) -> Option<usize> {
         self.pruning_metric("row_groups_pruned_statistics")
-            .map(|(_pruned, matched)| matched)
+            .map(|(_pruned, matched, _fully)| matched)
     }
 
     /// The number of row_groups fully matched by statistics
     fn row_groups_fully_matched_statistics(&self) -> Option<usize> {
-        self.metric_value("row_groups_fully_matched_statistics")
+        self.pruning_metric("row_groups_pruned_statistics")
+            .map(|(_pruned, _, fully)| fully)
     }
 
     /// The number of row_groups pruned by statistics
     fn row_groups_pruned_statistics(&self) -> Option<usize> {
         self.pruning_metric("row_groups_pruned_statistics")
-            .map(|(pruned, _matched)| pruned)
+            .map(|(pruned, _matched, _fully)| pruned)
     }
 
     /// Metric `files_ranges_pruned_statistics` tracks both pruned and matched count,
     /// for testing purpose, here it only aggregate the `pruned` count.
     fn files_ranges_pruned_statistics(&self) -> Option<usize> {
         self.pruning_metric("files_ranges_pruned_statistics")
-            .map(|(pruned, _matched)| pruned)
+            .map(|(pruned, _matched, _fully)| pruned)
     }
 
     /// The number of row_groups matched by bloom filter or statistics
@@ -207,26 +211,27 @@ impl TestOutput {
     /// count.
     fn row_groups_matched(&self) -> Option<usize> {
         self.row_groups_bloom_filter()
-            .map(|(_pruned, matched)| matched)
+            .map(|(_pruned, matched, _fully)| matched)
     }
 
     /// The number of row_groups pruned
     fn row_groups_pruned(&self) -> Option<usize> {
         self.row_groups_bloom_filter()
-            .map(|(pruned, _matched)| pruned)
+            .map(|(pruned, _matched, _fully)| pruned)
             .zip(self.row_groups_pruned_statistics())
             .map(|(a, b)| a + b)
     }
 
     /// The number of row pages pruned
     fn row_pages_pruned(&self) -> Option<usize> {
         self.pruning_metric("page_index_rows_pruned")
-            .map(|(pruned, _matched)| pruned)
+            .map(|(pruned, _matched, _fully)| pruned)
     }
 
     /// The number of row groups pruned by limit pruning
     fn limit_pruned_row_groups(&self) -> Option<usize> {
-        self.metric_value("limit_pruned_row_groups")
+        self.pruning_metric("limit_pruned_row_groups")
+            .map(|(pruned, _, _)| pruned)
     }
 
     fn description(&self) -> String {

datafusion/core/tests/parquet/row_group_pruning.rs

@@ -157,12 +157,12 @@ impl RowGroupPruningTest {
         );
         let bloom_filter_metrics = output.row_groups_bloom_filter();
         assert_eq!(
-            bloom_filter_metrics.map(|(_pruned, matched)| matched),
+            bloom_filter_metrics.map(|(_pruned, matched, _)| matched),
             self.expected_row_group_matched_by_bloom_filter,
             "mismatched row_groups_matched_bloom_filter",
         );
         assert_eq!(
-            bloom_filter_metrics.map(|(pruned, _matched)| pruned),
+            bloom_filter_metrics.map(|(pruned, _matched, _)| pruned),
             self.expected_row_group_pruned_by_bloom_filter,
             "mismatched row_groups_pruned_bloom_filter",
         );

datafusion/datasource-parquet/src/opener.rs

@@ -68,6 +68,8 @@ pub(super) struct ParquetOpener {
     pub batch_size: usize,
     /// Optional limit on the number of rows to read
     pub limit: Option<usize>,
+    /// limit order sensitivity
+    pub limit_order_sensitive: bool,
     /// Optional predicate to apply during the scan
     pub predicate: Option<Arc<dyn PhysicalExpr>>,
     /// Schema of the output table without partition columns.
@@ -163,6 +165,8 @@ impl FileOpener for ParquetOpener {
         let encryption_context = self.get_encryption_context();
         let max_predicate_cache_size = self.max_predicate_cache_size;
 
+        let limit_order_sensitive = self.limit_order_sensitive;
+
         Ok(Box::pin(async move {
             #[cfg(feature = "parquet_encryption")]
             let file_decryption_properties = encryption_context
@@ -407,8 +411,8 @@ impl FileOpener for ParquetOpener {
                     .add_matched(n_remaining_row_groups);
             }
 
-            // Prune by limit
-            if let Some(limit) = limit {
+            // Prune by limit if limit is set and limit order is not sensitive
+            if let (Some(limit), false) = (limit, limit_order_sensitive) {
                 row_groups.prune_by_limit(limit, rg_metadata, &file_metrics);
             }
 
@@ -881,6 +885,7 @@ mod test {
                 projection: Arc::new([0, 1]),
                 batch_size: 1024,
                 limit: None,
+                limit_order_sensitive: false,
                 predicate: Some(predicate),
                 logical_file_schema: schema.clone(),
                 metadata_size_hint: None,
@@ -950,6 +955,7 @@ mod test {
                 projection: Arc::new([0]),
                 batch_size: 1024,
                 limit: None,
+                limit_order_sensitive: false,
                 predicate: Some(predicate),
                 logical_file_schema: file_schema.clone(),
                 metadata_size_hint: None,
@@ -1039,6 +1045,7 @@ mod test {
                 projection: Arc::new([0]),
                 batch_size: 1024,
                 limit: None,
+                limit_order_sensitive: false,
                 predicate: Some(predicate),
                 logical_file_schema: file_schema.clone(),
                 metadata_size_hint: None,
@@ -1131,6 +1138,7 @@ mod test {
                 projection: Arc::new([0]),
                 batch_size: 1024,
                 limit: None,
+                limit_order_sensitive: false,
                 predicate: Some(predicate),
                 logical_file_schema: file_schema.clone(),
                 metadata_size_hint: None,
@@ -1223,6 +1231,7 @@ mod test {
                 projection: Arc::new([0]),
                 batch_size: 1024,
                 limit: None,
+                limit_order_sensitive: false,
                 predicate: Some(predicate),
                 logical_file_schema: file_schema.clone(),
                 metadata_size_hint: None,
@@ -1377,6 +1386,7 @@ mod test {
             projection: Arc::new([0, 1]),
             batch_size: 1024,
             limit: None,
+            limit_order_sensitive: false,
             predicate: Some(predicate),
             logical_file_schema: Arc::clone(&table_schema),
             metadata_size_hint: None,

datafusion/datasource-parquet/src/source.rs

@@ -583,6 +583,7 @@ impl FileSource for ParquetSource {
                 .batch_size
                 .expect("Batch size must set before creating ParquetOpener"),
             limit: base_config.limit,
+            limit_order_sensitive: base_config.limit_order_sensitive,
             predicate: self.predicate.clone(),
             logical_file_schema: Arc::clone(base_config.file_schema()),
             partition_fields: base_config.table_partition_cols().clone(),

datafusion/datasource/src/file_scan_config.rs

@@ -175,6 +175,8 @@ pub struct FileScanConfig {
     /// The maximum number of records to read from this plan. If `None`,
     /// all records after filtering are returned.
     pub limit: Option<usize>,
+    /// Whether the scan's limit is order sensitive
+    pub limit_order_sensitive: bool,
     /// All equivalent lexicographical orderings that describe the schema.
     pub output_ordering: Vec<LexOrdering>,
     /// File compression type
@@ -256,6 +258,7 @@ pub struct FileScanConfigBuilder {
     object_store_url: ObjectStoreUrl,
     file_source: Arc<dyn FileSource>,
     limit: Option<usize>,
+    limit_order_sensitive: bool,
     projection_indices: Option<Vec<usize>>,
     constraints: Option<Constraints>,
     file_groups: Vec<FileGroup>,
@@ -287,6 +290,7 @@ impl FileScanConfigBuilder {
             file_compression_type: None,
             new_lines_in_values: None,
             limit: None,
+            limit_order_sensitive: false,
             projection_indices: None,
             constraints: None,
             batch_size: None,
@@ -301,6 +305,12 @@ impl FileScanConfigBuilder {
         self
     }
 
+    /// Set whether the limit should be order-sensitive.
+    pub fn with_limit_order_sensitive(mut self, order_sensitive: bool) -> Self {
+        self.limit_order_sensitive = order_sensitive;
+        self
+    }
+
     /// Set the file source for scanning files.
     ///
     /// This method allows you to change the file source implementation (e.g. ParquetSource, CsvSource, etc.)
@@ -428,6 +438,7 @@ impl FileScanConfigBuilder {
             object_store_url,
             file_source,
             limit,
+            limit_order_sensitive,
             projection_indices,
             constraints,
             file_groups,
@@ -461,6 +472,7 @@ impl FileScanConfigBuilder {
             object_store_url,
             file_source,
             limit,
+            limit_order_sensitive,
             projection_exprs,
             constraints,
             file_groups,
@@ -485,6 +497,7 @@ impl From<FileScanConfig> for FileScanConfigBuilder {
             file_compression_type: Some(config.file_compression_type),
             new_lines_in_values: Some(config.new_lines_in_values),
             limit: config.limit,
+            limit_order_sensitive: config.limit_order_sensitive,
             projection_indices: config
                 .projection_exprs
                 .map(|p| p.ordered_column_indices()),

datafusion/expr/src/logical_plan/builder.rs

@@ -2708,12 +2708,12 @@ mod tests {
 
         assert_snapshot!(plan, @r"
         Union
-          Cross Join: 
+          Cross Join:
             SubqueryAlias: left
               Values: (Int32(1))
             SubqueryAlias: right
               Values: (Int32(1))
-          Cross Join: 
+          Cross Join:
             SubqueryAlias: left
               Values: (Int32(1))
             SubqueryAlias: right