feat(datafusion): file pruning based on pushdown limit for partition cols filters

Signed-off-by: Adrian Tanase <[email protected]>

Author: aditanase

crates/core/src/delta_datafusion/mod.rs

@@ -47,7 +47,7 @@ use datafusion::execution::context::{SessionConfig, SessionContext, SessionState
 use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::execution::FunctionRegistry;
 use datafusion::optimizer::simplify_expressions::ExprSimplifier;
-use datafusion::physical_optimizer::pruning::PruningPredicate;
+use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningStatistics};
 use datafusion_common::scalar::ScalarValue;
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion_common::{
@@ -581,31 +581,64 @@ impl<'a> DeltaScanBuilder<'a> {
                 (files, files_scanned, 0)
             }
             None => {
-                if let Some(predicate) = &logical_filter {
-                    let pruning_predicate =
-                        PruningPredicate::try_new(predicate.clone(), logical_schema.clone())?;
-                    let files_to_prune = pruning_predicate.prune(self.snapshot)?;
-                    let mut files_pruned = 0usize;
-                    let files = self
+                // early return in case we have no push down filters or limit
+                if logical_filter.is_none() && self.limit.is_none() {
+                    let files = self.snapshot.file_actions()?;
+                    let files_scanned = files.len();
+                    (files, files_scanned, 0)
+                } else {
+                    let num_containers = self.snapshot.num_containers();
+
+                    let files_to_prune = if let Some(predicate) = &logical_filter {
+                        let pruning_predicate =
+                            PruningPredicate::try_new(predicate.clone(), logical_schema.clone())?;
+                        pruning_predicate.prune(self.snapshot)?
+                    } else {
+                        vec![true; num_containers]
+                    };
+
+                    // needed to enforce limit and deal with missing statistics
+                    // rust port of https://github.com/delta-io/delta/pull/1495
+                    let mut pruned_without_stats = vec![];
+                    let mut rows_collected = 0;
+                    let mut files = vec![];
+
+                    for (action, keep) in self
                         .snapshot
                         .file_actions_iter()?
                         .zip(files_to_prune.into_iter())
-                        .filter_map(|(action, keep)| {
-                            if keep {
-                                Some(action.to_owned())
+                    {
+                        // prune file based on predicate pushdown
+                        if keep {
+                            // prune file based on limit pushdown
+                            if let Some(limit) = self.limit {
+                                if let Some(stats) = action.get_stats()? {
+                                    if rows_collected <= limit as i64 {
+                                        rows_collected += stats.num_records;
+                                        files.push(action.to_owned());
+                                    } else {
+                                        break;
+                                    }
+                                } else {
+                                    // some files are missing stats; skipping but storing them
+                                    // in a list in case we can't reach the target limit
+                                    pruned_without_stats.push(action.to_owned());
+                                }
                             } else {
-                                files_pruned += 1;
-                                None
+                                files.push(action.to_owned());
                             }
-                        })
-                        .collect::<Vec<_>>();
+                        }
+                    }
+
+                    if let Some(limit) = self.limit {
+                        if rows_collected < limit as i64 {
+                            files.extend(pruned_without_stats);
+                        }
+                    }
 
                     let files_scanned = files.len();
+                    let files_pruned = num_containers - files_scanned;
                     (files, files_scanned, files_pruned)
-                } else {
-                    let files = self.snapshot.file_actions()?;
-                    let files_scanned = files.len();
-                    (files, files_scanned, 0)
                 }
             }
         };

crates/core/tests/integration_datafusion.rs

@@ -285,6 +285,34 @@ mod local {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_files_scanned_pushdown_limit() -> Result<()> {
+        use datafusion::prelude::*;
+        let ctx = SessionContext::new();
+        let state = ctx.state();
+        let table = open_table("../test/tests/data/delta-0.8.0")
+            .await?;
+
+        // Simple Equality test, we only exercise the limit in this test
+        let e = col("value").eq(lit(2));
+        let metrics = get_scan_metrics(&table, &state, &[e.clone()]).await?;
+        assert_eq!(metrics.num_scanned_files(), 2);
+        assert_eq!(metrics.num_scanned_files(), metrics.keep_count);
+        assert_eq!(metrics.skip_count, 0);
+
+        let metrics = get_scan_metrics_with_limit(&table, &state, &[e.clone()], Some(1)).await?;
+        assert_eq!(metrics.num_scanned_files(), 1);
+        assert_eq!(metrics.num_scanned_files(), metrics.keep_count);
+        assert_eq!(metrics.skip_count, 1);
+
+        let metrics = get_scan_metrics_with_limit(&table, &state, &[e.clone()], Some(3)).await?;
+        assert_eq!(metrics.num_scanned_files(), 2);
+        assert_eq!(metrics.num_scanned_files(), metrics.keep_count);
+        assert_eq!(metrics.skip_count, 0);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_datafusion_write_from_serialized_delta_scan() -> Result<()> {
         // Build an execution plan for scanning a DeltaTable and serialize it to bytes.
@@ -513,13 +541,14 @@ mod local {
         Ok(())
     }
 
-    async fn get_scan_metrics(
+    async fn get_scan_metrics_with_limit(
         table: &DeltaTable,
         state: &SessionState,
         e: &[Expr],
+        limit: Option<usize>,
     ) -> Result<ExecutionMetricsCollector> {
         let mut metrics = ExecutionMetricsCollector::default();
-        let scan = table.scan(state, None, e, None).await?;
+        let scan = table.scan(state, None, e, limit).await?;
         if scan.properties().output_partitioning().partition_count() > 0 {
             let plan = CoalescePartitionsExec::new(scan);
             let task_ctx = Arc::new(TaskContext::from(state));
@@ -534,6 +563,14 @@ mod local {
         Ok(metrics)
     }
 
+    async fn get_scan_metrics(
+        table: &DeltaTable,
+        state: &SessionState,
+        e: &[Expr],
+    ) -> Result<ExecutionMetricsCollector> {
+        get_scan_metrics_with_limit(table, state, e, None).await
+    }
+
     fn create_all_types_batch(
         not_null_rows: usize,
         null_rows: usize,