diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index e39b0583abf..a09888d102d 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/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 = `, or per table with /// `OPTIONS(scan_concurrency '')`. pub scan_concurrency: Option, default = None diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 62455fd7b6b..d53af9a398d 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/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::() { scan_builder = vortex_plan.apply_to_builder(scan_builder); diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 37bde295a9a..439ba8230ab 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/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, } diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index f90e4d42742..43ac485e68b 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/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()); diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index e08e4cb3f31..d75c250de8b 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/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. /// diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 4200d00614a..422d14d8e9a 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/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 { 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, /// Maximal number of rows to read (after filtering). @@ -137,6 +140,7 @@ impl RepeatedScan { selection: Selection, splits: Splits, concurrency: usize, + threading: ScanThreading, map_fn: MapFn, limit: Option, dtype: DType, @@ -151,6 +155,7 @@ impl RepeatedScan { selection, splits, concurrency, + threading, map_fn, limit, dtype, @@ -240,7 +245,7 @@ impl RepeatedScan { &self, row_range: Option>, ) -> VortexResult>> { - 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(); diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 814c0759453..d33759150a2 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/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 { /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare) /// uses them instead of walking the layout. natural_splits: Option>, - /// 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 VortexResult + Send + Sync>, metrics_registry: Option>, @@ -96,9 +105,8 @@ impl ScanBuilder { 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 ScanBuilder { 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>) -> Self { self.metrics_registry = metrics; @@ -293,6 +317,7 @@ impl ScanBuilder { 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 ScanBuilder { self.selection, splits, self.concurrency, + self.threading, self.map_fn, self.limit, dtype, diff --git a/vortex-layout/src/scan/threading.rs b/vortex-layout/src/scan/threading.rs new file mode 100644 index 00000000000..c97c22b40b1 --- /dev/null +++ b/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) + ); + } +}