feat(datafusion): optmize partition pruning, pushdown full predicates for DF integration

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

Author: aditanase

crates/core/src/delta_datafusion/mod.rs

@@ -27,14 +27,17 @@ use std::fmt::{self, Debug};
 use std::sync::Arc;
 
 use arrow_array::types::UInt16Type;
-use arrow_array::{Array, DictionaryArray, RecordBatch, StringArray, TypedDictionaryArray};
+use arrow_array::{
+    Array, BooleanArray, DictionaryArray, RecordBatch, StringArray, TypedDictionaryArray,
+};
 use arrow_cast::display::array_value_to_string;
 use arrow_cast::{cast_with_options, CastOptions};
 use arrow_schema::{
     ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema, SchemaRef,
     SchemaRef as ArrowSchemaRef, TimeUnit,
 };
 use arrow_select::concat::concat_batches;
+use arrow_select::filter::filter_record_batch;
 use async_trait::async_trait;
 use chrono::{DateTime, TimeZone, Utc};
 use datafusion::catalog::memory::DataSourceExec;
@@ -59,8 +62,11 @@ use datafusion_common::{
 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::utils::{conjunction, split_conjunction};
+use datafusion_expr::{
+    col, BinaryExpr, Expr, Extension, LogicalPlan, Operator, TableProviderFilterPushDown,
+    Volatility,
+};
 use datafusion_physical_expr::{create_physical_expr, PhysicalExpr};
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
@@ -86,7 +92,9 @@ use url::Url;
 use crate::delta_datafusion::expr::parse_predicate_expression;
 use crate::delta_datafusion::schema_adapter::DeltaSchemaAdapterFactory;
 use crate::errors::{DeltaResult, DeltaTableError};
-use crate::kernel::{Add, DataCheck, EagerSnapshot, Invariant, Snapshot, StructTypeExt};
+use crate::kernel::{
+    Add, DataCheck, EagerSnapshot, Invariant, LogDataHandler, Snapshot, StructTypeExt,
+};
 use crate::logstore::LogStoreRef;
 use crate::table::builder::ensure_table_uri;
 use crate::table::state::DeltaTableState;
@@ -541,6 +549,16 @@ impl<'a> DeltaScanBuilder<'a> {
             for idx in used_columns {
                 fields.push(logical_schema.field(*idx).to_owned());
             }
+            // partition filters with Exact pushdown were removed from projection by DF optimizer,
+            // we need to add them back for the predicate pruning to work
+            if let Some(expr) = &self.filter {
+                for c in expr.column_refs() {
+                    let idx = logical_schema.index_of(c.name.as_str())?;
+                    if !used_columns.contains(&idx) {
+                        fields.push(logical_schema.field(idx).to_owned());
+                    }
+                }
+            }
             Arc::new(ArrowSchema::new(fields))
         } else {
             logical_schema
@@ -549,32 +567,48 @@ 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());
-            let simplifier = ExprSimplifier::new(simplify_context).with_max_cycles(10);
-            let simplified = simplifier.simplify(expr).unwrap();
+        let logical_filter = self
+            .filter
+            .clone()
+            .map(|expr| simplify_expr(&context, &df_schema, expr));
+        // only inexact filters should be pushed down to the data source, doing otherwise
+        // will make stats inexact and disable datafusion optimizations like AggregateStatistics
+        let pushdown_filter = self
+            .filter
+            .and_then(|expr| {
+                let predicates = split_conjunction(&expr);
+                let pushdown_filters = get_pushdown_filters(
+                    &predicates,
+                    self.snapshot.metadata().partition_columns.as_slice(),
+                );
 
-            context
-                .create_physical_expr(simplified, &df_schema)
-                .unwrap()
-        });
+                let filtered_predicates = predicates
+                    .into_iter()
+                    .zip(pushdown_filters.into_iter())
+                    .filter_map(|(filter, pushdown)| {
+                        if pushdown == TableProviderFilterPushDown::Inexact {
+                            Some(filter.clone())
+                        } else {
+                            None
+                        }
+                    });
+                conjunction(filtered_predicates)
+            })
+            .map(|expr| simplify_expr(&context, &df_schema, expr));
 
         // Perform Pruning of files to scan
