fix: Update persistent module to use simplified expression handling from PR #8

lukekim · lukekim · commit 995c5e39fd6a · 2026-02-26T18:17:49.000-08:00
diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs
@@ -6,7 +6,7 @@
 use std::fmt::Debug;
 
 use datafusion_common::stats::Precision as DFPrecision;
-use vortex::expr::stats::Precision;
+use vortex::stats::Precision;
 
 mod convert;
 mod persistent;
diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs
@@ -2,8 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::ops::Range;
-use std::sync::Arc;
-use std::sync::Weak;
+use std::sync::{Arc, Weak};
 
 use arrow_schema::Schema;
 use datafusion_common::DataFusionError;
@@ -43,8 +42,7 @@ use vortex::metrics::Label;
 use vortex::metrics::MetricsRegistry;
 use vortex::scan::ScanBuilder;
 use vortex::session::VortexSession;
-use vortex_utils::aliases::dash_map::DashMap;
-use vortex_utils::aliases::dash_map::Entry;
+use vortex_utils::aliases::dash_map::{DashMap, Entry};
 
 use crate::VortexAccessPlan;
 use crate::convert::exprs::ExpressionConvertor;
@@ -90,6 +88,7 @@ pub(crate) struct VortexOpener {
     pub layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
     /// Whether the query has output ordering specified
     pub has_output_ordering: bool,
+}
 
     pub expression_convertor: Arc<dyn ExpressionConvertor>,
     pub file_metadata_cache: Option<Arc<dyn FileMetadataCache>>,
