Reduce V2 scan filtering, I/O, and batch-tail copying

benchmarks/datafusion-bench/src/main.rs

@@ -120,7 +120,7 @@ struct Args {
     #[arg(long, value_delimiter = ',', value_parser = value_parser!(Format))]
     formats: Vec<Format>,
 
-    /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H).
+    /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H, TPC-DS and ClickBench).
     #[arg(long, action = clap::ArgAction::Set)]
     scan_v2: Option<bool>,
 
@@ -137,8 +137,10 @@ async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
     let opts = Opts::from(args.options);
     let scan_v2 = args.scan_v2.unwrap_or_else(|| {
-        matches!(args.benchmark, BenchmarkArg::TpcH)
-            && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0")
+        matches!(
+            args.benchmark,
+            BenchmarkArg::TpcH | BenchmarkArg::TpcDS | BenchmarkArg::ClickBench
+        ) && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0")
     });
 
     set_join_set_tracer(get_static_tracer())?;

vortex-datafusion/src/persistent/format.rs

@@ -4,6 +4,7 @@
 use std::fmt::Debug;
 use std::fmt::Formatter;
 use std::sync::Arc;
+use std::sync::Weak;
 
 use arrow_schema::DataType;
 use arrow_schema::Schema;
@@ -35,6 +36,7 @@ use datafusion_datasource::file_sink_config::FileSinkConfig;
 use datafusion_datasource::sink::DataSinkExec;
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry;
+use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_expr::dml::InsertOp;
 use datafusion_physical_expr::LexRequirement;
 use datafusion_physical_plan::ExecutionPlan;
@@ -67,9 +69,11 @@ use vortex::scalar::Scalar;
 use vortex::scalar::ScalarValue as VortexScalarValue;
 use vortex::session::VortexSession;
 use vortex_arrow::ArrowSessionExt;
+use vortex_utils::aliases::dash_map::DashMap;
 
 use super::cache::CachedVortexMetadata;
 use super::sink::VortexSink;
+use super::source::SharedScanState;
 use super::source::VortexSource;
 use crate::PrecisionExt as _;
 use crate::convert::ExpressionConverter;
