delta_datafusion: prune Add file actions

Author: asubiotto

crates/core/src/delta_datafusion/mod.rs

@@ -101,6 +101,7 @@ pub mod logical;
 pub mod physical;
 pub mod planner;
 
+use crate::kernel::transaction::state::AddContainer;
 pub use cdf::scan::DeltaCdfTableProvider;
 
 mod schema_adapter;
@@ -562,41 +563,55 @@ impl<'a> DeltaScanBuilder<'a> {
                 .unwrap()
         });
 
-        // Perform Pruning of files to scan
-        let (files, files_scanned, files_pruned) = match self.files {
-            Some(files) => {
-                let files = files.to_owned();
-                let files_scanned = files.len();
-                (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
+        let (files, files_scanned, files_pruned) = if let Some(predicate) = &logical_filter {
+            let pruning_predicate =
+                PruningPredicate::try_new(predicate.clone(), logical_schema.clone())?;
+
+            match self.files {
+                Some(files) => {
+                    let files_vec = files.to_vec();
+                    let container = AddContainer::new(
+                        &files_vec,
+                        &self.snapshot.metadata().partition_columns,
+                        logical_schema.clone(),
+                    );
+
+                    let keep_mask = pruning_predicate.prune(&container)?;
+                    let total_files = files.len();
+
+                    let pruned_files: Vec<_> = files
+                        .iter()
+                        .zip(keep_mask)
+                        .filter(|&(_, keep)| keep)
+                        .map(|(action, _)| action.to_owned())
+                        .collect();
+
+                    let files_kept = pruned_files.len();
+                    (pruned_files, files_kept, total_files - files_kept)
+                }
+                None => {
+                    let keep_mask = pruning_predicate.prune(self.snapshot)?;
+                    let total_files = keep_mask.len();
+
+                    let pruned_files: Vec<_> = self
                         .snapshot
                         .file_actions_iter()?
-                        .zip(files_to_prune.into_iter())
-                        .filter_map(|(action, keep)| {
-                            if keep {
-                                Some(action.to_owned())
-                            } else {
-                                files_pruned += 1;
-                                None
-                            }
-                        })
-                        .collect::<Vec<_>>();
+                        .zip(keep_mask)
+                        .filter(|&(_, keep)| keep)
+                        .map(|(action, _)| action.to_owned())
+                        .collect();
 
-                    let files_scanned = files.len();
-                    (files, files_scanned, files_pruned)
-                } else {
-                    let files = self.snapshot.file_actions()?;
-                    let files_scanned = files.len();
-                    (files, files_scanned, 0)
+                    let files_kept = pruned_files.len();
+                    (pruned_files, files_kept, total_files - files_kept)
                 }
             }
+        } else {
+            let files = match self.files {
+                Some(f) => f.to_vec(),
+                None => self.snapshot.file_actions()?,
+            };
+            let files_scanned = files.len();
+            (files, files_scanned, 0)
         };
 
         // TODO we group files together by their partition values. If the table is partitioned
@@ -2282,15 +2297,15 @@ mod tests {
 
         let df = ctx.sql("select * from test").await.unwrap();
         let actual = df.collect().await.unwrap();
-        let expected = vec! [
-                "+----+----+----+-------------------------------------------------------------------------------+",
-                "| c3 | c1 | c2 | file_source                                                                   |",
-                "+----+----+----+-------------------------------------------------------------------------------+",
-                "| 4  | 6  | a  | c1=6/c2=a/part-00011-10619b10-b691-4fd0-acc4-2a9608499d7c.c000.snappy.parquet |",
-                "| 5  | 4  | c  | c1=4/c2=c/part-00003-f525f459-34f9-46f5-82d6-d42121d883fd.c000.snappy.parquet |",
-                "| 6  | 5  | b  | c1=5/c2=b/part-00007-4e73fa3b-2c88-424a-8051-f8b54328ffdb.c000.snappy.parquet |",
-                "+----+----+----+-------------------------------------------------------------------------------+",
-            ];
+        let expected = vec![
+            "+----+----+----+-------------------------------------------------------------------------------+",
+            "| c3 | c1 | c2 | file_source                                                                   |",
+            "+----+----+----+-------------------------------------------------------------------------------+",
+            "| 4  | 6  | a  | c1=6/c2=a/part-00011-10619b10-b691-4fd0-acc4-2a9608499d7c.c000.snappy.parquet |",
+            "| 5  | 4  | c  | c1=4/c2=c/part-00003-f525f459-34f9-46f5-82d6-d42121d883fd.c000.snappy.parquet |",
+            "| 6  | 5  | b  | c1=5/c2=b/part-00007-4e73fa3b-2c88-424a-8051-f8b54328ffdb.c000.snappy.parquet |",
+            "+----+----+----+-------------------------------------------------------------------------------+",
+        ];
         assert_batches_sorted_eq!(&expected, &actual);
     }
 

crates/core/src/kernel/transaction/mod.rs

@@ -109,7 +109,7 @@ pub(crate) mod application;
 mod conflict_checker;
 mod protocol;
 #[cfg(feature = "datafusion")]
-mod state;
+pub(crate) mod state;
 
 const DELTA_LOG_FOLDER: &str = "_delta_log";
 pub(crate) const DEFAULT_RETRIES: usize = 15;

crates/core/tests/integration_datafusion.rs

@@ -24,7 +24,7 @@ use datafusion_proto::bytes::{
     logical_plan_from_bytes_with_extension_codec, logical_plan_to_bytes_with_extension_codec,
 };
 use deltalake_core::delta_datafusion::DeltaScan;
-use deltalake_core::kernel::{DataType, MapType, PrimitiveType, StructField, StructType};
+use deltalake_core::kernel::{Add, DataType, MapType, PrimitiveType, StructField, StructType};
 use deltalake_core::operations::create::CreateBuilder;
 use deltalake_core::protocol::SaveMode;
 use deltalake_core::writer::{DeltaWriter, RecordBatchWriter};
@@ -429,8 +429,6 @@ mod local {
             let _result = collect(plan.execute(0, task_ctx)?).await?;
             visit_execution_plan(&plan, &mut metrics).unwrap();
         } else {
-            // if scan produces no output from ParquetExec, we still want to visit DeltaScan
-            // to check its metrics
             visit_execution_plan(scan.as_ref(), &mut metrics).unwrap();
         }
 
@@ -859,6 +857,108 @@ mod local {
         Ok(())
     }
 
+    async fn get_scan_metrics_with_files(
+        table: &DeltaTable,
+        state: &SessionState,
+        files: Option<Vec<Add>>,
+        e: &[Expr],
+    ) -> Result<ExecutionMetricsCollector> {
+        use deltalake_core::delta_datafusion::{DeltaScanConfig, DeltaTableProvider};
+
+        let mut provider = DeltaTableProvider::try_new(
+            table.snapshot().unwrap().clone(),
+            table.log_store(),
+            DeltaScanConfig::default(),
+        )?;
+
+        if let Some(f) = files {
+            provider = provider.with_files(f);
+        }
+
+        let mut metrics = ExecutionMetricsCollector::default();
+        let scan = provider.scan(state, None, e, None).await?;
+
+        if scan.properties().output_partitioning().partition_count() > 0 {
+            let plan = CoalescePartitionsExec::new(scan);
+            let task_ctx = Arc::new(TaskContext::from(state));
+            let _result = collect(plan.execute(0, task_ctx)?).await?;
+            visit_execution_plan(&plan, &mut metrics).unwrap();
+        } else {
+            visit_execution_plan(scan.as_ref(), &mut metrics).unwrap();
+        }
+
+        Ok(metrics)
+    }
+
+    #[tokio::test]
+    async fn test_files_scanned_with_files() -> Result<()> {
+        use datafusion::prelude::*;
+        let ctx = SessionContext::new();
+        let state = ctx.state();
+
+        let batch1 = create_all_types_batch(3, 0, 0); // values 0-2
+        let batch2 = create_all_types_batch(3, 0, 4); // values 4-6
+        let batch3 = create_all_types_batch(3, 0, 7); // values 7-9
+
+        let (_tmp, mut table) = prepare_table(vec![batch1], SaveMode::Overwrite, vec![]).await;
+        let files_before_1 = table.snapshot().unwrap().file_actions().unwrap();
+
+        table = DeltaOps(table)
+            .write(vec![batch2])
+            .with_save_mode(SaveMode::Append)
+            .await
+            .unwrap();
+        let files_before_2 = table.snapshot().unwrap().file_actions().unwrap();
+
+        table = DeltaOps(table)
+            .write(vec![batch3])
+            .with_save_mode(SaveMode::Append)
+            .await
+            .unwrap();
+        let all_files = table.snapshot().unwrap().file_actions().unwrap();
+
+        assert_eq!(all_files.len(), 3);
+
+        let file_0_2 = files_before_1[0].clone();
+        let file_4_6 = files_before_2
+            .iter()
+            .find(|f| !files_before_1.iter().any(|f1| f1.path == f.path))
+            .unwrap()
+            .clone();
+        let file_7_9 = all_files
+            .iter()
+            .find(|f| !files_before_2.iter().any(|f2| f2.path == f.path))
+            .unwrap()
+            .clone();
+
+        // Test without with_files (normal snapshot pruning)
+        let e = col("int64").eq(lit(5i64));
+        let metrics = get_scan_metrics(&table, &state, &[e]).await?;
+        assert_eq!(metrics.num_scanned_files(), 1);
+        assert_eq!(metrics.skip_count, 2);
+
+        // Test with with_files providing all files (should behave the same)
+        let e = col("int64").eq(lit(5i64));
+        let metrics =
+            get_scan_metrics_with_files(&table, &state, Some(all_files.clone()), &[e]).await?;
+        assert_eq!(metrics.num_scanned_files(), 1);
+        assert_eq!(metrics.skip_count, 2);
+
+        let subset_files = vec![file_0_2.clone(), file_7_9.clone()];
+        let e = col("int64").gt(lit(6i64));
+        let metrics = get_scan_metrics_with_files(&table, &state, Some(subset_files), &[e]).await?;
+        assert_eq!(metrics.num_scanned_files(), 1);
+        assert_eq!(metrics.skip_count, 1);
+
+        let subset_files = vec![file_0_2.clone(), file_4_6.clone()];
+        let e = col("int64").gt(lit(6i64));
+        let metrics = get_scan_metrics_with_files(&table, &state, Some(subset_files), &[e]).await?;
+        assert_eq!(metrics.num_scanned_files(), 0);
+        assert_eq!(metrics.skip_count, 2);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_datafusion_partitioned_types() -> Result<()> {
         let ctx = SessionContext::new();