[HSTACK] [DF] - pushdown limit and predicates for partition cols

Author: aditanase

crates/core/src/delta_datafusion/mod.rs

@@ -41,15 +41,14 @@ use datafusion::catalog::{Session, TableProviderFactory};
 use datafusion::config::TableParquetOptions;
 use datafusion::datasource::memory::DataSourceExec;
 use datafusion::datasource::physical_plan::{
-    wrap_partition_type_in_dict, wrap_partition_value_in_dict, FileScanConfig,
-    ParquetSource,
+    wrap_partition_type_in_dict, wrap_partition_value_in_dict, FileScanConfig, ParquetSource,
 };
 use datafusion::datasource::{listing::PartitionedFile, MemTable, TableProvider, TableType};
 use datafusion::execution::context::{SessionConfig, SessionContext, SessionState, TaskContext};
 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::{
@@ -60,7 +59,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};
@@ -556,6 +557,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 {
@@ -567,18 +572,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()
         });
 
@@ -590,31 +594,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)
                 }
             }
         };
@@ -810,17 +847,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 {
@@ -926,10 +993,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> {