Delegate to datafusion's FilePruner to know if it has any work to do

vortex-datafusion/src/persistent/opener.rs

@@ -11,7 +11,6 @@ use arrow_schema::Schema;
 use datafusion_common::DataFusionError;
 use datafusion_common::Result as DFResult;
 use datafusion_common::ScalarValue;
-use datafusion_common::Statistics;
 use datafusion_common::arrow::array::AsArray;
 use datafusion_common::arrow::array::RecordBatch;
 use datafusion_common::exec_datafusion_err;
@@ -25,7 +24,6 @@ use datafusion_physical_expr::PhysicalExprRef;
 use datafusion_physical_expr::projection::ProjectionExprs;
 use datafusion_physical_expr::simplifier::PhysicalExprSimplifier;
 use datafusion_physical_expr::split_conjunction;
-use datafusion_physical_expr::utils::collect_columns;
 use datafusion_physical_expr::utils::reassign_expr_columns;
 use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
 use datafusion_physical_expr_adapter::replace_columns_with_literals;
@@ -159,13 +157,6 @@ impl FileOpener for VortexOpener {
             .zip(file.partition_values.clone())
             .collect::<std::collections::HashMap<String, ScalarValue>>();
 
-        let predicate_uses_partition_columns =
-            file_pruning_predicate.as_ref().is_some_and(|predicate| {
-                collect_columns(predicate)
-                    .iter()
-                    .any(|column| literal_value_cols.contains_key(column.name()))
-            });
-
         if !literal_value_cols.is_empty() {
             projection = projection.try_map_exprs(|expr| {
                 replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols)
@@ -179,26 +170,15 @@ impl FileOpener for VortexOpener {
         }
 
         Ok(async move {
-            // FilePruner requires a statistics object even when the rewritten predicate
-            // only contains partition literals. Supply unknown file-column statistics in
-            // that case so static and dynamic partition predicates can still prune.
-            let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns)
-                .then(|| {
-                    file.clone()
-                        .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema)))
-                });
-            let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file);
-
-            let mut file_pruner = file_pruning_predicate
-                .filter(|_| file.has_statistics() || predicate_uses_partition_columns)
-                .and_then(|predicate| {
-                    FilePruner::try_new(
-                        Arc::clone(&predicate),
-                        &unified_file_schema,
-                        pruning_file,
-                        predicate_creation_errors,
-                    )
-                });
+            // FilePruner now has all the logic to figure out if its worthwhile, and will return `None` if not.
+            let mut file_pruner = file_pruning_predicate.and_then(|predicate| {
+                FilePruner::try_new(
+                    Arc::clone(&predicate),
+                    &unified_file_schema,
+                    &file,
+                    predicate_creation_errors,
+                )
+            });
 
             // Check if this file should be pruned based on statistics/partition values.
             // Returns empty stream if file can be skipped entirely.
@@ -663,12 +643,12 @@ mod tests {
     use datafusion::physical_expr::planner::logical2physical;
     use datafusion::physical_expr_adapter::DefaultPhysicalExprAdapterFactory;
     use datafusion::scalar::ScalarValue;
+    use datafusion_common::Statistics;
     use datafusion_common::stats::Precision;
     use datafusion_execution::cache::default_cache::DefaultCache;
     use datafusion_expr::Operator;
     use datafusion_physical_expr::PhysicalExpr;
     use datafusion_physical_expr::expressions as df_expr;
-    use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr;
     use datafusion_physical_expr::projection::ProjectionExpr;
     use insta::assert_snapshot;
     use itertools::Itertools;
@@ -1009,48 +989,6 @@ mod tests {
         Ok(())
     }
 
-    #[tokio::test]
-    async fn test_file_pruning_replaces_partition_columns_without_file_statistics()
-    -> anyhow::Result<()> {
-        let object_store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
-        let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
-        let table_schema = TableSchema::builder(Arc::clone(&file_schema))
-            .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))])
-            .build();
-
-        let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef;
-        let predicate = Arc::new(df_expr::BinaryExpr::new(
-            Arc::clone(&partition_column),
-            Operator::Gt,
-            df_expr::lit(ScalarValue::Int32(Some(1))),
-        )) as PhysicalExprRef;
-        let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new(
-            vec![partition_column],
-            predicate,
-        )) as PhysicalExprRef;
-
-        let mut opener = make_opener(object_store, table_schema, None);
-        opener.file_pruning_predicate = Some(dynamic_predicate);
-        let df_metrics = opener.df_metrics.clone();
-
-        // The file does not exist and has no statistics. Replacing `part` with 1
-        // makes the predicate false, so pruning must happen before any file I/O.
-        let mut file = PartitionedFile::new("missing.vortex", 1);
-        file.partition_values = vec![ScalarValue::Int32(Some(1))];
-        let batches = opener.open(file)?.await?.try_collect::<Vec<_>>().await?;
-
-        assert!(batches.is_empty());
-        assert_eq!(
-            df_metrics
-                .clone_inner()
-                .sum_by_name("num_predicate_creation_errors")
-                .map(|metric| metric.as_usize()),
-            Some(0)
-        );
-
-        Ok(())
-    }
-
     #[tokio::test]
     async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> {
         let object_store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;