@@ -99,7 +98,7 @@ pub(crate) struct VortexOpener {
 }
 
 impl FileOpener for VortexOpener {
-    fn open(&self, file: PartitionedFile) -> DFResult<FileOpenFuture> {
+    fn open(&self, file_meta: FileMeta, file: PartitionedFile) -> DFResult<FileOpenFuture> {
         let session = self.session.clone();
         let metrics_registry = self.metrics_registry.clone();
         let labels = vec![
@@ -158,7 +157,7 @@ impl FileOpener for VortexOpener {
             // - Partition column values (e.g., date=2024-01-01)
             // - File-level statistics (min/max values per column)
             let mut file_pruner = file_pruning_predicate
-                .filter(|p| {
+                .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.
                     is_dynamic_physical_expr(p) || file.has_statistics()
@@ -269,13 +268,13 @@ impl FileOpener for VortexOpener {
             let projector = leftover_projection.make_projector(&stream_schema)?;
 
             // We share our layout readers with others partitions in the scan, so we can only need to read each layout in each file once.
-            let layout_reader = match layout_reader.entry(file.object_meta.location.clone()) {
+            let layout_reader = match layout_reader.entry(file_meta.object_meta.location.clone()) {
                 Entry::Occupied(mut occupied_entry) => {
                     if let Some(reader) = occupied_entry.get().upgrade() {
-                        tracing::trace!("reusing layout reader for {}", occupied_entry.key());
+                        log::trace!("reusing layout reader for {}", occupied_entry.key());
                         reader
                     } else {
-                        tracing::trace!("creating layout reader for {}", occupied_entry.key());
+                        log::trace!("creating layout reader for {}", occupied_entry.key());
                         let reader = vxf.layout_reader().map_err(|e| {
                             DataFusionError::Execution(format!(
                                 "Failed to create layout reader: {e}"
@@ -286,7 +285,7 @@ impl FileOpener for VortexOpener {
                     }
                 }
                 Entry::Vacant(vacant_entry) => {
-                    tracing::trace!("creating layout reader for {}", vacant_entry.key());
+                    log::trace!("creating layout reader for {}", vacant_entry.key());
                     let reader = vxf.layout_reader().map_err(|e| {
                         DataFusionError::Execution(format!("Failed to create layout reader: {e}"))
                     })?;
@@ -296,29 +295,22 @@ impl FileOpener for VortexOpener {
                 }
             };
 
-            let mut scan_builder = ScanBuilder::new(session.clone(), layout_reader);
-
-            if let Some(extensions) = file.extensions
-                && let Some(vortex_plan) = extensions.downcast_ref::<VortexAccessPlan>()
-            {
-                scan_builder = vortex_plan.apply_to_builder(scan_builder);
-            }
-
-            if let Some(file_range) = file.range {
+            let mut scan_builder = ScanBuilder::new(session, layout_reader);
+            if let Some(file_range) = file_meta.range {
                 scan_builder = apply_byte_range(
                     file_range,
-                    file.object_meta.size,
+                    file_meta.object_meta.size,
                     vxf.row_count(),
                     scan_builder,
                 );
             }
 
             let filter = filter
                 .and_then(|f| {
-                    // Verify that all filters we've accepted from DataFusion get pushed down.
-                    // This will only fail if the user has not configured a suitable
-                    // PhysicalExprAdapterFactory on the file source to handle rewriting the
-                    // expression to handle missing/reordered columns in the Vortex file.
+                    let exprs = split_conjunction(&f)
+                        .into_iter()
+                        .filter(|expr| can_be_pushed_down(expr, &predicate_file_schema))
+                        .collect::<Vec<_>>();
 
                     let (pushed, unpushed): (Vec<PhysicalExprRef>, Vec<PhysicalExprRef>) =
                         split_conjunction(&f)
@@ -391,7 +383,7 @@ impl FileOpener for VortexOpener {
                 .map_err(move |e: VortexError| {
                     DataFusionError::External(Box::new(e.with_context(format!(
                         "Failed to read Vortex file: {}",
-                        file.object_meta.location
+                        file_meta.object_meta.location
                     ))))
                 })
                 .try_flatten()
@@ -405,17 +397,21 @@ impl FileOpener for VortexOpener {
                 .boxed();
 
             if let Some(file_pruner) = file_pruner {
-                Ok(PrunableStream::new(file_pruner, stream).boxed())
+                Ok(Box::pin(EarlyStoppingStream::new(
+                    stream,
+                    file_pruner,
+                    Count::new(),
+                )))
             } else {
-                Ok(stream)
+                Ok(Box::pin(stream))
             }
         }
         .in_current_span()
         .boxed())
     }
 }
 
-/// If the file has a [`FileRange`], we translate it into a row range in the file for the scan.
+/// If the file has a [`FileRange`](datafusion::datasource::listing::FileRange), we translate it into a row range in the file for the scan.
 fn apply_byte_range(
     file_range: FileRange,
     total_size: u64,
@@ -447,7 +443,6 @@ mod tests {
     use std::sync::Arc;
     use std::sync::LazyLock;
 
-    use arrow_schema::Field;
     use arrow_schema::Fields;
     use arrow_schema::SchemaRef;
     use datafusion::arrow::array::DictionaryArray;
@@ -458,7 +453,6 @@ mod tests {
     use datafusion::arrow::datatypes::Schema;
     use datafusion::arrow::datatypes::UInt32Type;
     use datafusion::arrow::util::display::FormatOptions;
-    use datafusion::arrow::util::pretty::pretty_format_batches_with_options;
     use datafusion::common::record_batch;
     use datafusion::logical_expr::col;
     use datafusion::logical_expr::lit;
@@ -595,7 +589,7 @@ mod tests {
 
         // filter matches partition value
         let filter = col("part").eq(lit(1));
-        let filter = logical2physical(&filter, table_schema.table_schema());
+        let filter = logical2physical(&filter, table_schema.as_ref());
 
         let opener = make_opener(object_store.clone(), table_schema.clone(), Some(filter));
         let stream = opener.open(file.clone()).unwrap().await.unwrap();
@@ -604,11 +598,11 @@ mod tests {
         let num_batches = data.len();
         let num_rows = data.iter().map(|rb| rb.num_rows()).sum::<usize>();
 
-        assert_eq!((num_batches, num_rows), (1, 3));
+        assert_eq!((num_batches, num_rows), expected_result1);
 
         // filter doesn't matches partition value
         let filter = col("part").eq(lit(2));
-        let filter = logical2physical(&filter, table_schema.table_schema());
+        let filter = logical2physical(&filter, table_schema.as_ref());
 
         let opener = make_opener(object_store.clone(), table_schema.clone(), Some(filter));
         let stream = opener.open(file.clone()).unwrap().await.unwrap();
@@ -625,29 +619,18 @@ mod tests {
     #[tokio::test]
     async fn test_open_files_different_table_schema() -> anyhow::Result<()> {
         let object_store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+        let file1_path = "/path/file1.vortex";
+        let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap();
+        let data_size1 = write_arrow_to_vortex(object_store.clone(), file1_path, batch1).await?;
+        let file1 = PartitionedFile::new(file1_path.to_string(), data_size1);
 
-        let file1 = {
-            let file1_path = "/path/file1.vortex";
-            let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap();
-            let data_size1 =
-                write_arrow_to_vortex(object_store.clone(), file1_path, batch1).await?;
-            PartitionedFile::new(file1_path.to_string(), data_size1)
-        };
-
-        let file2 = {
-            let file2_path = "/path/file2.vortex";
-            let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap();
-            let data_size2 =
-                write_arrow_to_vortex(object_store.clone(), file2_path, batch2).await?;
-            PartitionedFile::new(file2_path.to_string(), data_size2)
-        };
+        let file2_path = "/path/file2.vortex";
+        let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap();
+        let data_size2 = write_arrow_to_vortex(object_store.clone(), file2_path, batch2).await?;
+        let file2 = PartitionedFile::new(file1_path.to_string(), data_size1);
 
         // Table schema has can accommodate both files
-        let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new(
-            "a",
-            DataType::Int32,
-            true,
-        )])));
+        let table_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
 
         let make_opener = |filter| VortexOpener {
             partition: 1,
@@ -670,10 +653,12 @@ mod tests {
         };
 
         let filter = col("a").lt(lit(100_i32));
-        let filter = logical2physical(&filter, table_schema.table_schema());
+        let filter = logical2physical(&filter, table_schema.as_ref());
 
         let opener1 = make_opener(filter.clone());
-        let stream = opener1.open(file1)?.await?;
+        let stream = opener1
+            .open(make_meta(file1_path, data_size1), file1)?
+            .await?;
 
         let format_opts = FormatOptions::new().with_types_info(true);
 
@@ -690,7 +675,9 @@ mod tests {
         ");
 
         let opener2 = make_opener(filter.clone());
-        let stream = opener2.open(file2)?.await?;
+        let stream = opener2
+            .open(make_meta(file2_path, data_size2), file2)?
+            .await?;
 
         let data = stream.try_collect::<Vec<_>>().await?;
         assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r"
@@ -806,7 +793,7 @@ mod tests {
         let data_size = write_arrow_to_vortex(object_store.clone(), file_path, batch).await?;
 
         // Table schema has an extra utf8 field.
-        let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new(
+        let table_schema = Arc::new(Schema::new(vec![Field::new(
             "my_struct",
             DataType::Struct(Fields::from(vec![
                 Field::new(
@@ -822,23 +809,25 @@ mod tests {
                 Field::new("field3", DataType::Utf8, true),
             ])),
             true,
-        )])));
+        )]));
 
-        let opener = make_opener(
-            object_store.clone(),
-            table_schema.clone(),
-            // expression references my_struct column which has different fields in each
-            // field.
-            Some(logical2physical(
+        let opener = VortexOpener {
+            session: SESSION.clone(),
+            object_store: object_store.clone(),
+            projection: None,
+            filter: Some(logical2physical(
                 &col("my_struct").is_not_null(),
-                table_schema.table_schema(),
+                &table_schema,
             )),
         );
 
         // The opener should be able to open the file with a filter on the
         // struct column.
         let data = opener
-            .open(PartitionedFile::new(file_path.to_string(), data_size))?
+            .open(
+                make_meta(file_path, data_size),
+                PartitionedFile::new(file_path.to_string(), data_size),
+            )?
             .await?
             .try_collect::<Vec<_>>()
             .await?;
diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs
@@ -3,12 +3,11 @@
 
 use std::any::Any;
 use std::fmt::Formatter;
-use std::sync::Arc;
-use std::sync::Weak;
+use std::sync::{Arc, Weak};
 
 use datafusion_common::Result as DFResult;
 use datafusion_common::config::ConfigOptions;
-use datafusion_datasource::TableSchema;
+use datafusion_common::{Result as DFResult, Statistics};
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_datasource::file_stream::FileOpener;
@@ -18,12 +17,11 @@ use datafusion_physical_expr::conjunction;
 use datafusion_physical_expr::projection::ProjectionExprs;
 use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
-use datafusion_physical_plan::DisplayFormatType;
-use datafusion_physical_plan::PhysicalExpr;
-use datafusion_physical_plan::filter_pushdown::FilterPushdownPropagation;
-use datafusion_physical_plan::filter_pushdown::PushedDown;
-use datafusion_physical_plan::filter_pushdown::PushedDownPredicate;
+use datafusion_physical_plan::filter_pushdown::{
+    FilterPushdownPropagation, PushedDown, PushedDownPredicate,
+};
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
+use datafusion_physical_plan::{DisplayFormatType, PhysicalExpr};
 use object_store::ObjectStore;
 use object_store::path::Path;
 use vortex::error::VortexExpect;
