first changes

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

vortex-datafusion/src/convert/exprs.rs

@@ -6,8 +6,10 @@ use std::sync::Arc;
 use arrow_schema::{DataType, Schema};
 use datafusion_expr::Operator as DFOperator;
 use datafusion_functions::core::getfield::GetFieldFunc;
-use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr};
-use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, is_dynamic_physical_expr};
+use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef, ScalarFunctionExpr};
+use datafusion_physical_expr_common::physical_expr::{
+    is_dynamic_physical_expr, snapshot_physical_expr,
+};
 use datafusion_physical_plan::expressions as df_expr;
 use itertools::Itertools;
 use vortex::compute::LikeOptions;
@@ -24,11 +26,25 @@ use crate::convert::{FromDataFusion, TryFromDataFusion};
 
 /// Tries to convert the expressions into a vortex conjunction. Will return Ok(None) iff the input conjunction is empty.
 pub(crate) fn make_vortex_predicate(
-    predicate: &[&Arc<dyn PhysicalExpr>],
+    predicate: &[Arc<dyn PhysicalExpr>],
 ) -> VortexResult<Option<Expression>> {
     let exprs = predicate
         .iter()
-        .map(|e| Expression::try_from_df(e.as_ref()))
+        .filter_map(|e| {
+            if is_dynamic_physical_expr(e) {
+                let e = snapshot_physical_expr(e.clone()).expect("do");
+                match Expression::try_from_df(e.as_ref()) {
+                    Ok(e) => Some(Ok(e)),
+                    Err(_) => {
+                        // If we fail to convert the expression to Vortex, its safe
+                        // to drop it as we don't declare it as pushed down
+                        None
+                    }
+                }
+            } else {
+                Some(Expression::try_from_df(e.as_ref()))
+            }
+        })
         .collect::<VortexResult<Vec<_>>>()?;
 
     Ok(exprs.into_iter().reduce(and))
@@ -323,15 +339,15 @@ mod tests {
     #[test]
     fn test_make_vortex_predicate_single() {
         let col_expr = Arc::new(df_expr::Column::new("test", 0)) as Arc<dyn PhysicalExpr>;
-        let result = make_vortex_predicate(&[&col_expr]).unwrap();
+        let result = make_vortex_predicate(&[col_expr]).unwrap();
         assert!(result.is_some());
     }
 
     #[test]
     fn test_make_vortex_predicate_multiple() {
         let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc<dyn PhysicalExpr>;
         let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc<dyn PhysicalExpr>;
-        let result = make_vortex_predicate(&[&col1, &col2]).unwrap();
+        let result = make_vortex_predicate(&[col1, col2]).unwrap();
         assert!(result.is_some());
         // Result should be an AND expression combining the two columns
     }

vortex-datafusion/src/persistent/mod.rs

@@ -195,7 +195,7 @@ mod tests {
 
     #[tokio::test]
     async fn create_table_ordered_by() -> anyhow::Result<()> {
-        let dir = TempDir::new().unwrap();
+        let dir = TempDir::new()?;
 
         let factory: VortexFormatFactory = VortexFormatFactory::new();
         let mut session_state_builder = SessionStateBuilder::new().with_default_features();
@@ -238,6 +238,8 @@ mod tests {
         let (state, plan) = df.clone().into_parts();
         let physical_plan = state.create_physical_plan(&plan).await?;
 
+        dbg!(physical_plan.as_ref());
+
         insta::assert_snapshot!(DisplayableExecutionPlan::new(physical_plan.as_ref())
                 .tree_render().to_string(), @r"
         ┌───────────────────────────┐
@@ -249,8 +251,9 @@ mod tests {
         ┌─────────────┴─────────────┐
         │       DataSourceExec      │
         │    --------------------   │
-        │          files: 3         │
+        │         files: 14         │
         │       format: vortex      │
+        │      predicate: true      │
         └───────────────────────────┘
         ");
 

vortex-datafusion/src/persistent/opener.rs

@@ -40,13 +40,9 @@ pub(crate) struct VortexOpener {
     pub object_store: Arc<dyn ObjectStore>,
     /// Projection by index of the file's columns
     pub projection: Option<Arc<[usize]>>,
-    /// Filter expression optimized for pushdown into Vortex scan operations.
-    /// This may be a subset of file_pruning_predicate containing only expressions
-    /// that Vortex can efficiently evaluate.
-    pub filter: Option<PhysicalExprRef>,
     /// Filter expression used by DataFusion's FilePruner to eliminate files based on
     /// statistics and partition values without opening them.
-    pub file_pruning_predicate: Option<PhysicalExprRef>,
+    pub predicate: Option<PhysicalExprRef>,
     pub expr_adapter_factory: Option<Arc<dyn PhysicalExprAdapterFactory>>,
     pub schema_adapter_factory: Arc<dyn SchemaAdapterFactory>,
     /// Hive-style partitioning columns
@@ -149,8 +145,7 @@ impl FileOpener for VortexOpener {
         let session = self.session.clone();
         let object_store = self.object_store.clone();
         let projection = self.projection.clone();
-        let mut filter = self.filter.clone();
-        let file_pruning_predicate = self.file_pruning_predicate.clone();
+        let mut predicate = self.predicate.clone();
         let expr_adapter_factory = self.expr_adapter_factory.clone();
         let partition_fields = self.partition_fields.clone();
         let file_cache = self.file_cache.clone();
@@ -178,7 +173,8 @@ impl FileOpener for VortexOpener {
             // opening them based on:
             // - Partition column values (e.g., date=2024-01-01)
             // - File-level statistics (min/max values per column)
-            let mut file_pruner = file_pruning_predicate
+            let mut file_pruner = predicate
+                .clone()
                 .map(|predicate| {
                     // Only create pruner if we have dynamic expressions or file statistics
                     // to work with. Static predicates without stats won't benefit from pruning.
@@ -225,15 +221,16 @@ impl FileOpener for VortexOpener {
 
                 // The adapter rewrites the expression to the local file schema, allowing
                 // for schema evolution and divergence between the table's schema and individual files.
-                filter = filter
-                    .map(|filter| {
+                predicate = predicate
+                    .clone()
+                    .map(|expr| {
                         let logical_file_schema =
                             compute_logical_file_schema(&physical_file_schema, &logical_schema);
 
                         let expr = expr_adapter_factory
                             .create(logical_file_schema, physical_file_schema.clone())
                             .with_partition_values(partition_values)
-                            .rewrite(filter)?;
+                            .rewrite(expr)?;
 
                         // Expression might now reference columns that don't exist in the file, so we can give it
                         // another simplification pass.
@@ -294,11 +291,15 @@ impl FileOpener for VortexOpener {
                 );
             }
 
-            let filter = filter
+            let filter = predicate
                 .and_then(|f| {
                     let exprs = split_conjunction(&f)
                         .into_iter()
-                        .filter(|expr| can_be_pushed_down(expr, &predicate_file_schema))
+                        .cloned()
+                        .filter(|expr| {
+                            is_dynamic_physical_expr(expr)
+                                || can_be_pushed_down(expr, &predicate_file_schema)
+                        })
                         .collect::<Vec<_>>();
 
                     make_vortex_predicate(&exprs).transpose()
@@ -521,8 +522,7 @@ mod tests {
             session: SESSION.clone(),
             object_store: object_store.clone(),
             projection: Some([0].into()),
-            filter: Some(filter),
-            file_pruning_predicate: None,
+            predicate: Some(filter),
             expr_adapter_factory: expr_adapter_factory.clone(),
             schema_adapter_factory: Arc::new(DefaultSchemaAdapterFactory),
             partition_fields: vec![Arc::new(Field::new("part", DataType::Int32, false))],
@@ -602,8 +602,7 @@ mod tests {
             session: SESSION.clone(),
             object_store: object_store.clone(),
             projection: Some([0].into()),
-            filter: Some(filter),
-            file_pruning_predicate: None,
+            predicate: Some(filter),
             expr_adapter_factory: expr_adapter_factory.clone(),
             schema_adapter_factory: Arc::new(DefaultSchemaAdapterFactory),
             partition_fields: vec![],
@@ -710,11 +709,10 @@ mod tests {
             session: SESSION.clone(),
             object_store: object_store.clone(),
             projection: None,
-            filter: Some(logical2physical(
+            predicate: Some(logical2physical(
                 &col("my_struct").is_not_null(),
                 &table_schema,
             )),
-            file_pruning_predicate: None,
             expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory) as _),
             schema_adapter_factory: Arc::new(DefaultSchemaAdapterFactory),
             partition_fields: vec![],

vortex-datafusion/src/persistent/source.rs

@@ -17,9 +17,7 @@ use datafusion_physical_expr_adapter::{
     DefaultPhysicalExprAdapterFactory, PhysicalExprAdapterFactory,
 };
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
-use datafusion_physical_plan::filter_pushdown::{
-    FilterPushdownPropagation, PushedDown, PushedDownPredicate,
-};
+use datafusion_physical_plan::filter_pushdown::{FilterPushdownPropagation, PushedDown};
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::{DisplayFormatType, PhysicalExpr};
 use object_store::ObjectStore;
@@ -45,10 +43,7 @@ pub struct VortexSource {
     pub(crate) file_cache: VortexFileCache,
     /// Combined predicate expression containing all filters from DataFusion query planning.
     /// Used with FilePruner to skip files based on statistics and partition values.
-    pub(crate) full_predicate: Option<PhysicalExprRef>,
-    /// Subset of predicates that can be pushed down into Vortex scan operations.
-    /// These are expressions that Vortex can efficiently evaluate during scanning.
-    pub(crate) vortex_predicate: Option<PhysicalExprRef>,
+    pub(crate) predicate: Option<PhysicalExprRef>,
     pub(crate) batch_size: Option<usize>,
     pub(crate) projected_statistics: Option<Statistics>,
     /// This is the file schema the table expects, which is the table's schema without partition columns, and **not** the file's physical schema.
@@ -67,8 +62,7 @@ impl VortexSource {
         Self {
             session,
             file_cache,
-            full_predicate: None,
-            vortex_predicate: None,
+            predicate: None,
             batch_size: None,
             projected_statistics: None,
             arrow_file_schema: None,
@@ -140,8 +134,8 @@ impl FileSource for VortexSource {
             session: self.session.clone(),
             object_store,
             projection,
-            filter: self.vortex_predicate.clone(),
-            file_pruning_predicate: self.full_predicate.clone(),
+
+            predicate: self.predicate.clone(),
             expr_adapter_factory,
             schema_adapter_factory,
             partition_fields: base_config.table_partition_cols.clone(),
@@ -184,7 +178,7 @@ impl FileSource for VortexSource {
     }
 
     fn filter(&self) -> Option<Arc<dyn PhysicalExpr>> {
-        self.vortex_predicate.clone()
+        self.predicate.clone()
     }
 
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
@@ -197,7 +191,7 @@ impl FileSource for VortexSource {
             .clone()
             .vortex_expect("projected_statistics must be set");
 
-        if self.vortex_predicate.is_some() {
+        if self.predicate.is_some() {
             Ok(statistics.to_inexact())
         } else {
             Ok(statistics)
@@ -211,13 +205,13 @@ impl FileSource for VortexSource {
     fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
         match t {
             DisplayFormatType::Default | DisplayFormatType::Verbose => {
-                if let Some(ref predicate) = self.vortex_predicate {
+                if let Some(ref predicate) = self.predicate {
                     write!(f, ", predicate: {predicate}")?;
                 }
             }
             // Use TreeRender style key=value formatting to display the predicate
             DisplayFormatType::TreeRender => {
-                if let Some(ref predicate) = self.vortex_predicate {
+                if let Some(ref predicate) = self.predicate {
                     writeln!(f, "predicate={}", fmt_sql(predicate.as_ref()))?;
                 };
             }
@@ -238,57 +232,29 @@ impl FileSource for VortexSource {
 
         let mut source = self.clone();
 
-        // Combine new filters with existing predicate for file pruning.
-        // This full predicate is used by FilePruner to eliminate files.
-        source.full_predicate = match source.full_predicate {
-            Some(predicate) => Some(conjunction(
-                std::iter::once(predicate).chain(filters.clone()),
-            )),
-            None => Some(conjunction(filters.clone())),
-        };
-
-        let supported_filters = filters
-            .into_iter()
+        let supported = filters
+            .iter()
             .map(|expr| {
-                if can_be_pushed_down(&expr, schema) {
-                    PushedDownPredicate::supported(expr)
+                if can_be_pushed_down(expr, schema) {
+                    PushedDown::Yes
                 } else {
-                    PushedDownPredicate::unsupported(expr)
+                    PushedDown::No
                 }
             })
             .collect::<Vec<_>>();
 
-        if supported_filters
-            .iter()
-            .all(|p| matches!(p.discriminant, PushedDown::No))
-        {
-            return Ok(FilterPushdownPropagation::with_parent_pushdown_result(
-                vec![PushedDown::No; supported_filters.len()],
-            )
-            .with_updated_node(Arc::new(source) as _));
-        }
-
-        let supported = supported_filters
-            .iter()
-            .filter_map(|p| match p.discriminant {
-                PushedDown::Yes => Some(&p.predicate),
-                PushedDown::No => None,
-            })
-            .cloned();
-
-        let predicate = match source.vortex_predicate {
-            Some(predicate) => conjunction(std::iter::once(predicate).chain(supported)),
-            None => conjunction(supported),
+        // Combine new filters with existing predicate
+        source.predicate = match source.predicate {
+            Some(predicate) => Some(conjunction(
+                std::iter::once(predicate).chain(filters.clone()),
+            )),
+            None => Some(conjunction(filters.clone())),
         };
 
-        tracing::debug!(%predicate, "Saving predicate");
-
-        source.vortex_predicate = Some(predicate);
-
-        Ok(FilterPushdownPropagation::with_parent_pushdown_result(
-            supported_filters.iter().map(|f| f.discriminant).collect(),
+        Ok(
+            FilterPushdownPropagation::with_parent_pushdown_result(supported)
+                .with_updated_node(Arc::new(source) as _),
         )
-        .with_updated_node(Arc::new(source) as _))
     }
 
     fn with_schema_adapter_factory(