Add ScanThreading to stop double-counting scan concurrency under DataFusion

vortex-datafusion/src/persistent/format.rs

@@ -219,12 +219,14 @@ extensions_options! {
         /// Enabled by default.
         pub predicate_pushdown: bool, default = true
         /// The intra-partition scan concurrency, controlling the number of row splits to process
-        /// concurrently per-thread within each file.
+        /// concurrently within each file. DataFusion drives the parallelism across partitions,
+        /// so this is not scaled by the core count.
         ///
         /// This does not affect the overall parallelism
         /// across partitions, which is controlled by DataFusion's execution configuration.
         ///
-        /// Leave as `None` to use Vortex's scan default. Override per session
+        /// Leave as `None` to split the default budget of a machine-wide Vortex scan evenly
+        /// across the scan's partitions. Override per session
         /// with `SET vortex.scan_concurrency = <n>`, or per table with
         /// `OPTIONS(scan_concurrency '<n>')`.
         pub scan_concurrency: Option<usize>, default = None

vortex-datafusion/src/persistent/opener.rs

@@ -52,6 +52,7 @@ use vortex::file::OpenOptionsSessionExt;
 use vortex::io::InstrumentedReadAt;
 use vortex::layout::LayoutReader;
 use vortex::layout::scan::scan_builder::ScanBuilder;
+use vortex::layout::scan::threading::ScanThreading;
 use vortex::metrics::Label;
 use vortex::metrics::MetricsRegistry;
 use vortex::session::VortexSession;
@@ -395,7 +396,10 @@ impl FileOpener for VortexOpener {
                 }
             };
 
-            let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader));
+            // DataFusion already runs one file stream per partition, so the scan must not
+            // scale its concurrency by the core count again.
+            let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader))
+                .with_threading(ScanThreading::External);
 
             if let Some(vortex_plan) = file.extensions.get::<VortexAccessPlan>() {
                 scan_builder = vortex_plan.apply_to_builder(scan_builder);

vortex-datafusion/src/persistent/source.rs

@@ -32,10 +32,12 @@ use object_store::ObjectStore;
 use object_store::path::Path;
 use vortex::file::VORTEX_FILE_EXTENSION;
 use vortex::layout::LayoutReader;
+use vortex::layout::scan::scan_builder::DEFAULT_SCAN_CONCURRENCY;
 use vortex::metrics::DefaultMetricsRegistry;
 use vortex::metrics::MetricsRegistry;
 use vortex::session::VortexSession;
 use vortex_utils::aliases::dash_map::DashMap;
+use vortex_utils::parallelism::get_available_parallelism;
 
 use super::opener::NaturalSplits;
 use super::opener::VortexOpener;
@@ -306,9 +308,10 @@ impl VortexSource {
         self
     }
 
-    /// Sets the per-file Vortex scan concurrency.
+    /// Sets the number of row splits each file scan keeps in flight.
     ///
-    /// This is separate from DataFusion's partition-level parallelism.
+    /// This is separate from DataFusion's partition-level parallelism. When unset, the default
+    /// budget of a machine-wide scan is split evenly across the scan's partitions.
     pub fn with_scan_concurrency(mut self, scan_concurrency: usize) -> Self {
         self.options.scan_concurrency = Some(scan_concurrency);
         self
@@ -365,13 +368,30 @@ impl VortexSource {
             expression_converter: Arc::clone(&self.expression_converter),
             file_metadata_cache: self.file_metadata_cache.clone(),
             projection_pushdown: self.options.projection_pushdown,
-            scan_concurrency: self.options.scan_concurrency,
+            scan_concurrency: Some(self.options.scan_concurrency.unwrap_or_else(|| {
+                partition_scan_concurrency(
+                    base_config.file_groups.len(),
+                    get_available_parallelism().unwrap_or(1),
+                )
+            })),
         };
 
         Ok(opener)
     }
 }
 