@@ -127,6 +131,7 @@ const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usi
 /// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html
 /// [`ListingOptions`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingOptions.html
 pub struct VortexFormat {
+    scan_states: DashMap<(ObjectStoreUrl, usize), Weak<SharedScanState>>,
     session: VortexSession,
     opts: VortexTableOptions,
     expression_converter: Option<Arc<dyn ExpressionConverter>>,
@@ -459,12 +464,27 @@ impl VortexFormat {
     /// Creates a format with explicit [`VortexTableOptions`].
     pub fn new_with_options(session: VortexSession, opts: VortexTableOptions) -> Self {
         Self {
+            scan_states: DashMap::default(),
             session,
             opts,
             expression_converter: None,
         }
     }
 
+    fn shared_scan_state(
+        &self,
+        store_url: ObjectStoreUrl,
+        store: Arc<dyn ObjectStore>,
+    ) -> Arc<SharedScanState> {
+        let store_id = Arc::as_ptr(&store) as *const () as usize;
+        let mut entry = self.scan_states.entry((store_url, store_id)).or_default();
+        entry.upgrade().unwrap_or_else(|| {
+            let shared = Arc::new(SharedScanState::for_store(store));
+            *entry = Arc::downgrade(&shared);
+            shared
+        })
+    }
+
     /// Returns the format-specific configuration that will be copied into the
     /// [`VortexSource`] created for a scan.
     pub fn options(&self) -> &VortexTableOptions {
@@ -748,6 +768,11 @@ impl FileFormat for VortexFormat {
 
         source = source
             .with_file_metadata_cache(state.runtime_env().cache_manager.get_file_metadata_cache());
+        if source.options().scan_v2 && source.vortex_reader_factory.is_none() {
+            let store_url = file_scan_config.object_store_url.clone();
+            let store = state.runtime_env().object_store(&store_url)?;
+            source = source.with_scan_state(self.shared_scan_state(store_url, store));
+        }
 
         let conf = FileScanConfigBuilder::from(file_scan_config)
             .with_source(Arc::new(source))
@@ -821,6 +846,7 @@ mod tests {
     use datafusion_physical_expr::expressions as df_expr;
     use datafusion_physical_expr::projection::ProjectionExprs;
     use datafusion_physical_plan::filter_pushdown::PushedDown;
+    use object_store::memory::InMemory;
     use rstest::rstest;
     use vortex::array::stats::StatsSet;
     use vortex::dtype::Nullability;
@@ -834,6 +860,28 @@ mod tests {
     use crate::convert::ProcessedProjection;
     use crate::convert::scalar_from_df;
 
+    #[test]
+    fn scan_state_is_scoped_to_store_instance_and_released() {
+        let format = VortexFormat::new(VortexSession::default());
+        let url = ObjectStoreUrl::local_filesystem();
+        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
+        let weak_store = Arc::downgrade(&store);
+        let first = format.shared_scan_state(url.clone(), Arc::clone(&store));
+        let second = format.shared_scan_state(url.clone(), Arc::clone(&store));
+        assert!(Arc::ptr_eq(&first, &second));
+
+        let other = format.shared_scan_state(url, Arc::new(InMemory::new()));
+        assert!(!Arc::ptr_eq(&first, &other));
+
+        let weak_state = Arc::downgrade(&first);
+        drop(store);
+        assert!(weak_store.upgrade().is_some());
+        drop(first);
+        drop(second);
+        assert!(weak_state.upgrade().is_none());
+        assert!(weak_store.upgrade().is_none());
+    }
+
     #[test]
     fn lossy_bound_cast_does_not_prove_constant() -> VortexResult<()> {
         let dtype = DType::Extension(

vortex-datafusion/src/persistent/opener.rs

@@ -45,6 +45,7 @@ use futures::FutureExt;
 use futures::StreamExt;
 use futures::TryStreamExt;
 use futures::stream;
+use object_store::ObjectMeta;
 use object_store::path::Path;
 use tracing::Instrument;
 use vortex::array::VortexSessionExecute;
@@ -83,6 +84,8 @@ use crate::metrics::PATH_LABEL;
 use crate::persistent::cache::CachedVortexMetadata;
 use crate::persistent::reader::VortexReaderFactory;
 use crate::persistent::scan_v2::ScanDriver;
+use crate::persistent::scan_v2::rebalance_batches;
+use crate::persistent::source::FileVersion;
 use crate::persistent::stream::PrunableStream;
 
 #[derive(Clone)]
@@ -120,10 +123,10 @@ pub(crate) struct VortexOpener {
     /// Partitions of one file read through one source, so a chunk straddling a partition
     /// boundary and the file's metadata segments are read once, as they are for V1 through the
     /// shared layout reader.
-    pub segment_sources: Arc<DashMap<Path, Weak<dyn SegmentSource>>>,
+    pub segment_sources: Arc<DashMap<FileVersion, Weak<dyn SegmentSource>>>,
     /// The V2 scan cache each open file's partitions share: its lowered layout, filter and
     /// projection plans, pruned zones and dictionary values, built once per file.
-    pub scan_caches: Arc<DashMap<Path, WeakScanCache>>,
+    pub scan_caches: Arc<DashMap<FileVersion, WeakScanCache>>,
     /// Shared V1 full-file natural splits keyed by file path.
     pub natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
     /// Shared physical Arrow schemas of each file, keyed by file path.
@@ -140,6 +143,7 @@ pub(crate) struct VortexOpener {
     /// How many partitions the engine runs the scan on, so a V2 scan parallelises inside a
     /// partition only as far as the engine has not.
     pub partitions: usize,
+    pub batch_size: usize,
 }
 
 impl FileOpener for VortexOpener {
@@ -179,6 +183,7 @@ impl FileOpener for VortexOpener {
         let has_output_ordering = self.has_output_ordering;
         let scan_concurrency = self.scan_concurrency;
         let partitions = self.partitions.max(1);
+        let batch_size = self.batch_size;
         let scan_v2 = self.scan_v2;
 
         let expr_converter = Arc::clone(&self.expression_converter);
@@ -277,6 +282,34 @@ impl FileOpener for VortexOpener {
                 return Ok(stream::empty().boxed());
             }
 
+            let byte_range = file
+                .range
+                .as_ref()
+                .map(|range| -> DFResult<Range<u64>> {
+                    Ok(Range {
+                        start: u64::try_from(range.start).map_err(|_| {
+                            exec_datafusion_err!("Vortex file range start is negative")
+                        })?,
+                        end: u64::try_from(range.end).map_err(|_| {
+                            exec_datafusion_err!("Vortex file range end is negative")
+                        })?,
+                    })
+                })
+                .transpose()?;
+
+            // Coarse row boundaries are independent of the projection and predicates. Empty
+            // byte partitions can finish before adapting or converting those expressions.
+            let v2_rows = if scan_v2 {
+                byte_range
+                    .clone()
+                    .map(|bytes| partition_row_range(bytes, file.object_meta.size, vxf.row_count()))
+            } else {
+                None
+            };
+            if v2_rows.as_ref().is_some_and(Range::is_empty) {
+                return Ok(stream::empty().boxed());
+            }
+
             // This is the expected arrow types of the actual columns in the file, which might have different types
             // from the unified logical schema or miss. Every partition opening this file computes
             // the same schema, so it is shared across them.
@@ -403,21 +436,6 @@ impl FileOpener for VortexOpener {
                 .transpose()?;
             let projector = leftover_projection.make_projector(&stream_schema)?;
 
-            let byte_range = file
-                .range
-                .as_ref()
-                .map(|range| -> DFResult<Range<u64>> {
-                    Ok(Range {
-                        start: u64::try_from(range.start).map_err(|_| {
-                            exec_datafusion_err!("Vortex file range start is negative")
-                        })?,
-                        end: u64::try_from(range.end).map_err(|_| {
-                            exec_datafusion_err!("Vortex file range end is negative")
-                        })?,
-                    })
-                })
-                .transpose()?;
-
             let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false);
             let residual_path = file.path().clone();
             let native_limit = limit
@@ -433,23 +451,16 @@ impl FileOpener for VortexOpener {
                 )?;
                 Ok(RecordBatch::from(arrow.as_struct().clone()))
             };
+            let rebalance = limit.is_none() && native_filter.is_none() && residual_filter.is_none();
             let stream = if scan_v2 {
                 let mut scan_file = vxf.scan_file();
-                scan_file.segments = shared_segment_source(
-                    &segment_sources,
-                    &file.object_meta.location,
-                    scan_file.segments,
-                );
-                scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta.location));
+                scan_file.segments =
+                    shared_segment_source(&segment_sources, &file.object_meta, scan_file.segments);
+                scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta));
                 let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file)
                     .with_projection(scan_projection)
                     .with_some_filter(native_filter);
-                if let Some(bytes) = byte_range {
-                    // Every partition maps the same byte boundary to the same row boundary.
-                    let rows = partition_row_range(bytes, file.object_meta.size, vxf.row_count());
-                    if rows.is_empty() {
-                        return Ok(stream::empty().boxed());
-                    }
+                if let Some(rows) = v2_rows {
                     builder = builder.with_row_range(rows);
                 }
                 if let Some(selection) = file
@@ -467,6 +478,7 @@ impl FileOpener for VortexOpener {
                         .div_ceil(partitions)
                         .max(1)
                 });
+                builder = builder.with_prefetch_batches(if concurrency <= 1 { 4 } else { 1 });
                 let handle = vxf.session().handle();
                 // Lowering and planning a wide file is real CPU: it runs where compute runs.
                 let scan = handle
@@ -495,7 +507,12 @@ impl FileOpener for VortexOpener {
                 };
 
                 let driver = ScanDriver::new(handle, concurrency, has_output_ordering);
-                Ok(driver.run(scan, root, map_batch))
+                let batches = driver.run(scan, root, map_batch);
+                Ok(if rebalance {
+                    rebalance_batches(batches, batch_size)
+                } else {
+                    batches
+                })
             } else {
                 // Share V1 layout readers across partitions of the file.
                 let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) {
@@ -693,10 +710,15 @@ fn contains_lambda(projection: &ProjectionExprs) -> DFResult<bool> {
     Ok(false)
 }
 
-/// Returns the scan cache other partitions of `path` are planning with, or registers a new one.
-/// A scan holds its cache until its stream ends, so the cache lives as long as one scan.
-fn shared_scan_cache(scan_caches: &DashMap<Path, WeakScanCache>, path: &Path) -> ScanCache {
-    match scan_caches.entry(path.clone()) {
+/// Returns the scan cache for this object version, or registers a new one.
+/// Live streams for this file version own the cache; the registry holds only a weak reference.
+/// The last stream releases the cache, even if the multi-file scan is still running.
+fn shared_scan_cache(
+    scan_caches: &DashMap<FileVersion, WeakScanCache>,
+    meta: &ObjectMeta,
+) -> ScanCache {
+    let path = &meta.location;
+    match scan_caches.entry(FileVersion(meta.clone())) {
         Entry::Occupied(mut entry) => {
             if let Some(shared) = entry.get().upgrade() {
                 tracing::trace!("reusing scan cache for {path}");
@@ -715,14 +737,15 @@ fn shared_scan_cache(scan_caches: &DashMap<Path, WeakScanCache>, path: &Path) ->
     }
 }
 
-/// Returns the segment source other partitions of `path` are reading through, or registers
+/// Returns the segment source other scans of this object version read through, or registers
 /// `fresh` as that source. Like the layout readers, a source lives only as long as one scan.
 fn shared_segment_source(
-    segment_sources: &DashMap<Path, Weak<dyn SegmentSource>>,
-    path: &Path,
+    segment_sources: &DashMap<FileVersion, Weak<dyn SegmentSource>>,
+    meta: &ObjectMeta,
     fresh: Arc<dyn SegmentSource>,
 ) -> Arc<dyn SegmentSource> {
-    match segment_sources.entry(path.clone()) {
+    let path = &meta.location;
+    match segment_sources.entry(FileVersion(meta.clone())) {
         Entry::Occupied(mut entry) => {
             if let Some(shared) = entry.get().upgrade() {
                 tracing::trace!("reusing segment source for {path}");
@@ -1148,6 +1171,7 @@ mod tests {
             segment_sources: Default::default(),
             scan_caches: Default::default(),
             partitions: 1,
+            batch_size: 8192,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1558,6 +1582,7 @@ mod tests {
             segment_sources: Default::default(),
             scan_caches: Default::default(),
             partitions: 1,
+            batch_size: 8192,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1650,6 +1675,7 @@ mod tests {
             segment_sources: Default::default(),
             scan_caches: Default::default(),
             partitions: 1,
+            batch_size: 8192,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1833,6 +1859,7 @@ mod tests {
             segment_sources: Default::default(),
             scan_caches: Default::default(),
             partitions: 1,
+            batch_size: 8192,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1898,6 +1925,7 @@ mod tests {
             segment_sources: Default::default(),
             scan_caches: Default::default(),
             partitions: 1,
+            batch_size: 8192,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -2114,6 +2142,7 @@ mod tests {
             segment_sources: Default::default(),
             scan_caches: Default::default(),
             partitions: 1,
+            batch_size: 8192,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,

vortex-datafusion/src/persistent/reader.rs

@@ -69,11 +69,14 @@ impl VortexReaderFactory for DefaultVortexReaderFactory {
         file: &PartitionedFile,
         session: &VortexSession,
     ) -> DFResult<Arc<dyn VortexReadAt>> {
-        Ok(Arc::new(ObjectStoreReadAt::new_with_allocator(
-            Arc::clone(&self.object_store),
-            file.path().clone(),
-            session.handle(),
-            session.allocator(),
-        )) as _)
+        Ok(Arc::new(
+            ObjectStoreReadAt::new_with_allocator(
+                Arc::clone(&self.object_store),
+                file.path().clone(),
+                session.handle(),
+                session.allocator(),
+            )
+            .with_cached_file_handle(),
+        ) as _)
     }
 }