Merge pull request JanKaul#321 from splitgraph/fix-eq-del-predicates

fix: ensure we only push down filters to delete files that are compatible with their schema

Author: JanKaul

datafusion_iceberg/src/table.rs

@@ -424,15 +424,15 @@ async fn table_scan(
         .unwrap();
 
     // If there is a filter expression the manifests to read are pruned based on the pruning statistics available in the manifest_list file.
-    let physical_predicate = if let Some(predicate) = conjunction(filters.iter().cloned()) {
-        Some(create_physical_expr(
-            &predicate,
-            &arrow_schema.as_ref().clone().try_into()?,
-            session.execution_props(),
-        )?)
-    } else {
-        None
-    };
+    let physical_predicate = conjunction(filters.iter().cloned())
+        .map(|predicate| {
+            create_physical_expr(
+                &predicate,
+                &arrow_schema.as_ref().clone().try_into()?,
+                session.execution_props(),
+            )
+        })
+        .transpose()?;
 
     let mut table_partition_cols = datafusion_partition_columns(partition_fields)?;
 
@@ -628,32 +628,24 @@ async fn table_scan(
     }
 
     let file_source = {
-        let physical_predicate = physical_predicate.clone();
         let table_schema = TableSchema::new(
             file_schema.clone(),
             table_partition_cols.iter().cloned().map(Arc::new).collect(),
         );
-        async move {
-            Arc::new(
-                if let Some(physical_predicate) = physical_predicate.clone() {
-                    ParquetSource::new(table_schema)
-                        .with_predicate(physical_predicate)
-                        .with_pushdown_filters(true)
-                } else {
-                    ParquetSource::new(table_schema)
-                },
-            )
+        let mut source = ParquetSource::new(table_schema);
+        if let Some(ref predicate) = physical_predicate {
+            source = source
+                .with_predicate(predicate.clone())
+                .with_pushdown_filters(true);
         }
-        .instrument(tracing::debug_span!("datafusion_iceberg::file_source"))
-        .await
+        Arc::new(source)
     };
 
     // Create plan for every partition with delete files
     let mut plans = stream::iter(equality_delete_file_groups.into_iter())
         .then(|(partition_value, mut delete_files)| {
             let object_store_url = object_store_url.clone();
             let statistics = statistics.clone();
-            let physical_predicate = physical_predicate.clone();
             let schema = &schema;
             let file_source = file_source.clone();
             let projection_expr = projection_expr.clone();
@@ -707,7 +699,6 @@ async fn table_scan(
                     .try_fold(None, |acc, delete_manifest| {
                         let object_store_url = object_store_url.clone();
                         let statistics = statistics.clone();
-                        let physical_predicate = physical_predicate.clone();
                         let schema = &schema;
                         let file_source = file_source.clone();
                         let mut data_files = Vec::new();
@@ -733,6 +724,7 @@ async fn table_scan(
                             .unwrap();
                             data_files.push(data_file);
                         }
+
                         async move {
                             let delete_schema = schema.project(
                                 delete_manifest
@@ -745,6 +737,27 @@ async fn table_scan(
                             let delete_file_schema: SchemaRef =
                                 Arc::new((delete_schema.fields()).try_into().unwrap());
 
+                            // Scope down the filter expressions used in the query to only those
+                            // which are completely contained in the delete file schema.
+                            let delete_physical_predicate = conjunction(
+                                filters
+                                    .iter()
+                                    .filter(|expr| {
+                                        expr.column_refs().iter().all(|col| {
+                                            delete_file_schema.field_with_name(col.name()).is_ok()
+                                        })
+                                    })
+                                    .cloned(),
+                            )
+                            .map(|predicate| {
+                                create_physical_expr(
+                                    &predicate,
+                                    &delete_file_schema.as_ref().clone().try_into()?,
+                                    session.execution_props(),
+                                )
+                            })
+                            .transpose()?;
+
                             let last_updated_ms = table.metadata().last_updated_ms;
                             let manifest_path = if enable_manifest_file_path_column {
                                 Some(delete_manifest.0.clone())
@@ -759,15 +772,13 @@ async fn table_scan(
                                 manifest_path,
                             )?;
 
-                            let delete_file_source = Arc::new(
-                                if let Some(physical_predicate) = physical_predicate.clone() {
-                                    ParquetSource::new(delete_file_schema)
-                                        .with_predicate(physical_predicate)
-                                        .with_pushdown_filters(true)
-                                } else {
-                                    ParquetSource::new(delete_file_schema)
-                                },
-                            );
+                            let mut delete_source = ParquetSource::new(delete_file_schema);
+                            if let Some(pred) = delete_physical_predicate {
+                                delete_source = delete_source
+                                    .with_predicate(pred)
+                                    .with_pushdown_filters(true);
+                            }
+                            let delete_file_source = Arc::new(delete_source);
 
                             let delete_file_scan_config = FileScanConfigBuilder::new(
                                 object_store_url.clone(),

datafusion_iceberg/tests/equality_delete.rs

@@ -1,3 +1,4 @@
+use datafusion::common::test_util::batches_to_string;
 use datafusion::{
     arrow::{error::ArrowError, record_batch::RecordBatch},
     assert_batches_eq,
@@ -25,6 +26,15 @@ use object_store::local::LocalFileSystem;
 use std::sync::Arc;
 use tempfile::TempDir;
 
+async fn run_query(query: &str, ctx: &SessionContext) -> Vec<RecordBatch> {
+    ctx.sql(query)
+        .await
+        .expect("Failed to create plan for query")
+        .collect()
+        .await
+        .expect("Failed to execute query")
+}
+
 /// Convert DuckDB's Arrow RecordBatch to DataFusion's Arrow RecordBatch
 /// using Arrow IPC format as an interchange format.
 /// This allows compatibility between different Arrow versions.
@@ -120,28 +130,19 @@ pub async fn test_equality_delete() {
 
     ctx.register_catalog("warehouse", datafusion_catalog);
 
-    ctx.sql(
-        "INSERT INTO warehouse.test.orders (id, customer_id, product_id, date, amount) VALUES 
+    run_query(
+        "INSERT INTO warehouse.test.orders (id, customer_id, product_id, date, amount) VALUES
                 (1, 1, 1, '2020-01-01', 1),
                 (2, 2, 1, '2020-01-01', 1),
                 (3, 3, 1, '2020-01-01', 3),
                 (4, 1, 2, '2020-02-02', 1),
                 (5, 1, 1, '2020-02-02', 2),
                 (6, 3, 3, '2020-02-02', 3);",
+        &ctx,
     )
-    .await
-    .expect("Failed to create query plan for insert")
-    .collect()
-    .await
-    .expect("Failed to insert values into table");
+    .await;
 
-    let batches = ctx
-        .sql("select * from warehouse.test.orders order by id")
-        .await
-        .expect("Failed to create plan for select")
-        .collect()
-        .await
-        .expect("Failed to execute select query");
+    let batches = run_query("select * from warehouse.test.orders order by id", &ctx).await;
 
     let expected = [
         "+----+-------------+------------+------------+--------+",
@@ -157,15 +158,7 @@ pub async fn test_equality_delete() {
     ];
     assert_batches_eq!(expected, &batches);
 
-    let batches = ctx
-        .sql(
-            "SELECT id, customer_id, product_id, date FROM warehouse.test.orders WHERE customer_id = 1 order by id",
-        )
-        .await
-        .expect("Failed to create query plan for insert")
-        .collect()
-        .await
-        .expect("Failed to insert values into table");
+    let batches = run_query("SELECT id, customer_id, product_id, date FROM warehouse.test.orders WHERE customer_id = 1 order by id", &ctx).await;
 
     let expected = [
         "+----+-------------+------------+------------+",
@@ -204,13 +197,7 @@ pub async fn test_equality_delete() {
         .await
         .unwrap();
 
-    let batches = ctx
-        .sql("select * from warehouse.test.orders order by id")
-        .await
-        .expect("Failed to create plan for select")
-        .collect()
-        .await
-        .expect("Failed to execute select query");
+    let batches = run_query("select * from warehouse.test.orders order by id", &ctx).await;
 
     let expected = [
         "+----+-------------+------------+------------+--------+",
@@ -237,25 +224,16 @@ pub async fn test_equality_delete() {
     assert_batches_eq!(expected, &duckdb_batches);
 
     // Test that projecting a column that is not included in equality deletes works
-    ctx.sql(
+    run_query(
         "INSERT INTO warehouse.test.orders (id, customer_id, product_id, date, amount) VALUES
                 (7, 3, 2, '2020-01-01', 2),
                 (8, 2, 1, '2020-02-02', 3),
                 (9, 1, 3, '2020-01-01', 1);",
+        &ctx,
     )
-    .await
-    .expect("Failed to create query plan for insert")
-    .collect()
-    .await
-    .expect("Failed to insert values into table");
+    .await;
 
-    let batches = ctx
-        .sql("select sum(amount) from warehouse.test.orders")
-        .await
-        .expect("Failed to create plan for select")
-        .collect()
-        .await
-        .expect("Failed to execute select query");
+    let batches = run_query("select sum(amount) from warehouse.test.orders", &ctx).await;
 
     let expected = [
         "+-----------------------------------+",
@@ -265,4 +243,20 @@ pub async fn test_equality_delete() {
         "+-----------------------------------+",
     ];
     assert_batches_eq!(expected, &batches);
+
+    // Test that using a filter on a column that is not included in equality deletes works
+    let query = "select count(*) from warehouse.test.orders where product_id = 1 and (amount = 1 or customer_id = 3)";
+    let batches = run_query(query, &ctx).await;
+    let expected = [
+        "+----------+",
+        "| count(*) |",
+        "+----------+",
+        "| 2        |",
+        "+----------+",
+    ];
+    assert_batches_eq!(expected, &batches);
+
+    // Ensure we only pushed down predicates that have matching columns with delete file schemas (i.e. amount was not pushed down).
+    let batches = run_query(&format!("explain {query}"), &ctx).await;
+    assert!(batches_to_string(&batches).contains("projection=[id, customer_id, product_id, date], file_type=parquet, predicate=product_id@2 = 1,"));
 }