+/// The default number of splits each partition's scan keeps in flight.
+///
+/// DataFusion runs one file stream per partition, and each stream scans with
+/// [`ScanThreading::External`](vortex::layout::scan::threading::ScanThreading::External). Splitting
+/// the budget of a single machine-wide scan across the partitions keeps the total number of splits
+/// in flight the same however many partitions DataFusion plans.
+fn partition_scan_concurrency(partitions: usize, cores: usize) -> usize {
+    (DEFAULT_SCAN_CONCURRENCY * cores)
+        .div_ceil(partitions.max(1))
+        .max(1)
+}
+
 impl FileSource for VortexSource {
     fn create_file_opener(
         &self,
@@ -570,11 +590,26 @@ mod tests {
     use datafusion_physical_expr::expressions as df_expr;
     use datafusion_physical_expr::expressions::Column;
     use object_store::memory::InMemory;
+    use rstest::rstest;
     use vortex::VortexSessionDefault;
 
     use super::*;
     use crate::convert::exprs::ProcessedProjection;
 
+    #[rstest]
+    #[case::one_partition_per_core(96, 96, 4)]
+    #[case::single_partition(1, 96, 384)]
+    #[case::fewer_partitions_than_cores(16, 96, 24)]
+    #[case::more_partitions_than_cores(200, 96, 2)]
+    #[case::no_partitions(0, 8, 32)]
+    fn partition_scan_concurrency_splits_machine_budget(
+        #[case] partitions: usize,
+        #[case] cores: usize,
+        #[case] expected: usize,
+    ) {
+        assert_eq!(partition_scan_concurrency(partitions, cores), expected);
+    }
+
     struct TrackingExpressionConverter {
         inner: DefaultExpressionConverter,
     }

vortex-duckdb/src/file_reader.rs

@@ -26,6 +26,7 @@ use vortex::io::runtime::BlockingRuntime as _;
 use vortex::io::std_file::StdFileSystem;
 use vortex::layout::LayoutReaderRef;
 use vortex::layout::scan::scan_builder::ScanBuilder;
+use vortex::layout::scan::threading::ScanThreading;
 use vortex::mask::Mask;
 use vortex::session::SessionExt as _;
 
@@ -200,6 +201,7 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor
     let reader = Arc::clone(file.reader());
     let filter = &global.filter;
     let builder = ScanBuilder::new(SESSION.clone(), reader)
+        .with_threading(ScanThreading::External)
         .with_projection(global.projection.clone())
         .with_some_filter(filter.filter.clone())
         .with_selection(filter.row_selection.clone());

vortex-layout/src/scan/mod.rs

@@ -13,6 +13,7 @@ mod splits;
 mod tasks;
 #[cfg(test)]
 mod test;
+pub mod threading;
 
 /// A heuristic for an ideal split size.
 ///

vortex-layout/src/scan/repeated_scan.rs

@@ -33,7 +33,6 @@ use vortex_io::session::RuntimeSessionExt;
 use vortex_scan::row_mask::RowMask;
 use vortex_scan::selection::Selection;
 use vortex_session::VortexSession;
-use vortex_utils::parallelism::get_available_parallelism;
 
 use crate::LayoutReaderRef;
 use crate::scan::filter::FilterExpr;
@@ -44,6 +43,7 @@ use crate::scan::tasks::SplitFuture;
 use crate::scan::tasks::TaskContext;
 use crate::scan::tasks::split_exec;
 use crate::scan::tasks::split_projection;
+use crate::scan::threading::ScanThreading;
 
 /// A projected subset (by indices, range, and filter) of rows from a Vortex data source.
 ///
@@ -61,8 +61,11 @@ pub struct RepeatedScan<A> {
     selection: Selection,
     /// The natural splits of the file.
     splits: Splits,
-    /// The number of splits to make progress on concurrently **per-thread**.
+    /// The number of splits to make progress on concurrently, per thread when
+    /// [`ScanThreading::Internal`] and per scan when [`ScanThreading::External`].
     concurrency: usize,
+    /// Whether the scan scales its concurrency by the available parallelism.
+    threading: ScanThreading,
     /// Function to apply to each [`ArrayRef`] within the spawned split tasks.
     map_fn: MapFn<A>,
     /// Maximal number of rows to read (after filtering).
@@ -137,6 +140,7 @@ impl<A: 'static + Send> RepeatedScan<A> {
         selection: Selection,
         splits: Splits,
         concurrency: usize,
+        threading: ScanThreading,
         map_fn: MapFn<A>,
         limit: Option<ScanLimit>,
         dtype: DType,
@@ -151,6 +155,7 @@ impl<A: 'static + Send> RepeatedScan<A> {
             selection,
             splits,
             concurrency,
+            threading,
             map_fn,
             limit,
             dtype,
@@ -240,7 +245,7 @@ impl<A: 'static + Send> RepeatedScan<A> {
         &self,
         row_range: Option<Range<u64>>,
     ) -> VortexResult<BoxStream<'static, VortexResult<A>>> {
-        let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1);
+        let concurrency = self.threading.splits_in_flight(self.concurrency);
         let handle = self.session.handle();
         let ordered = self.ordered;
         let ctx = self.task_context();

vortex-layout/src/scan/scan_builder.rs

@@ -42,6 +42,12 @@ use crate::scan::repeated_scan::RepeatedScan;
 use crate::scan::split_by::SplitBy;
 use crate::scan::splits::Splits;
 use crate::scan::splits::attempt_split_ranges;
+use crate::scan::threading::ScanThreading;
+
+/// The default number of splits a scan keeps in flight per worker thread.
+///
+/// Four allows for some I/O lookahead without too much impact on work-stealing.
+pub const DEFAULT_SCAN_CONCURRENCY: usize = 4;
 
 /// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs.
 ///
@@ -70,8 +76,11 @@ pub struct ScanBuilder<A> {
     /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare)
     /// uses them instead of walking the layout.
     natural_splits: Option<Arc<[u64]>>,
-    /// The number of splits to make progress on concurrently **per-thread**.
+    /// The number of splits to make progress on concurrently, per thread when
+    /// [`ScanThreading::Internal`] and per scan when [`ScanThreading::External`].
     concurrency: usize,
+    /// Whether the scan scales its concurrency by the available parallelism.
+    threading: ScanThreading,
     /// Function to apply to each [`ArrayRef`] within the spawned split tasks.
     map_fn: Arc<dyn Fn(ArrayRef) -> VortexResult<A> + Send + Sync>,
     metrics_registry: Option<Arc<dyn MetricsRegistry>>,
@@ -96,9 +105,8 @@ impl ScanBuilder<ArrayRef> {
             selection: Default::default(),
             split_by: SplitBy::default(),
             natural_splits: None,
-            // We default to four tasks per worker thread, which allows for some I/O lookahead
-            // without too much impact on work-stealing.
-            concurrency: 4,
+            concurrency: DEFAULT_SCAN_CONCURRENCY,
+            threading: ScanThreading::default(),
             map_fn: Arc::new(Ok),
             metrics_registry: None,
             limit: None,
@@ -225,14 +233,30 @@ impl<A: 'static + Send> ScanBuilder<A> {
         self.concurrency
     }
 
-    /// The number of row splits to make progress on concurrently per-thread, must
-    /// be greater than 0.
+    /// The number of row splits to make progress on concurrently, must be greater than 0.
+    ///
+    /// This is per thread under [`ScanThreading::Internal`] and per scan under
+    /// [`ScanThreading::External`].
     pub fn with_concurrency(mut self, concurrency: usize) -> Self {
         assert!(concurrency > 0);
         self.concurrency = concurrency;
         self
     }
 
+    /// Returns who drives the parallelism of the scan.
+    pub fn threading(&self) -> ScanThreading {
+        self.threading
+    }
+
+    /// Set who drives the parallelism of the scan.
+    ///
+    /// Use [`ScanThreading::External`] when a query engine already runs one scan per worker
+    /// thread, so the per-thread concurrency is not scaled by the core count a second time.
+    pub fn with_threading(mut self, threading: ScanThreading) -> Self {
+        self.threading = threading;
+        self
+    }
+
     /// Add or clear the metrics registry used by scan execution.
     pub fn with_some_metrics_registry(mut self, metrics: Option<Arc<dyn MetricsRegistry>>) -> Self {
         self.metrics_registry = metrics;
@@ -293,6 +317,7 @@ impl<A: 'static + Send> ScanBuilder<A> {
             split_by: self.split_by,
             natural_splits: self.natural_splits,
             concurrency: self.concurrency,
+            threading: self.threading,
             metrics_registry: self.metrics_registry,
             limit: self.limit,
             row_offset: self.row_offset,
@@ -357,6 +382,7 @@ impl<A: 'static + Send> ScanBuilder<A> {
             self.selection,
             splits,
             self.concurrency,
+            self.threading,
             self.map_fn,
             self.limit,
             dtype,

vortex-layout/src/scan/threading.rs

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_utils::parallelism::get_available_parallelism;
+
+/// Who drives the parallelism of a scan.
+///
+/// A scan stream keeps a number of splits in flight at once. When the scan owns the machine it
+/// scales that number by the available parallelism so every core has work. When a query engine
+/// such as DataFusion or DuckDB already partitions the work across its own threads, each scan is
+/// one of many running at once, and scaling by the core count again would multiply the in-flight
+/// splits by the core count a second time.
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
+pub enum ScanThreading {
+    /// The scan is the only consumer of the machine and keeps `concurrency` splits in flight per
+    /// available core.
+    #[default]
+    Internal,
+    /// An external engine runs many scans in parallel, so each scan keeps only `concurrency`
+    /// splits in flight.
+    External,
+}
+
+impl ScanThreading {
+    /// The number of splits to keep in flight for the given per-thread `concurrency`.
+    pub(crate) fn splits_in_flight(self, concurrency: usize) -> usize {
+        match self {
+            ScanThreading::Internal => concurrency * get_available_parallelism().unwrap_or(1),
+            ScanThreading::External => concurrency,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_utils::parallelism::get_available_parallelism;
+
+    use super::ScanThreading;
+
+    #[test]
+    fn external_threading_does_not_scale_by_cores() {
+        assert_eq!(ScanThreading::External.splits_in_flight(4), 4);
+        assert_eq!(
+            ScanThreading::Internal.splits_in_flight(4),
+            4 * get_available_parallelism().unwrap_or(1)
+        );
+    }
+}