Fix and speed up the V2 scan's planner chain

vortex-datafusion/src/persistent/opener.rs

@@ -137,9 +137,6 @@ pub(crate) struct VortexOpener {
     pub projection_pushdown: bool,
     pub scan_concurrency: Option<usize>,
     pub scan_v2: bool,
-    /// 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,
 }
 
 impl FileOpener for VortexOpener {
@@ -178,7 +175,6 @@ impl FileOpener for VortexOpener {
         let file_schemas = Arc::clone(&self.file_schemas);
         let has_output_ordering = self.has_output_ordering;
         let scan_concurrency = self.scan_concurrency;
-        let partitions = self.partitions.max(1);
         let scan_v2 = self.scan_v2;
 
         let expr_converter = Arc::clone(&self.expression_converter);
@@ -444,6 +440,11 @@ impl FileOpener for VortexOpener {
                 let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file)
                     .with_projection(scan_projection)
                     .with_some_filter(native_filter);
+                // How many partitions split the file, assuming they cut it evenly.
+                let readers = byte_range.as_ref().map_or(1, |bytes| {
+                    let len = (bytes.end - bytes.start).max(1);
+                    usize::try_from(file.object_meta.size.div_ceil(len)).unwrap_or(usize::MAX)
+                });
                 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());
@@ -459,14 +460,16 @@ impl FileOpener for VortexOpener {
                 {
                     builder = builder.with_selection(selection.clone());
                 }
-                // Unless set, chunks run beside each other only as far as the engine's partitions
-                // leave cores idle: at one partition per core they run in the scan's own run.
-                let concurrency = scan_concurrency.unwrap_or_else(|| {
+                // Unless set, chunks run beside each other only as far as the cores outnumber the
+                // partitions reading this file: a file split a partition per core is read in each
+                // scan's own run, a file read whole by one partition on every core.
+                let concurrency = scan_concurrency.filter(|&c| c > 0).unwrap_or_else(|| {
                     get_available_parallelism()
                         .unwrap_or(1)
-                        .div_ceil(partitions)
+                        .div_ceil(readers)
                         .max(1)
                 });
+                builder = builder.with_concurrency(concurrency);
                 let handle = vxf.session().handle();
                 // Lowering and planning a wide file is real CPU: it runs where compute runs.
                 let scan = handle
@@ -1147,7 +1150,6 @@ mod tests {
             layout_readers: Default::default(),
             segment_sources: Default::default(),
             scan_caches: Default::default(),
-            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1557,7 +1559,6 @@ mod tests {
             layout_readers: Default::default(),
             segment_sources: Default::default(),
             scan_caches: Default::default(),
-            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1649,7 +1650,6 @@ mod tests {
             layout_readers: Default::default(),
             segment_sources: Default::default(),
             scan_caches: Default::default(),
-            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1832,7 +1832,6 @@ mod tests {
             layout_readers: Default::default(),
             segment_sources: Default::default(),
             scan_caches: Default::default(),
-            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1897,7 +1896,6 @@ mod tests {
             layout_readers: Default::default(),
             segment_sources: Default::default(),
             scan_caches: Default::default(),
-            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -2113,7 +2111,6 @@ mod tests {
             layout_readers: Default::default(),
             segment_sources: Default::default(),
             scan_caches: Default::default(),
-            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,

vortex-datafusion/src/persistent/scan_v2.rs

@@ -25,22 +25,21 @@ use futures::stream::BoxStream;
 use vortex::array::ArrayRef;
 use vortex::error::VortexResult;
 use vortex::error::vortex_err;
-use vortex::io::request::IoSource;
+use vortex::io::request::IoService;
 use vortex::io::runtime::Handle;
 use vortex::layout::scan::v2::PreparedScan;
 use vortex::scan::planning::driver::Progress;
 use vortex::scan::planning::driver::Run;
 use vortex::scan::planning::planner::Planner;
 use vortex::scan::planning::planner::WorkScope;
-use vortex_utils::parallelism::get_available_parallelism;
 
 /// Maps a scan's batches to Arrow on the task that produced them.
 type MapBatch = Arc<dyn Fn(ArrayRef) -> VortexResult<RecordBatch> + Send + Sync>;
 
 /// Runs one partition's V2 scan on DataFusion's runtime.
 pub(crate) struct ScanDriver {
     handle: Handle,
-    /// Chunks run beside each other per worker thread; at 1 they run on the scan's own run.
+    /// The scan's chunks run at once; at 1 they run on the scan's own run.
     concurrency: usize,
     /// Whether batches are yielded in row order, which holds them until the scan is done.
     ordered: bool,
@@ -64,16 +63,16 @@ impl ScanDriver {
         map: impl Fn(ArrayRef) -> VortexResult<RecordBatch> + Send + Sync + 'static,
     ) -> BoxStream<'static, VortexResult<RecordBatch>> {
         let map: MapBatch = Arc::new(map);
-        let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1);
-        let inline = in_flight <= 1;
+        let in_flight = self.concurrency;
+        let inline = self.concurrency <= 1;
         let scope = WorkScope {
             file_ordinal: 0,
             rows: scan.rows().clone(),
         };
         let io = scan.io();
         let (chunk_handle, chunk_io, chunk_map) =
             (self.handle.clone(), Arc::clone(&io), Arc::clone(&map));
-        let batches = run_stream(self.handle, root, scope, io.session(), map, inline)
+        let batches = run_stream(self.handle, root, scope, Arc::clone(&io), map, inline)
             .map_ok(move |yielded| match yielded {
                 Yielded::Batch(start, batch) => {
                     stream::once(future::ready(Ok((start, batch)))).boxed()
@@ -82,7 +81,7 @@ impl ScanDriver {
                     &chunk_handle,
                     chunk,
                     scope,
-                    chunk_io.session(),
+                    Arc::clone(&chunk_io),
                     Arc::clone(&chunk_map),
                 ),
             })
@@ -128,7 +127,8 @@ struct RunState {
     handle: Handle,
     /// Taken while a CPU step runs on its task and put back with its result.
     run: Option<Run>,
-    io: Arc<dyn IoSource>,
+    /// Gives every root of the run a session of its own: retiring a root clears its session.
+    io: Arc<dyn IoService>,
     map: MapBatch,
     /// Whether a spawned chunk runs on this run, as another root, instead of being yielded.
     inline: bool,
@@ -143,12 +143,12 @@ fn run_stream(
     handle: Handle,
     root: Box<dyn Planner>,
     scope: WorkScope,
-    io: Arc<dyn IoSource>,
+    io: Arc<dyn IoService>,
     map: MapBatch,
     inline: bool,
 ) -> BoxStream<'static, VortexResult<Yielded>> {
     let mut run = Run::new();
-    run.admit(root, scope, Arc::clone(&io));
+    run.admit(root, scope, io.session());
     let state = RunState {
         handle,
         run: Some(run),
@@ -187,7 +187,7 @@ fn run_stream(
                     return Ok(Some((Yielded::Batch(start, batch), state)));
                 }
                 Step::Spawn(scope, child) if state.inline => {
-                    run.admit(child, scope, Arc::clone(&state.io));
+                    run.admit(child, scope, state.io.session());
                 }
                 Step::Spawn(scope, child) => {
                     return Ok(Some((Yielded::Spawn(scope, child), state)));
@@ -211,13 +211,13 @@ fn chunk_stream(
     handle: &Handle,
     root: Box<dyn Planner>,
     scope: WorkScope,
-    io: Arc<dyn IoSource>,
+    io: Arc<dyn IoService>,
     map: MapBatch,
 ) -> BoxStream<'static, VortexResult<(u64, RecordBatch)>> {
     let (mut send, recv) = mpsc::channel(1);
     let task = handle.spawn(async move {
         let mut run = Run::new();
-        run.admit(root, scope, Arc::clone(&io));
+        run.admit(root, scope, io.session());
         loop {
             match run.advance()? {
                 Progress::Batch(batch) => {
@@ -229,7 +229,7 @@ fn chunk_stream(
                     }
                 }
                 Progress::Spawn(scope, child) => {
-                    run.admit(child, scope, Arc::clone(&io));
+                    run.admit(child, scope, io.session());
                 }
                 Progress::RootDone(_) if run.live_roots() > 0 => {}
                 Progress::RootDone(_) | Progress::Idle => return Ok(()),

vortex-datafusion/src/persistent/source.rs

@@ -49,6 +49,31 @@ use crate::convert::exprs::ExpressionConverter;
 use crate::persistent::reader::DefaultVortexReaderFactory;
 use crate::persistent::reader::VortexReaderFactory;
 
+/// The fewest rows a filtered V2 scan partition reads when the engine splits files. Each
+/// partition proves zone maps, plans its filter and compiles its plans once, so a partition much
+/// smaller than a few output windows spends more on that than on reading rows. An unfiltered scan
+/// keeps the engine's partitions: every row it reads flows into the operators above it, which
+/// run a task per partition.
+const V2_MIN_ROWS_PER_PARTITION: usize = 256 * 1024;
+
+/// Caps `target_partitions` so no V2 partition reads fewer than [`V2_MIN_ROWS_PER_PARTITION`]
+/// rows, when the files' row counts are known.
+fn v2_target_partitions(target_partitions: usize, config: &FileScanConfig) -> usize {
+    let rows = config
+        .file_groups
+        .iter()
+        .flat_map(|group| group.iter())
+        .map(|file| {
+            file.statistics
+                .as_ref()
+                .and_then(|statistics| statistics.num_rows.get_value().copied())
+        })
+        .sum::<Option<usize>>();
+    rows.map_or(target_partitions, |rows| {
+        target_partitions.min(rows.div_ceil(V2_MIN_ROWS_PER_PARTITION).max(1))
+    })
+}
+
 /// File scan implementation for reading one or more `.vortex` files.
 ///
 /// `VortexSource` is the lower-level read component underneath
@@ -378,7 +403,6 @@ impl VortexSource {
             projection_pushdown: self.options.projection_pushdown,
             scan_concurrency: self.options.scan_concurrency,
             scan_v2: self.options.scan_v2,
-            partitions: base_config.file_groups.len(),
         };
 
         Ok(opener)
@@ -468,6 +492,11 @@ impl FileSource for VortexSource {
         if config.file_compression_type.is_compressed() {
             return Ok(None);
         }
+        let target_partitions = if self.options.scan_v2 && self.vortex_predicate.is_some() {
+            v2_target_partitions(target_partitions, config)
+        } else {
+            target_partitions
+        };
         let Some(file_groups) = FileGroupPartitioner::new()
             .with_target_partitions(target_partitions)
             .with_repartition_file_min_size(repartition_file_min_size)

vortex-datafusion/src/persistent/tests.rs

@@ -486,9 +486,10 @@ async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits(
     let session = VortexSession::default();
     let path = object_store::path::Path::parse("/split-aligned-repartition.vortex")?;
 
-    let chunk_1_len = 2_000;
-    let chunk_2_len = 5_000;
-    let chunk_3_len = 6_000;
+    // Enough rows that a V2 scan, which keeps a minimum of rows per partition, still splits.
+    let chunk_1_len = 200_000;
+    let chunk_2_len = 250_000;
+    let chunk_3_len = 300_000;
     let row_count = chunk_1_len + chunk_2_len + chunk_3_len;
 
     let chunk_1 = StructArray::try_new(
@@ -582,6 +583,38 @@ async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits(
     Ok(())
 }
 
+/// A filtered V2 scan of a file smaller than the minimum rows per partition reads it in one
+/// partition. An unfiltered scan keeps the engine's partitions: every row it reads flows into the
+/// operators above it.
+#[tokio::test]
+async fn v2_scan_keeps_partitions_unless_filtered() -> anyhow::Result<()> {
+    let store = Arc::new(InMemory::new()) as _;
+    let session = VortexSession::default();
+    let path = object_store::path::Path::parse("/small.vortex")?;
+    let table = StructArray::try_new(
+        ["value"].into(),
+        vec![Buffer::from_iter(0_i32..100_000).into_array()],
+        100_000,
+        Validity::NonNullable,
+    )?
+    .into_array();
+    let mut writer = ObjectStoreWrite::new(Arc::clone(&store), &path).await?;
+    session
+        .write_options()
+        .write(&mut writer, table.to_array_stream())
+        .await?;
+    writer.shutdown().await?;
+
+    let ctx = make_session(store, true, true);
+    let unfiltered = count_query_partitions(&ctx, "SELECT value FROM '/small.vortex'").await?;
+    let filtered =
+        count_query_partitions(&ctx, "SELECT value FROM '/small.vortex' WHERE value > 10").await?;
+
+    assert_eq!(unfiltered, 4);
+    assert_eq!(filtered, 1);
+    Ok(())
+}
+
 /// Roundtrip an `arrow.uuid` extension column through a Vortex file: write the column directly
 /// via the session-aware Arrow→Vortex conversion, then `SELECT *` and assert both the field
 /// metadata and the underlying values survive the trip.