-        let (files, files_scanned, files_pruned) = match self.files {
+        let (files, files_scanned, files_pruned, pruning_mask) = match self.files {
             Some(files) => {
                 let files = files.to_owned();
                 let files_scanned = files.len();
-                (files, files_scanned, 0)
+                (files, files_scanned, 0, None)
             }
             None => {
                 // 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)
+                    (files, files_scanned, 0, None)
                 } else {
                     let num_containers = self.snapshot.num_containers();
 
@@ -595,7 +629,7 @@ impl<'a> DeltaScanBuilder<'a> {
                     for (action, keep) in self
                         .snapshot
                         .file_actions_iter()?
-                        .zip(files_to_prune.into_iter())
+                        .zip(files_to_prune.iter().cloned())
                     {
                         // prune file based on predicate pushdown
                         if keep {
@@ -627,7 +661,7 @@ impl<'a> DeltaScanBuilder<'a> {
 
                     let files_scanned = files.len();
                     let files_pruned = num_containers - files_scanned;
-                    (files, files_scanned, files_pruned)
+                    (files, files_scanned, files_pruned, Some(files_to_prune))
                 }
             }
         };
@@ -684,10 +718,18 @@ impl<'a> DeltaScanBuilder<'a> {
             ));
         }
 
-        let stats = self
-            .snapshot
-            .datafusion_table_statistics()
-            .unwrap_or(Statistics::new_unknown(&schema));
+        // FIXME - where is the correct place to marry file pruning with statistics pruning?
+        //  Temporarily re-generating the log handler, just so that we can compute the stats.
+        //  Should we update datafusion_table_statistics to optionally take the mask?
+        let stats = if let Some(mask) = pruning_mask {
+            let es = self.snapshot.snapshot();
+            let pruned_stats = prune_file_statistics(&es.files, mask);
+            LogDataHandler::new(&pruned_stats, es.metadata(), es.schema()).statistics()
+        } else {
+            self.snapshot.datafusion_table_statistics()
+        };
+
+        let stats = stats.unwrap_or(Statistics::new_unknown(&schema));
 
         let parquet_options = TableParquetOptions {
             global: self.session.config().options().execution.parquet.clone(),
@@ -700,7 +742,7 @@ impl<'a> DeltaScanBuilder<'a> {
         // Sometimes (i.e Merge) we want to prune files that don't make the
         // filter and read the entire contents for files that do match the
         // filter
-        if let Some(predicate) = logical_filter {
+        if let Some(predicate) = pushdown_filter {
             if config.enable_parquet_pushdown {
                 file_source = file_source.with_predicate(Arc::clone(&file_schema), predicate);
             }
@@ -746,6 +788,43 @@ impl<'a> DeltaScanBuilder<'a> {
     }
 }
 
+fn simplify_expr(
+    context: &SessionContext,
+    df_schema: &DFSchema,
+    expr: Expr,
+) -> Arc<dyn PhysicalExpr> {
+    // Simplify the expression first
+    let props = ExecutionProps::new();
+    let simplify_context = SimplifyContext::new(&props).with_schema(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)
+        .unwrap()
+}
+
+fn prune_file_statistics(
+    record_batches: &Vec<RecordBatch>,
+    pruning_mask: Vec<bool>,
+) -> Vec<RecordBatch> {
+    let mut filtered_batches = Vec::new();
+    let mut mask_offset = 0;
+
+    for batch in record_batches {
+        let num_rows = batch.num_rows();
+        let batch_mask = &pruning_mask[mask_offset..mask_offset + num_rows];
+        mask_offset += num_rows;
+
+        let boolean_mask = BooleanArray::from(batch_mask.to_vec());
+        let filtered_batch =
+            filter_record_batch(batch, &boolean_mask).expect("Failed to filter RecordBatch");
+        filtered_batches.push(filtered_batch);
+    }
+
+    filtered_batches
+}
+
 // TODO: implement this for Snapshot, not for DeltaTable
 #[async_trait]
 impl TableProvider for DeltaTable {
@@ -793,17 +872,81 @@ 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.as_slice();
+        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: &[String],
+) -> Vec<TableProviderFilterPushDown> {
+    filter
+        .iter()
+        .cloned()
+        .map(|expr| {
+            let applicable = expr_is_exact_predicate_for_cols(partition_cols, expr);
+            if !expr.column_refs().is_empty() && applicable {
+                TableProviderFilterPushDown::Exact
+            } else {
+                TableProviderFilterPushDown::Inexact
+            }
+        })
+        .collect()
+}
+
+// inspired from datafusion::listing::helpers, but adapted to only stats based pruning
+fn expr_is_exact_predicate_for_cols(partition_cols: &[String], expr: &Expr) -> bool {
+    let mut is_applicable = true;
+    expr.apply(|expr| match expr {
+        Expr::Column(Column { ref name, .. }) => {
+            is_applicable &= partition_cols.contains(&name);
+
+            // TODO: decide if we should constrain this to Utf8 columns (including views, dicts etc)
+
+            if is_applicable {
+                Ok(TreeNodeRecursion::Jump)
+            } else {
+                Ok(TreeNodeRecursion::Stop)
+            }
+        }
+        Expr::BinaryExpr(BinaryExpr { ref op, .. }) => {
+            is_applicable &= matches!(
+                op,
+                Operator::And
+                    | Operator::Or
+                    | Operator::NotEq
+                    | Operator::Eq
+                    | Operator::Gt
+                    | Operator::GtEq
+                    | Operator::Lt
+                    | Operator::LtEq
+            );
+            if is_applicable {
+                Ok(TreeNodeRecursion::Continue)
+            } else {
+                Ok(TreeNodeRecursion::Stop)
+            }
+        }
+        Expr::Literal(_)
+        | Expr::Not(_)
+        | Expr::IsNotNull(_)
+        | Expr::IsNull(_)
+        | Expr::Between(_)
+        | Expr::InList(_) => Ok(TreeNodeRecursion::Continue),
+        _ => {
+            is_applicable = false;
+            Ok(TreeNodeRecursion::Stop)
+        }
+    })
+    .unwrap();
+    is_applicable
+}
+
 /// A Delta table provider that enables additional metadata columns to be included during the scan
 #[derive(Debug)]
 pub struct DeltaTableProvider {
@@ -885,10 +1028,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.as_slice();
+        Ok(get_pushdown_filters(filter, partition_cols))
     }
 
     fn statistics(&self) -> Option<Statistics> {

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

@@ -891,14 +891,53 @@ 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 optional but will optimize partition column pruning
+        fn contained(&self, column: &Column, value: &HashSet<ScalarValue>) -> Option<BooleanArray> {
+            if value.is_empty() || !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());
+
+            // TODO: this was inspired by parquet's BloomFilter pruning, decide if we should
+            //  just convert to Vec<String> for a subset of column types and use .contains
+            fn check_scalar(pv: &str, value: &ScalarValue) -> bool {
+                match value {
+                    ScalarValue::Utf8(Some(v))
+                    | ScalarValue::Utf8View(Some(v))
+                    | ScalarValue::LargeUtf8(Some(v)) => pv == v,
+
+                    ScalarValue::Dictionary(_, inner) => check_scalar(pv, inner),
+                    // FIXME: is this a good enough default or should we sync this with
+                    //  expr_applicable_for_cols and bail out with None
+                    _ => value.to_string() == pv,
+                }
+            }
+
+            for i in 0..partition_values.len() {
+                if partition_values.is_null(i) {
+                    contains.push(false);
+                } else {
+                    contains.push(
+                        value
+                            .iter()
+                            .any(|scalar| check_scalar(partition_values.value(i), scalar)),
+                    );
+                }
+            }
+
+            Some(BooleanArray::from(contains))
         }
     }
 }