fix(datafusion): add partition pruning, pushdown full predicates for DF integration

Author: aditanase

crates/core/src/delta_datafusion/mod.rs

@@ -58,7 +58,9 @@ use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::logical_plan::CreateExternalTable;
 use datafusion_expr::simplify::SimplifyContext;
 use datafusion_expr::utils::conjunction;
-use datafusion_expr::{col, Expr, Extension, LogicalPlan, TableProviderFilterPushDown, Volatility};
+use datafusion_expr::{
+    col, BinaryExpr, Expr, Extension, LogicalPlan, TableProviderFilterPushDown, Volatility,
+};
 use datafusion_physical_expr::{create_physical_expr, PhysicalExpr};
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
@@ -534,6 +536,10 @@ impl<'a> DeltaScanBuilder<'a> {
             Some(schema.clone()),
         )?;
 
+        // TODO temporarily using full schema to generate pruning predicates
+        //  should we optimize this by only including fields referenced from predicates?
+        let filter_df_schema = logical_schema.clone().to_dfschema()?;
+
         let logical_schema = if let Some(used_columns) = self.projection {
             let mut fields = vec![];
             for idx in used_columns {
@@ -545,18 +551,17 @@ impl<'a> DeltaScanBuilder<'a> {
         };
 
         let context = SessionContext::new();
-        let df_schema = logical_schema.clone().to_dfschema()?;
 
         let logical_filter = self.filter.map(|expr| {
             // Simplify the expression first
             let props = ExecutionProps::new();
             let simplify_context =
-                SimplifyContext::new(&props).with_schema(df_schema.clone().into());
+                SimplifyContext::new(&props).with_schema(filter_df_schema.clone().into());
             let simplifier = ExprSimplifier::new(simplify_context).with_max_cycles(10);
             let simplified = simplifier.simplify(expr).unwrap();
 
             context
-                .create_physical_expr(simplified, &df_schema)
+                .create_physical_expr(simplified, &filter_df_schema)
                 .unwrap()
         });
 
@@ -757,17 +762,47 @@ impl TableProvider for DeltaTable {
         &self,
         filter: &[&Expr],
     ) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
-        Ok(filter
-            .iter()
-            .map(|_| TableProviderFilterPushDown::Inexact)
-            .collect())
+        let partition_cols = self.snapshot()?.metadata().partition_columns.clone();
+        Ok(get_pushdown_filters(filter, partition_cols))
     }
 
     fn statistics(&self) -> Option<Statistics> {
         self.snapshot().ok()?.datafusion_table_statistics()
     }
 }
 
+fn get_pushdown_filters(
+    filter: &[&Expr],
+    partition_cols: Vec<String>,
+) -> Vec<TableProviderFilterPushDown> {
+    filter
+        .iter()
+        .map(|filter| {
+            let columns = extract_columns(filter);
+            if !columns.is_empty() && columns.iter().all(|col| partition_cols.contains(col)) {
+                TableProviderFilterPushDown::Exact
+            } else {
+                TableProviderFilterPushDown::Inexact
+            }
+        })
+        .collect()
+}
+
+fn extract_columns(expr: &Expr) -> Vec<String> {
+    let mut columns = Vec::new();
+    match expr {
+        Expr::Column(col) => columns.push(col.name.clone()),
+        Expr::BinaryExpr(BinaryExpr { left, right, .. }) => {
+            let left_columns = extract_columns(left);
+            let right_columns = extract_columns(right);
+            columns.extend(left_columns);
+            columns.extend(right_columns);
+        }
+        _ => {}
+    }
+    columns
+}
+
 /// A Delta table provider that enables additional metadata columns to be included during the scan
 #[derive(Debug)]
 pub struct DeltaTableProvider {
@@ -849,10 +884,8 @@ impl TableProvider for DeltaTableProvider {
         &self,
         filter: &[&Expr],
     ) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
-        Ok(filter
-            .iter()
-            .map(|_| TableProviderFilterPushDown::Inexact)
-            .collect())
+        let partition_cols = self.snapshot.metadata().partition_columns.clone();
+        Ok(get_pushdown_filters(filter, partition_cols))
     }
 
     fn statistics(&self) -> Option<Statistics> {

crates/core/src/kernel/snapshot/log_data.rs

@@ -891,14 +891,37 @@ mod datafusion {
             arrow_cast::cast(batch.column_by_name("output")?, &ArrowDataType::UInt64).ok()
         }
 
-        // This function is required since DataFusion 35.0, but is implemented as a no-op
-        // https://github.com/apache/arrow-datafusion/blob/ec6abece2dcfa68007b87c69eefa6b0d7333f628/datafusion/core/src/datasource/physical_plan/parquet/page_filter.rs#L550
-        fn contained(
-            &self,
-            _column: &Column,
-            _value: &HashSet<ScalarValue>,
-        ) -> Option<BooleanArray> {
-            None
+        // This function is required for partition column pruning to be executed correctly
+        fn contained(&self, column: &Column, value: &HashSet<ScalarValue>) -> Option<BooleanArray> {
+            // Check if the column is a partition column
+            if !self.metadata.partition_columns.contains(&column.name) {
+                return None;
+            }
+
+            // Retrieve the partition values for the column
+            let partition_values = self.pick_stats(column, "__dummy__")?;
+
+            let partition_values = partition_values
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .ok_or(DeltaTableError::generic(
+                    "failed to downcast string result to StringArray.",
+                ))
+                .ok()?;
+
+            let mut contains = Vec::with_capacity(partition_values.len());
+
+            for i in 0..partition_values.len() {
+                if partition_values.is_null(i) {
+                    contains.push(false);
+                } else {
+                    let partition_value =
+                        ScalarValue::Utf8(Some(partition_values.value(i).to_string()));
+                    contains.push(value.contains(&partition_value));
+                }
+            }
+
+            Some(BooleanArray::from(contains))
         }
     }
 }