From fae2626075c57469b60e14980eb477d9887f3068 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 15:01:46 +0100 Subject: [PATCH 01/22] Expose a V2 scan's planners and prepared plans; drive them elsewhere A planner says what it needs, compute or reads. Where either runs is an engine's decision, so this crate stops driving: `ScanBuilder::prepare` lowers and plans a scan without IO into a `PreparedScan`, which exposes its rows, selection, filter plans and IO service, and `PrunePlanner` and `ChunkPlanner` are public for an engine to compose, the prune planner continuing into whatever it is built with and passing `None` when no proof ran. The builder's ordering, concurrency, mapping and streams go with the driver, and the layout's tests drive their scans on a run of their own. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-layout/src/scan/execution/chunk.rs | 8 +- vortex-layout/src/scan/execution/filter.rs | 6 +- vortex-layout/src/scan/execution/mod.rs | 91 ++--- vortex-layout/src/scan/execution/prune.rs | 23 +- vortex-layout/src/scan/v2/io.rs | 6 +- vortex-layout/src/scan/v2/mod.rs | 398 +++------------------ vortex-layout/src/scan/v2/tests.rs | 230 +++++------- 7 files changed, 210 insertions(+), 552 deletions(-) diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 6808c707c48..7ef7d24bcba 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -59,9 +59,9 @@ pub(crate) fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec, rows: Range, chunks: IntoIter>, @@ -71,7 +71,9 @@ pub(in crate::scan) struct ChunkPlanner { } impl ChunkPlanner { - pub(super) fn new( + /// The planner of the chunks of `rows`, planning only what `selection` and `kept`, the rows + /// zone pruning left over `rows`, leave something of. + pub fn new( scan: Arc, rows: Range, selection: Selection, diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 92ef0e66c1e..1873e5af62a 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -45,7 +45,7 @@ use crate::scan::filter::FilterExpr; use crate::segments::SegmentId; /// Plans and adaptive conjunct ordering shared by every chunk of a scan. -pub(in crate::scan) struct FilterPlans { +pub struct FilterPlans { expression: FilterExpr, predicates: Vec, /// A metadata-only proof per conjunct, where the conjunct has one. @@ -53,7 +53,7 @@ pub(in crate::scan) struct FilterPlans { } impl FilterPlans { - pub(in crate::scan) fn new( + pub(crate) fn new( expression: BoundExpression, root: &PlanRef, session: &VortexSession, @@ -94,7 +94,7 @@ impl FilterPlans { } /// The pruning proofs, run once over a scan's rows before its chunks are planned. - pub(in crate::scan) fn proofs(&self) -> impl Iterator { + pub(crate) fn proofs(&self) -> impl Iterator { self.pruning.iter().flatten() } } diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index e371fdb539c..33200450e26 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -3,13 +3,13 @@ //! The planners of a V2 scan, and the prepared plans they run. //! -//! A [`PreparedScan`] holds a scan's plans over one file and makes its planners. The planners -//! form a chain, each emitting the next through the planning protocol: a [`PrunePlanner`] -//! proves the filter's zone maps over the scan's rows and emits one [`ChunkPlanner`] for the rows -//! they keep; the chunk planner cuts those rows into chunks and emits a [`FilterPlanner`] or a -//! [`ProjectionPlanner`] per chunk; a filter planner runs its chunk's conjuncts and emits a -//! projection planner for the rows that survive; a projection planner emits one morsel per piece -//! of its rows. A stage with nothing to do is never made. +//! A [`PreparedScan`] holds a scan's plans over one file. An engine composes the planners over it +//! and drives them: a [`PrunePlanner`] proves the filter's zone maps over the scan's rows and +//! emits the planner it was built to continue into, with the rows they keep; a [`ChunkPlanner`] +//! cuts those rows into chunks and spawns a [`FilterPlanner`] or a [`ProjectionPlanner`] per +//! chunk; a filter planner runs its chunk's conjuncts and emits a projection planner for the rows +//! that survive; a projection planner emits one morsel per piece of its rows. A planner only +//! says what it needs, compute or reads; where either runs is the engine's driver's decision. mod adaptor; mod chunk; @@ -22,6 +22,7 @@ use std::ops::Range; use std::sync::Arc; use vortex_buffer::Alignment; +use vortex_io::request::IoService; use vortex_io::request::IoTarget; use vortex_mask::Mask; use vortex_scan::planning::planner::Planner; @@ -29,7 +30,7 @@ use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_session::VortexSession; -pub(in crate::scan) use self::chunk::ChunkPlanner; +pub use self::chunk::ChunkPlanner; pub(crate) use self::chunk::FILTERED_ROWS_PER_CHUNK; pub(crate) use self::chunk::MAX_ROWS_PER_CHUNK; pub(crate) use self::chunk::ROW_BLOCK_SIZE; @@ -37,11 +38,16 @@ pub(crate) use self::chunk::ROW_BLOCK_SIZE; pub(crate) use self::chunk::row_chunks; pub(crate) use self::chunk::rows_per_chunk; use self::filter::FilterPlanner; -pub(super) use self::filter::FilterPlans; +pub use self::filter::FilterPlans; use self::projection::ProjectionPlanner; -use self::prune::PrunePlanner; +pub use self::prune::PruneNext; +pub use self::prune::PrunePlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; +use crate::scan::v2::ScanCache; +use crate::scan::v2::io::SegmentRanges; +use crate::scan::v2::io::SegmentScanIo; +use crate::segments::SegmentSource; /// Where a segment's bytes live in the file. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -67,9 +73,20 @@ impl SegmentLocation { /// A scan's plans over one file, built once from the file's footer and read by every planner of /// the scan. Nothing here touches a segment, and no planner plans again. -pub(crate) struct PreparedScan { +pub struct PreparedScan { /// The session used for decoding and expression evaluation. pub(super) session: VortexSession, + /// The rows of the root row domain the scan covers. + pub(super) rows: Range, + /// The rows of `rows` the scan selects. + pub(super) selection: Selection, + /// Fetches segment bytes by id. + pub(super) segments: Arc, + /// The segments' byte ranges, back to their ids. + pub(super) ranges: SegmentRanges, + /// Held for the scan's life, so scans of the file that start meanwhile share it. + #[expect(dead_code, reason = "kept alive, never read")] + pub(super) cache: ScanCache, /// The location of every segment, indexed by segment id. pub(super) locations: Arc<[SegmentLocation]>, /// The optimized projection plan, over the scan's root row domain. @@ -86,40 +103,28 @@ pub(crate) struct PreparedScan { } impl PreparedScan { - /// The first planner a scan of `rows` needs. Every later planner is emitted by the one before - /// it: the zone proofs run first when the filter has any, then one planner per chunk of the - /// rows they keep. - pub(crate) fn planner( - self: &Arc, - rows: Range, - selection: Selection, - ) -> Box { - let proofs = self - .filter - .as_ref() - .is_some_and(|filter| filter.proofs().next().is_some()); - if proofs && !rows.is_empty() { - let scan = Arc::clone(self); - let chunks = rows.clone(); - Box::new(PrunePlanner::new( - Arc::clone(self), - rows, - Box::new(move |kept| Box::new(scan.chunk_planner(chunks, selection, Some(kept)))), - )) - } else { - Box::new(self.chunk_planner(rows, selection, None)) - } + /// The rows of the file's row domain the scan covers. + pub fn rows(&self) -> &Range { + &self.rows + } + + /// The rows of [`rows`](Self::rows) the scan selects. + pub fn selection(&self) -> &Selection { + &self.selection + } + + /// The filter's plans, when the scan has a filter. + pub fn filter(&self) -> Option<&FilterPlans> { + self.filter.as_deref() } - /// The planner cutting `rows` into chunks, planning only the chunks `selection` and `kept` - /// leave something of. - pub(super) fn chunk_planner( - self: &Arc, - rows: Range, - selection: Selection, - kept: Option, - ) -> ChunkPlanner { - ChunkPlanner::new(Arc::clone(self), rows, selection, kept) + /// The service a driver takes IO sessions from: every read a planner of the scan requests is + /// served from it by segment. + pub fn io(&self) -> Arc { + Arc::new(SegmentScanIo::new( + Arc::clone(&self.segments), + Arc::clone(&self.ranges), + )) } /// The planner of one chunk's selected rows: its conjuncts when the scan has a filter, else diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 1fd153f989c..9389d9b2d72 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -48,17 +48,21 @@ fn is_proved(plan: &PlanRef, session: &VortexSession) -> VortexResult { Ok(true) } -/// The planner a prune planner continues into, given the rows the proofs kept. -pub(super) type Next = Box Box + Send>; +/// The planner a prune planner continues into, given the rows the proofs kept, or `None` when +/// no proof ran. +pub type PruneNext = Box) -> Box + Send>; /// Runs the filter's zone-map proofs over the scan's rows and emits the planner it was built /// to continue into, with the rows no proof excludes. The zone maps are the only reads, -/// started together. Rows every proof excludes end the scan here: nothing is emitted. -pub(super) struct PrunePlanner { +/// started together. Rows every proof excludes end the scan here: nothing is emitted. A filter +/// without proofs emits the next planner at once. +pub struct PrunePlanner { scan: Arc, rows: Range, /// What to emit once every proof has run, taken when it is. - next: Option, + next: Option, + /// Whether any proof ran, so the kept rows mean something. + proved: bool, /// The rows no proof has excluded yet, over `rows`. mask: Mask, proofs: IntoIter, @@ -71,7 +75,8 @@ pub(super) struct PrunePlanner { } impl PrunePlanner { - pub(super) fn new(scan: Arc, rows: Range, next: Next) -> Self { + /// The planner of the proofs of `scan`'s filter over `rows`, continuing into `next`. + pub fn new(scan: Arc, rows: Range, next: PruneNext) -> Self { let proofs: Vec<_> = scan .filter .iter() @@ -82,6 +87,7 @@ impl PrunePlanner { scan, rows, next: Some(next), + proved: false, mask: Mask::new_true(len), proofs: proofs.into_iter(), active: None, @@ -170,6 +176,7 @@ impl Planner for PrunePlanner { ); self.active = None; self.pieces.clear(); + self.proved = true; } if self.mask.all_false() { self.done = true; @@ -201,7 +208,9 @@ impl Planner for PrunePlanner { .next .take() .ok_or_else(|| vortex_err!("the prune planner has nothing to continue into"))?; - let kept = std::mem::replace(&mut self.mask, Mask::new_true(0)); + let kept = self + .proved + .then(|| std::mem::replace(&mut self.mask, Mask::new_true(0))); Ok(PlannerOutput::Planner( WorkScope { file_ordinal: 0, diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 2779739b281..0e690ebfde7 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -35,7 +35,7 @@ use crate::segments::SegmentId; use crate::segments::SegmentSource; /// Byte ranges of a file's segments, keyed by `(offset, length)`, back to their segment ids. -pub(super) type SegmentRanges = Arc>; +pub(crate) type SegmentRanges = Arc>; pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { let ranges = locations.iter().enumerate().map(|(id, location)| { @@ -46,13 +46,13 @@ pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { } /// Serves splits from any [`SegmentSource`], by segment id. -pub(super) struct SegmentScanIo { +pub(crate) struct SegmentScanIo { segments: Arc, ranges: SegmentRanges, } impl SegmentScanIo { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + pub(crate) fn new(segments: Arc, ranges: SegmentRanges) -> Self { Self { segments, ranges } } } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 5f8384fc36b..77c7c1fb29b 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -3,62 +3,39 @@ //! The planner-driven scan executor. //! -//! A V2 scan lowers the file's layout once into a physical plan and runs every split of it -//! through the planning protocol: the [`Run`] driver advances -//! the split's planners and morsels, and every segment read is a protocol request served from the -//! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. +//! A V2 scan lowers the file's layout once into a physical plan and runs it through the planning +//! protocol: an engine composes planners over a [`PreparedScan`], a [`PrunePlanner`] into a +//! [`ChunkPlanner`], admits the first to a [`Run`](vortex_scan::planning::driver::Run) of its own +//! driver, and every segment read is a protocol request served from the file's +//! [`SegmentSource`]. The layout reader is never asked to evaluate anything, and this crate never +//! decides where compute runs or how IO is awaited: that is the engine's driver. //! -//! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] -//! directly and supports projection, filters, row ranges, and selections. Unsupported layouts return an -//! error; limits are not part of this builder's API. Work is partitioned into fixed +//! Import this module's [`ScanBuilder`] to prepare a V2 scan. It accepts a [`ScanFile`] directly +//! and supports projection, filters, row ranges, and selections. Unsupported layouts return an +//! error when the plan is prepared; limits are not part of this builder's API. Work is cut into //! chunks of up to 65,536 rows, or 1,048,576 rows for filtered scans, independently of the stored -//! layout's chunk boundaries. -//! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. +//! layout's chunk boundaries. Full-file projections use smaller 8,192-row multiples when +//! serialized bytes per row are high. //! //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and proved //! zones, segment index, and dictionary values with every other scan of that file, like a cached //! layout reader does for V1 scans. mod cache; -mod io; +pub(crate) mod io; #[cfg(test)] mod tests; -use std::future; use std::ops::Range; use std::sync::Arc; -use std::task::Poll; -use arrow_array::RecordBatch; -use arrow_schema::ArrowError; -use arrow_schema::Field; -use arrow_schema::SchemaRef; -use futures::SinkExt; -use futures::Stream; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::channel::mpsc; -use futures::stream; -use futures::stream::BoxStream; -use vortex_array::ArrayRef; -use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_io::request::IoService; -use vortex_io::request::IoSource; -use vortex_io::runtime::Handle; -use vortex_io::session::RuntimeSessionExt; -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_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; pub use self::cache::ScanCache; pub use self::cache::WeakScanCache; @@ -66,15 +43,16 @@ use crate::LayoutRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; -use crate::scan::arrow::to_record_batch; +pub use crate::scan::execution::ChunkPlanner; use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; -use crate::scan::execution::FilterPlans; +pub use crate::scan::execution::FilterPlans; use crate::scan::execution::MAX_ROWS_PER_CHUNK; -use crate::scan::execution::PreparedScan; +pub use crate::scan::execution::PreparedScan; +pub use crate::scan::execution::PruneNext; +pub use crate::scan::execution::PrunePlanner; use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; use crate::scan::execution::rows_per_chunk; -use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::segments::SegmentSource; @@ -107,24 +85,21 @@ pub struct ScanFile { pub cache: Option, } -/// Builds a projection scan on the planner-driven V2 executor. +/// Prepares a scan on the planner-driven V2 executor. /// -/// This builder reads the file's layout and segments directly. Unsupported layouts return an -/// error when the stream is polled. Row limits are not supported. -pub struct ScanBuilder { +/// This builder reads the file's layout and segments directly and plans the scan without IO. +/// Driving the plan is the engine's: see [`PreparedScan`]. Row limits are not supported. +pub struct ScanBuilder { session: VortexSession, file: ScanFile, projection: BoundExpression, filter: Option, - ordered: bool, row_range: Option>, selection: Selection, - concurrency: usize, - map_fn: Arc VortexResult + Send + Sync>, row_offset: u64, } -impl ScanBuilder { +impl ScanBuilder { /// Creates a V2 scan over the file's layout and segment source. pub fn new(session: VortexSession, file: ScanFile) -> Self { Self { @@ -132,34 +107,13 @@ impl ScanBuilder { filter: None, session, file, - ordered: true, row_range: None, selection: Selection::default(), - concurrency: 4, - map_fn: Arc::new(Ok), row_offset: 0, } } - /// Converts each output batch to Arrow on the split's task. - pub fn into_record_batch_stream( - self, - schema: SchemaRef, - ) -> VortexResult> + Send + 'static> { - let struct_field = Field::new_struct("", schema.fields().clone(), false); - let session = self.session.clone(); - Ok(self - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - to_record_batch(chunk, &struct_field, &mut ctx) - }) - .into_stream()? - .map_err(|error| ArrowError::ExternalError(Box::new(error)))) - } -} - -impl ScanBuilder { - /// The output dtype after projection, before any output mapping. + /// The output dtype after projection. pub fn dtype(&self) -> &DType { self.projection.dtype() } @@ -176,12 +130,6 @@ impl ScanBuilder { self } - /// Configures whether batches are yielded in file order. - pub fn with_ordered(mut self, ordered: bool) -> Self { - self.ordered = ordered; - self - } - /// Restricts scanning to a contiguous file-local row range. pub fn with_row_range(mut self, row_range: Range) -> Self { self.row_range = Some(row_range); @@ -205,82 +153,25 @@ impl ScanBuilder { self } - /// Sets the number of chunks run at once per worker thread. Must be greater than zero. - pub fn with_concurrency(mut self, concurrency: usize) -> Self { - assert!(concurrency > 0); - self.concurrency = concurrency; - self - } - - /// Maps each output batch on its split's task. - pub fn map( - self, - map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, - ) -> ScanBuilder { - let old_map_fn = self.map_fn; - ScanBuilder { - session: self.session, - file: self.file, - projection: self.projection, - filter: self.filter, - ordered: self.ordered, - row_range: self.row_range, - selection: self.selection, - concurrency: self.concurrency, - row_offset: self.row_offset, - map_fn: Arc::new(move |array| old_map_fn(array).and_then(&map_fn)), - } - } - - /// Runs the scan on the session's runtime, preparing it on the first poll. - /// Preparation errors are returned as stream items. - pub fn into_stream(self) -> VortexResult>> { - let handle = self.session.handle(); - Ok(stream::once(async move { - handle - .spawn_cpu(move || ScanRun::try_new(self)?.execute_stream()) - .await - }) - .try_flatten() - .boxed()) - } -} - -/// A scan prepared over its file, ready to run. -struct ScanRun { - session: VortexSession, - scan: Arc, - /// Serves the scan's reads. - io: Arc, - /// The rows of the root row domain the scan covers. - rows: Range, - selection: Selection, - map_fn: Arc VortexResult + Send + Sync>, - ordered: bool, - concurrency: usize, - /// Held until the scan ends, so scans of the file that start meanwhile share it. - cache: ScanCache, -} - -impl ScanRun { /// Lowers the file's layout and plans the filter and the projection, from the file's cache - /// where an earlier scan planned them already. Touches no segment. - fn try_new(builder: ScanBuilder) -> VortexResult { - let file = builder.file; + /// where an earlier scan planned them already. Touches no segment. Lowering and optimizing a + /// wide file is real CPU, so a driver runs this where it runs compute. + pub fn prepare(self) -> VortexResult { + let file = self.file; let cache = file.cache.unwrap_or_default(); let root = match cache.root() { Some(root) => root, None => cache.insert_root(lower(&file.layout)?), }; let row_count = root.row_count(); - let filter = builder + let filter = self .filter .map(|filter| -> VortexResult<_> { match cache.filter_plans(&filter) { Some(plans) => Ok(plans), None => { let plans = - Arc::new(FilterPlans::new(filter.clone(), &root, &builder.session)?); + Arc::new(FilterPlans::new(filter.clone(), &root, &self.session)?); Ok(cache.insert_filter_plans(filter, plans)) } } @@ -288,15 +179,15 @@ impl ScanRun { .transpose()?; // Optimizing drops the zone tables a projection never reads and keeps the chunks of a // chunked column lazy, so an identity projection is optimized like any other. - let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype(); - let projection = match cache.projection(&builder.projection) { + let identity = self.projection.is_root() && self.projection.dtype() == root.dtype(); + let projection = match cache.projection(&self.projection) { Some(plan) => plan, None => { let plan = optimize( - plan_row_idx_expression(builder.projection.clone(), root)?, - &builder.session, + plan_row_idx_expression(self.projection.clone(), root)?, + &self.session, )?; - cache.insert_projection(builder.projection, plan) + cache.insert_projection(self.projection, plan) } }; let chunk_rows = if filter.is_some() { @@ -310,215 +201,20 @@ impl ScanRun { Some(ranges) => ranges, None => cache.insert_ranges(segment_ranges(&file.locations)), }; - let io = Arc::new(SegmentScanIo::new(file.segments, ranges)); - Ok(Self { - session: builder.session.clone(), - scan: Arc::new(PreparedScan { - session: builder.session, - locations: file.locations, - projection, - filter, - row_offset: builder.row_offset, - values: cache.values(), - chunk_rows, - }), - io, - rows: builder.row_range.unwrap_or(0..row_count), - selection: builder.selection, - map_fn: builder.map_fn, - ordered: builder.ordered, - concurrency: builder.concurrency, + let values = cache.values(); + Ok(PreparedScan { + session: self.session, + rows: self.row_range.unwrap_or(0..row_count), + selection: self.selection, + segments: file.segments, + ranges, cache, + locations: file.locations, + projection, + filter, + row_offset: self.row_offset, + values, + chunk_rows, }) } - - /// Runs the scan on planning runs the returned stream owns. - /// - /// The scan's first planner runs on one run the stream drives as it is polled, each of its - /// CPU steps a [`Handle::spawn_cpu`] task. Every chunk it spawns runs on a run and a session - /// of its own, on one task of the runtime that drives the run to its end and awaits only its - /// reads, up to `concurrency` chunks per worker thread at once; a chunk is started only when - /// the consumer has room for its batches. Every read is a `spawn_io` task of the segment - /// source. The engine thread polling the stream does no scan work itself. Batches of an - /// unordered scan are yielded as morsels finish; an ordered scan yields them in row order - /// once every morsel has. - fn execute_stream(self) -> VortexResult>> { - let handle = self.session.handle(); - let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); - let root = self.scan.planner(self.rows.clone(), self.selection); - let scope = WorkScope { - file_ordinal: 0, - rows: self.rows.clone(), - }; - let io = self.io; - let map_fn = self.map_fn; - let (chunk_handle, chunk_io, chunk_map) = - (handle.clone(), Arc::clone(&io), Arc::clone(&map_fn)); - let inline = in_flight <= 1; - let batches = run_stream(handle, root, scope, io.session(), map_fn, inline) - .map_ok(move |yielded| match yielded { - Yielded::Batch(start, batch) => { - stream::once(future::ready(Ok((start, batch)))).boxed() - } - Yielded::Spawn(scope, chunk) => chunk_stream( - &chunk_handle, - chunk, - scope, - chunk_io.session(), - Arc::clone(&chunk_map), - ), - }) - .try_flatten_unordered(in_flight); - let cache = self.cache; - let batches = batches.chain(stream::poll_fn(move |_| { - let _ = &cache; - Poll::Ready(None) - })); - if !self.ordered { - return Ok(batches.map_ok(|(_, batch)| batch).boxed()); - } - Ok(stream::once(async move { - let mut held: Vec<(u64, A)> = batches.try_collect().await?; - // Morsels finish in whatever order their reads arrive. - held.sort_by_key(|(start, _)| *start); - VortexResult::Ok(stream::iter(held.into_iter().map(|(_, batch)| Ok(batch)))) - }) - .try_flatten() - .boxed()) - } -} - -/// What the root run's stream yields: a mapped batch with the row it starts at, or a chunk a -/// planner of the run spawned. -enum Yielded { - Batch(u64, A), - Spawn(WorkScope, Box), -} - -/// One CPU step of a run, taken on a `spawn_cpu` task with the run. -enum Step { - Batch(u64, A), - Spawn(WorkScope, Box), - RootDone, - Idle, - Waiting, -} - -/// The state of the root run's stream. -struct RunState { - handle: Handle, - /// Taken while a CPU step runs on its task and put back with its result. - run: Option, - io: Arc, - map_fn: Arc VortexResult + Send + Sync>, - /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. - inline: bool, -} - -/// Runs `root` over `scope` on a planning run of its own, reading through `io`, each CPU step -/// on a `spawn_cpu` task, and yields its mapped batches and, unless `inline`, the chunks it -/// spawns. Inline, a spawned chunk runs on this run beside the others: the run's steps already -/// interleave every chunk's work, and a run per chunk would only add a hand-off per chunk for -/// parallelism the engine provides. -fn run_stream( - handle: Handle, - root: Box, - scope: WorkScope, - io: Arc, - map_fn: Arc VortexResult + Send + Sync>, - inline: bool, -) -> BoxStream<'static, VortexResult>> { - let mut run = Run::new(); - run.admit(root, scope, Arc::clone(&io)); - let state = RunState { - handle, - run: Some(run), - io, - map_fn, - inline, - }; - stream::try_unfold(state, |mut state| async move { - loop { - let run = state - .run - .take() - .ok_or_else(|| vortex_err!("the scan's run is missing"))?; - let map_fn = Arc::clone(&state.map_fn); - let (run, step) = state - .handle - .spawn_cpu(move || { - let mut run = run; - let step = run.advance().and_then(|progress| { - Ok(match progress { - Progress::Batch(batch) => { - Step::Batch(batch.scope.rows.start, map_fn(batch.array)?) - } - Progress::Spawn(scope, child) => Step::Spawn(scope, child), - Progress::RootDone(_) => Step::RootDone, - Progress::Idle => Step::Idle, - Progress::Waiting => Step::Waiting, - }) - }); - (run, step) - }) - .await; - let run = state.run.insert(run); - match step? { - Step::Batch(start, batch) => { - return Ok(Some((Yielded::Batch(start, batch), state))); - } - Step::Spawn(scope, child) if state.inline => { - run.admit(child, scope, Arc::clone(&state.io)); - } - Step::Spawn(scope, child) => { - return Ok(Some((Yielded::Spawn(scope, child), state))); - } - Step::RootDone if run.live_roots() > 0 => {} - Step::RootDone | Step::Idle => return Ok(None), - Step::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - }) - .boxed() -} - -/// Runs a spawned chunk's `root` over `scope` on a planning run of its own, on one task that -/// drives the run to its end and awaits only its reads, and yields its mapped batches with the -/// row each starts at. A planner the chunk spawns in turn runs as another root of the same run. -/// -/// One task per chunk, not one per step: a chunk's run waits for reads many times, and a task -/// hop at each wait is latency on the chunk's critical path that no other chunk's work fills. -fn chunk_stream( - handle: &Handle, - root: Box, - scope: WorkScope, - io: Arc, - map_fn: Arc VortexResult + Send + Sync>, -) -> BoxStream<'static, VortexResult<(u64, A)>> { - 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)); - loop { - match run.advance()? { - Progress::Batch(batch) => { - let start = batch.scope.rows.start; - let batch = map_fn(batch.array); - let failed = batch.is_err(); - if send.send(batch.map(|batch| (start, batch))).await.is_err() || failed { - return Ok(()); - } - } - Progress::Spawn(scope, child) => { - run.admit(child, scope, Arc::clone(&io)); - } - Progress::RootDone(_) if run.live_roots() > 0 => {} - Progress::RootDone(_) | Progress::Idle => return Ok(()), - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - }); - let errors = stream::once(task) - .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); - stream::select(recv, errors).boxed() } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index b69f5390704..822b0a2cafc 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -7,14 +7,11 @@ use std::sync::Arc; use std::sync::LazyLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use std::time::Duration; use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; -use futures::FutureExt; +use arrow_schema::Field; use futures::TryStreamExt; -use futures::channel::oneshot; -use futures::future::Shared; use futures::stream; use rstest::rstest; use vortex_array::ArrayContext; @@ -44,8 +41,8 @@ use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; use vortex_scan::planning::driver::Batch; use vortex_scan::planning::driver::Progress; use vortex_scan::planning::driver::Run; @@ -55,12 +52,14 @@ use vortex_scan::planning::planner::WorkScope; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; +use super::ChunkPlanner; use super::MAX_ROWS_PER_CHUNK; +use super::PreparedScan; +use super::PrunePlanner; use super::ROW_BLOCK_SIZE; use super::ScanBuilder; use super::ScanCache; use super::ScanFile; -use super::ScanRun; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; @@ -71,6 +70,7 @@ use crate::layouts::row_idx::row_idx; use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; +use crate::scan::arrow::to_record_batch; use crate::scan::execution::SegmentLocation; use crate::scan::execution::row_chunks; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; @@ -190,34 +190,50 @@ fn builder( session: &VortexSession, segments: &Arc, layout: &LayoutRef, -) -> VortexResult> { +) -> VortexResult { Ok(ScanBuilder::new( session.clone(), scan_file(segments, layout)?, )) } -/// Runs `rows` of a prepared scan on a planning run of its own and returns its batches, in row -/// order, each with the rows it stands for. -async fn run_batches( - scan: &ScanRun, - rows: Range, -) -> VortexResult> { +/// Composes a scan's planners the way an engine does: the proofs first when there is a filter, +/// continuing into the chunks of the rows they keep. +fn compose(scan: &Arc, rows: Range) -> Box { + let chunks = { + let scan = Arc::clone(scan); + let rows = rows.clone(); + move |kept: Option| -> Box { + let selection = scan.selection().clone(); + Box::new(ChunkPlanner::new(scan, rows, selection, kept)) + } + }; + if scan.filter().is_some() { + Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(chunks))) + } else { + chunks(None) + } +} + +/// Runs `rows` of a prepared scan on one planning run, every spawned chunk admitted to it, and +/// returns its batches in row order, each with the rows it stands for. +async fn run_batches(scan: &Arc, rows: Range) -> VortexResult> { + let io = scan.io(); let mut run = Run::new(); run.admit( - scan.scan.planner(rows.clone(), scan.selection.clone()), + compose(scan, rows.clone()), WorkScope { file_ordinal: 0, rows, }, - scan.io.session(), + io.session(), ); let mut batches = Vec::new(); loop { match run.advance()? { Progress::Batch(batch) => batches.push(batch), Progress::Spawn(scope, chunk) => { - run.admit(chunk, scope, scan.io.session()); + run.admit(chunk, scope, io.session()); } Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => break, @@ -228,11 +244,8 @@ async fn run_batches( Ok(batches) } -/// Runs `rows` of a prepared scan on a planning run of its own and returns its arrays. -async fn run_rows( - scan: &ScanRun, - rows: Range, -) -> VortexResult> { +/// Runs `rows` of a prepared scan and returns its arrays. +async fn run_rows(scan: &Arc, rows: Range) -> VortexResult> { Ok(run_batches(scan, rows) .await? .into_iter() @@ -240,6 +253,13 @@ async fn run_rows( .collect()) } +/// Prepares and runs a whole scan, returning its arrays. +async fn run_scan(builder: ScanBuilder) -> VortexResult> { + let scan = Arc::new(builder.prepare()?); + let rows = scan.rows().clone(); + run_rows(&scan, rows).await +} + async fn collect( dtype: DType, stream: impl futures::Stream>, @@ -248,10 +268,15 @@ async fn collect( Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) } +/// Prepares and runs a whole scan, returning its rows as one chunked array. +async fn collect_scan(dtype: DType, builder: ScanBuilder) -> VortexResult { + Ok(ChunkedArray::try_new(run_scan(builder).await?, dtype)?.into_array()) +} + /// Runs `builder` on both executors and checks the planner-driven one returns the same rows. async fn assert_same_as_v1( session: &VortexSession, - make: impl Fn() -> VortexResult>, + make: impl Fn() -> VortexResult, ) -> VortexResult<()> { let builder = make()?; let dtype = builder.dtype().clone(); @@ -265,13 +290,12 @@ async fn assert_same_as_v1( .with_projection(builder.projection.clone()) .with_some_filter(builder.filter.clone()) .with_selection(builder.selection.clone()) - .with_row_offset(builder.row_offset) - .with_ordered(builder.ordered); + .with_row_offset(builder.row_offset); if let Some(rows) = builder.row_range.clone() { v1 = v1.with_row_range(rows); } let expected = collect(dtype.clone(), v1.into_stream()?).await?; - let actual = collect(dtype, builder.into_stream()?).await?; + let actual = collect_scan(dtype, builder).await?; assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); Ok(()) } @@ -422,8 +446,14 @@ async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { let scan = builder(&session, &segments, &layout)? .with_projection(projection) .with_row_range(990..1010); - let schema = Arc::new(session.arrow().to_arrow_schema(scan.dtype())?); - let batches: Vec<_> = scan.into_record_batch_stream(schema)?.try_collect().await?; + let schema = session.arrow().to_arrow_schema(scan.dtype())?; + let field = Field::new_struct("", schema.fields().clone(), false); + let mut ctx = session.create_execution_ctx(); + let batches = run_scan(scan) + .await? + .into_iter() + .map(|chunk| to_record_batch(chunk, &field, &mut ctx)) + .collect::>>()?; let values: Vec<_> = batches .iter() .flat_map(|batch| { @@ -520,8 +550,8 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; - let chunks = row_chunks(scan.rows.clone(), scan.scan.chunk_rows); + let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); + let chunks = row_chunks(scan.rows().clone(), scan.chunk_rows); assert_eq!( chunks.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) @@ -559,10 +589,7 @@ async fn list_layouts_return_an_error() -> VortexResult<()> { let session = new_session().with_tokio(); let strategy = TableStrategy::new(chunked_flat(), chunked_flat()).with_list_layout(); let (segments, layout) = write(&session, &strategy).await?; - let builder = builder(&session, &segments, &layout)?; - let error = builder - .into_stream()? - .try_next() + let error = run_scan(builder(&session, &segments, &layout)?) .await .expect_err("V2 must reject an unsupported layout"); assert!(error.to_string().contains("vortex.plan.list_pack")); @@ -582,112 +609,25 @@ async fn projection_skips_unsupported_unselected_layouts() -> VortexResult<()> { } #[tokio::test] -async fn unordered_returns_every_row() -> VortexResult<()> { +async fn every_row_is_returned() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = ScanRun::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; - let chunks: Vec = scan.execute_stream()?.try_collect().await?; + let chunks = run_scan(builder(&session, &segments, &layout)?).await?; let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); assert_eq!(rows, CHUNKS * CHUNK_ROWS); Ok(()) } -struct GatedSegments { - inner: Arc, - blocked: SegmentId, - release: Shared>, -} - -impl SegmentSource for GatedSegments { - fn request(&self, id: SegmentId) -> SegmentFuture { - let bytes = self.inner.request(id); - if id != self.blocked { - return bytes; - } - let release = self.release.clone(); - async move { - let _ = release.await; - bytes.await - } - .boxed() - } -} - -#[rstest] -#[timeout(Duration::from_secs(5))] -#[tokio::test] -async fn unordered_yields_before_the_split_finishes() -> VortexResult<()> { - let session = new_session().with_tokio(); - let segments = Arc::new(TestSegments::default()); - let (mut sequence, eof) = SequenceId::root().split(); - let arrays = (0..2) - .map(|index| { - let start = i32::try_from(index * MAX_ROWS_PER_CHUNK)?; - Ok(( - sequence.advance(), - PrimitiveArray::from_iter(start..start + i32::try_from(MAX_ROWS_PER_CHUNK)?) - .into_array(), - )) - }) - .collect::>>()?; - let dtype = arrays[0].1.dtype().clone(); - let layout = chunked_flat() - .write_stream( - ArrayContext::empty().into(), - Arc::::clone(&segments), - SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))) - .sendable(), - eof, - &session, - ) - .await?; - let blocked = *layout - .slot(1)? - .ok_or_else(|| vortex_err!("second chunk is missing"))? - .segment_ids() - .first() - .ok_or_else(|| vortex_err!("second chunk has no segment"))?; - let (release, gate) = oneshot::channel(); - let segments: Arc = Arc::new(GatedSegments { - inner: segments, - blocked, - release: gate.shared(), - }); - let mut scan = ScanRun::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; - // Put both projection morsels in one chunk, as a dense filtered scan would. - Arc::get_mut(&mut scan.scan) - .ok_or_else(|| vortex_err!("the prepared scan is not shared yet"))? - .chunk_rows = usize::try_from(layout.row_count())?; - let mut batches = scan.execute_stream()?; - let first = batches.try_next().await; - let _ = release.send(()); - let first = first?.ok_or_else(|| vortex_err!("first batch is missing"))?; - assert_arrays_eq!( - first, - PrimitiveArray::from_iter(0i32..i32::try_from(MAX_ROWS_PER_CHUNK)?), - &mut session.create_execution_ctx() - ); - let rest: Vec<_> = batches.try_collect().await?; - assert_eq!( - rest.iter().map(|array| array.len()).sum::(), - MAX_ROWS_PER_CHUNK - ); - Ok(()) -} - #[tokio::test] async fn primitive_column() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; let dtype = layout.dtype().clone(); - let scan = ScanRun::try_new( + let actual = collect_scan( + DType::Primitive(vortex_array::dtype::PType::I32, false.into()), builder(&session, &segments, &layout)? .with_projection(get_item("ints", root()).bind(&dtype)?) .with_row_range(990..1010), - )?; - let actual = collect( - DType::Primitive(vortex_array::dtype::PType::I32, false.into()), - scan.execute_stream()?, ) .await?; assert_arrays_eq!( @@ -706,12 +646,12 @@ async fn work_chunks_ignore_layout_boundaries( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, strategy.get().as_ref()).await?; - let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); assert_eq!( - row_chunks(scan.rows.clone(), scan.scan.chunk_rows), + row_chunks(scan.rows().clone(), scan.chunk_rows), vec![0..(CHUNKS * CHUNK_ROWS) as u64] ); - let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; + let arrays = run_rows(&scan, scan.rows().clone()).await?; assert_eq!(arrays.len(), 1); assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) @@ -771,7 +711,7 @@ async fn scan_cache_reuses_values_between_scans( let before = counted.reads.load(Ordering::Relaxed); let scan = ScanBuilder::new(session.clone(), file.clone()) .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![12u64, 15]))?); - let actual = collect(layout.dtype().clone(), scan.into_stream()?).await?; + let actual = collect_scan(layout.dtype().clone(), scan).await?; assert_arrays_eq!( actual, PrimitiveArray::from_iter([12i32, 15]), @@ -788,13 +728,17 @@ async fn scan_cache_reuses_values_between_scans( async fn selected_split_scope_is_bounded_before_projection() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = ScanRun::try_new( + let scan = Arc::new( builder(&session, &segments, &layout)? - .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 2001]))?), - )?; - let mut chunks = scan - .scan - .chunk_planner(scan.rows.clone(), scan.selection.clone(), None); + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 2001]))?) + .prepare()?, + ); + let mut chunks = ChunkPlanner::new( + Arc::clone(&scan), + scan.rows().clone(), + scan.selection().clone(), + None, + ); let PlannerOutput::Spawn(scope, _) = chunks.compute()? else { vortex_bail!("the selected rows were not planned"); }; @@ -827,7 +771,7 @@ async fn repeated_projections_reuse_their_values() -> VortexResult<()> { let scan = ScanBuilder::new(session.clone(), file.clone()) .with_projection(projection.clone()) .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![12u64, 18]))?); - let actual = collect(projection.dtype().clone(), scan.into_stream()?).await?; + let actual = collect_scan(projection.dtype().clone(), scan).await?; assert_arrays_eq!( actual, vortex_array::arrays::BoolArray::from_iter([false, true]), @@ -873,12 +817,14 @@ async fn pruned_rows_get_no_chunk( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, zoned_table().as_ref()).await?; - let scan = ScanRun::try_new( - builder(&session, &segments, &layout)?.with_some_filter(Some( - gt(get_item("ints", root()), lit(above)).bind(layout.dtype())?, - )), - )?; - let batches = run_batches(&scan, scan.rows.clone()).await?; + let scan = Arc::new( + builder(&session, &segments, &layout)? + .with_some_filter(Some( + gt(get_item("ints", root()), lit(above)).bind(layout.dtype())?, + )) + .prepare()?, + ); + let batches = run_batches(&scan, scan.rows().clone()).await?; assert_eq!( batches .iter() @@ -913,7 +859,7 @@ async fn shared_cache_proves_zones_once() -> VortexResult<()> { let scan = ScanBuilder::new(session.clone(), file.clone()) .with_projection(projection.clone()) .with_some_filter(Some(filter.clone())); - let actual = collect(projection.dtype().clone(), scan.into_stream()?).await?; + let actual = collect_scan(projection.dtype().clone(), scan).await?; assert_eq!(actual.len(), 1_000); reads.push(counted.reads.load(Ordering::Relaxed) - before); } From 92f3f4d146ec249e93ea6ad6459165c5c39d1bbc Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 15:06:08 +0100 Subject: [PATCH 02/22] Compose and drive a V2 scan in DataFusion The opener composes a scan's planners, a prune planner into a chunk planner when the scan has a filter and the chunk planner alone when it does not, and a driver of this crate runs them: the root run stepped on `spawn_cpu` tasks as DataFusion polls the stream, each spawned chunk on a task of its own up to the concurrency, or on the root run when the engine's partitions already fill the cores, with the concurrency derived from cores per partition unless configured. Batches are mapped to Arrow on the task that produced them. Every spawning decision of a V2 scan under DataFusion is in this crate. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-datafusion/src/persistent/mod.rs | 1 + vortex-datafusion/src/persistent/opener.rs | 10 +- vortex-datafusion/src/persistent/scan_v2.rs | 284 ++++++++++++++++++++ 3 files changed, 291 insertions(+), 4 deletions(-) create mode 100644 vortex-datafusion/src/persistent/scan_v2.rs diff --git a/vortex-datafusion/src/persistent/mod.rs b/vortex-datafusion/src/persistent/mod.rs index 53fefacafa4..1df70f35edc 100644 --- a/vortex-datafusion/src/persistent/mod.rs +++ b/vortex-datafusion/src/persistent/mod.rs @@ -28,6 +28,7 @@ mod format; pub mod metrics; mod opener; pub mod reader; +mod scan_v2; mod sink; mod source; mod stream; diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 26620ea2522..f85482a31ba 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -53,6 +53,7 @@ use vortex::error::VortexExpect; use vortex::expr::and_collect; use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; +use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; @@ -77,6 +78,7 @@ use crate::metrics::PARTITION_LABEL; use crate::metrics::PATH_LABEL; use crate::persistent::cache::CachedVortexMetadata; use crate::persistent::reader::VortexReaderFactory; +use crate::persistent::scan_v2::ScanDriver; use crate::persistent::stream::PrunableStream; #[derive(Clone)] @@ -437,8 +439,7 @@ impl FileOpener for VortexOpener { scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta.location)); let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file) .with_projection(scan_projection) - .with_some_filter(native_filter) - .with_ordered(has_output_ordering); + .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()); @@ -462,8 +463,9 @@ impl FileOpener for VortexOpener { .div_ceil(partitions) .max(1) }); - builder = builder.with_concurrency(concurrency); - builder.map(map_batch).into_stream() + let driver = + ScanDriver::new(vxf.session().handle(), concurrency, has_output_ordering); + Ok(driver.run(builder, map_batch)) } else { // Share V1 layout readers across partitions of the file. let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { diff --git a/vortex-datafusion/src/persistent/scan_v2.rs b/vortex-datafusion/src/persistent/scan_v2.rs new file mode 100644 index 00000000000..10b36559f56 --- /dev/null +++ b/vortex-datafusion/src/persistent/scan_v2.rs @@ -0,0 +1,284 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Drives a V2 scan for DataFusion. +//! +//! A planner says what it needs: compute, or reads. This driver decides where each runs. It +//! composes the planners over a [`PreparedScan`], the zone proofs first when the scan has a +//! filter and the chunks of the rows they keep after, and runs them on planning runs it owns: +//! the root run stepped on `spawn_cpu` tasks as the stream is polled, each spawned chunk on a +//! task of its own, or on the root run when the engine's partitions already fill the cores. +//! Every read is a `spawn_io` task of the file's segment source. The DataFusion thread polling +//! the stream does no scan work itself. + +use std::future; +use std::sync::Arc; +use std::task::Poll; + +use arrow_array::RecordBatch; +use futures::SinkExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::channel::mpsc; +use futures::stream; +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::runtime::Handle; +use vortex::layout::scan::v2::ChunkPlanner; +use vortex::layout::scan::v2::PreparedScan; +use vortex::layout::scan::v2::PrunePlanner; +use vortex::layout::scan::v2::ScanBuilder; +use vortex::mask::Mask; +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 VortexResult + 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. + concurrency: usize, + /// Whether batches are yielded in row order, which holds them until the scan is done. + ordered: bool, +} + +impl ScanDriver { + pub(crate) fn new(handle: Handle, concurrency: usize, ordered: bool) -> Self { + Self { + handle, + concurrency, + ordered, + } + } + + /// Prepares `builder` where compute runs, composes its planners, and streams their batches + /// mapped by `map`. Preparation errors are the stream's first item. + pub(crate) fn run( + self, + builder: ScanBuilder, + map: impl Fn(ArrayRef) -> VortexResult + Send + Sync + 'static, + ) -> BoxStream<'static, VortexResult> { + let handle = self.handle.clone(); + let map: MapBatch = Arc::new(map); + stream::once(async move { + handle + .spawn_cpu(move || { + let scan = Arc::new(builder.prepare()?); + VortexResult::Ok(self.execute(scan, map)) + }) + .await + }) + .try_flatten() + .boxed() + } + + fn execute( + self, + scan: Arc, + map: MapBatch, + ) -> BoxStream<'static, VortexResult> { + let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); + let inline = in_flight <= 1; + let scope = WorkScope { + file_ordinal: 0, + rows: scan.rows().clone(), + }; + let io = scan.io(); + let root = compose(&scan); + 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) + .map_ok(move |yielded| match yielded { + Yielded::Batch(start, batch) => { + stream::once(future::ready(Ok((start, batch)))).boxed() + } + Yielded::Spawn(scope, chunk) => chunk_stream( + &chunk_handle, + chunk, + scope, + chunk_io.session(), + Arc::clone(&chunk_map), + ), + }) + .try_flatten_unordered(in_flight); + // The prepared scan holds the file's shared cache: keep it until the stream ends, so + // scans of the file that start meanwhile find it. + let batches = batches.chain(stream::poll_fn(move |_| { + let _ = &scan; + Poll::Ready(None) + })); + if !self.ordered { + return batches.map_ok(|(_, batch)| batch).boxed(); + } + stream::once(async move { + let mut held: Vec<(u64, RecordBatch)> = batches.try_collect().await?; + // Morsels finish in whatever order their reads arrive. + held.sort_by_key(|(start, _)| *start); + VortexResult::Ok(stream::iter(held.into_iter().map(|(_, batch)| Ok(batch)))) + }) + .try_flatten() + .boxed() + } +} + +/// The planners of a scan, composed: with a filter, its proofs first over the whole scan, +/// continuing into the chunks of the rows they keep; without, the chunks directly. +fn compose(scan: &Arc) -> Box { + let rows = scan.rows().clone(); + let chunks = { + let scan = Arc::clone(scan); + let rows = rows.clone(); + move |kept: Option| -> Box { + let selection = scan.selection().clone(); + Box::new(ChunkPlanner::new(scan, rows, selection, kept)) + } + }; + if scan.filter().is_some() { + Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(chunks))) + } else { + chunks(None) + } +} + +/// What the root run's stream yields: a mapped batch with the row it starts at, or a chunk a +/// planner of the run spawned. +enum Yielded { + Batch(u64, RecordBatch), + Spawn(WorkScope, Box), +} + +/// One CPU step of a run, taken on a `spawn_cpu` task with the run. +enum Step { + Batch(u64, RecordBatch), + Spawn(WorkScope, Box), + RootDone, + Idle, + Waiting, +} + +/// The state of the root run's stream. +struct RunState { + handle: Handle, + /// Taken while a CPU step runs on its task and put back with its result. + run: Option, + io: Arc, + map: MapBatch, + /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. + inline: bool, +} + +/// Runs `root` over `scope` on a planning run of its own, reading through `io`, each CPU step +/// on a `spawn_cpu` task, and yields its mapped batches and, unless `inline`, the chunks it +/// spawns. Inline, a spawned chunk runs on this run beside the others: the run's steps already +/// interleave every chunk's work, and a run per chunk would only add a hand-off per chunk for +/// parallelism the engine provides. +fn run_stream( + handle: Handle, + root: Box, + scope: WorkScope, + io: Arc, + map: MapBatch, + inline: bool, +) -> BoxStream<'static, VortexResult> { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + let state = RunState { + handle, + run: Some(run), + io, + map, + inline, + }; + stream::try_unfold(state, |mut state| async move { + loop { + let run = state + .run + .take() + .ok_or_else(|| vortex_err!("the scan's run is missing"))?; + let map = Arc::clone(&state.map); + let (run, step) = state + .handle + .spawn_cpu(move || { + let mut run = run; + let step = run.advance().and_then(|progress| { + Ok(match progress { + Progress::Batch(batch) => { + Step::Batch(batch.scope.rows.start, map(batch.array)?) + } + Progress::Spawn(scope, child) => Step::Spawn(scope, child), + Progress::RootDone(_) => Step::RootDone, + Progress::Idle => Step::Idle, + Progress::Waiting => Step::Waiting, + }) + }); + (run, step) + }) + .await; + let run = state.run.insert(run); + match step? { + Step::Batch(start, batch) => { + return Ok(Some((Yielded::Batch(start, batch), state))); + } + Step::Spawn(scope, child) if state.inline => { + run.admit(child, scope, Arc::clone(&state.io)); + } + Step::Spawn(scope, child) => { + return Ok(Some((Yielded::Spawn(scope, child), state))); + } + Step::RootDone if run.live_roots() > 0 => {} + Step::RootDone | Step::Idle => return Ok(None), + Step::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }) + .boxed() +} + +/// Runs a spawned chunk's `root` over `scope` on a planning run of its own, on one task that +/// drives the run to its end and awaits only its reads, and yields its mapped batches with the +/// row each starts at. A planner the chunk spawns in turn runs as another root of the same run. +/// +/// One task per chunk, not one per step: a chunk's run waits for reads many times, and a task +/// hop at each wait is latency on the chunk's critical path that no other chunk's work fills. +fn chunk_stream( + handle: &Handle, + root: Box, + scope: WorkScope, + io: Arc, + 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)); + loop { + match run.advance()? { + Progress::Batch(batch) => { + let start = batch.scope.rows.start; + let batch = map(batch.array); + let failed = batch.is_err(); + if send.send(batch.map(|batch| (start, batch))).await.is_err() || failed { + return Ok(()); + } + } + Progress::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&io)); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => return Ok(()), + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }); + let errors = stream::once(task) + .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); + stream::select(recv, errors).boxed() +} From c7d3c3f1ba32f916228da02486651994c3321b56 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 15:13:15 +0100 Subject: [PATCH 03/22] Drive a V2 scan of a file for callers without an engine `vortex_file::v2::ScanDriver` composes a scan's planners and runs them on the session's runtime, the way an engine's driver does, for the file API's callers: the compress and random-access benchmarks move to it. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- benchmarks/compress-bench/src/vortex.rs | 17 +- vortex-bench/src/random_access/take.rs | 8 +- vortex-file/src/v2/mod.rs | 2 + vortex-file/src/v2/scan.rs | 253 ++++++++++++++++++++++++ 4 files changed, 272 insertions(+), 8 deletions(-) create mode 100644 vortex-file/src/v2/scan.rs diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index fbdeb027554..cb22673cb6a 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -3,22 +3,25 @@ use std::io::Cursor; use std::path::Path; -use std::sync::Arc; use std::time::Duration; use std::time::Instant; use anyhow::Result; use anyhow::bail; +use arrow_schema::Field; use async_trait::async_trait; use bytes::Bytes; use futures::StreamExt; use futures::pin_mut; use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; use vortex::dtype::FieldNames; use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; +use vortex::file::v2::ScanDriver; +use vortex::io::session::RuntimeSessionExt; use vortex::layout::scan::v2::ScanBuilder; use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; @@ -96,13 +99,17 @@ async fn vortex_decompress_read(data: Bytes) -> Result<()> { .optimize_recursive()?; scan = scan.with_projection(projection); } - let schema = Arc::new(SESSION.arrow().to_arrow_schema(scan.dtype())?); + let schema = SESSION.arrow().to_arrow_schema(scan.dtype())?; + let field = Field::new_struct("", schema.fields().clone(), false); - let stream = scan.into_record_batch_stream(schema)?; + let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan); pin_mut!(stream); - while let Some(batch) = stream.next().await { - let _batch = batch?; + while let Some(chunk) = stream.next().await { + let mut ctx = SESSION.create_execution_ctx(); + let _batch = SESSION + .arrow() + .execute_arrow(chunk?, Some(&field), &mut ctx)?; } Ok(()) } diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 164c402ecc9..f57ee614c08 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -35,6 +35,8 @@ use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; +use vortex::file::v2::ScanDriver; +use vortex::io::session::RuntimeSessionExt; use vortex::layout::scan::v2::ScanBuilder; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; @@ -190,9 +192,9 @@ impl RandomAccessor for VortexRandomAccessor { async fn take(&self, indices: &[u64]) -> anyhow::Result { let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?; - let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file()) - .with_row_indices(indices) - .into_stream()?; + let scan = + ScanBuilder::new(SESSION.clone(), self.file.scan_file()).with_row_indices(indices); + let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan); let array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream) .read_all() .await?; diff --git a/vortex-file/src/v2/mod.rs b/vortex-file/src/v2/mod.rs index 988916b65d1..f9ffd62549d 100644 --- a/vortex-file/src/v2/mod.rs +++ b/vortex-file/src/v2/mod.rs @@ -2,5 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod file_stats_reader; +mod scan; pub use file_stats_reader::*; +pub use scan::ScanDriver; diff --git a/vortex-file/src/v2/scan.rs b/vortex-file/src/v2/scan.rs new file mode 100644 index 00000000000..7af6c934f1a --- /dev/null +++ b/vortex-file/src/v2/scan.rs @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Drives a V2 scan of a file for callers without an engine of their own. +//! +//! A planner says what it needs: compute, or reads. This driver decides where each runs. It +//! composes the planners over a [`PreparedScan`], the zone proofs first when the scan has a +//! filter and the chunks of the rows they keep after, and runs them on planning runs it owns: +//! the root run stepped on `spawn_cpu` tasks as the stream is polled, each spawned chunk on a +//! task of its own, up to the concurrency, or on the root run at a concurrency of one. Every +//! read is a `spawn_io` task of the file's segment source. + +use std::future; +use std::sync::Arc; +use std::task::Poll; + +use futures::SinkExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::channel::mpsc; +use futures::stream; +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::runtime::Handle; +use vortex_layout::scan::v2::ChunkPlanner; +use vortex_layout::scan::v2::PreparedScan; +use vortex_layout::scan::v2::PrunePlanner; +use vortex_layout::scan::v2::ScanBuilder; +use vortex_mask::Mask; +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; + +/// Runs a V2 scan on the session's runtime. +pub struct ScanDriver { + handle: Handle, + concurrency: usize, + ordered: bool, +} + +impl ScanDriver { + /// A driver spawning on `handle`, running `concurrency` chunks beside each other per worker + /// thread, or every chunk on the scan's own run at one, and yielding batches in row order + /// when `ordered`, which holds them until the scan is done. + pub fn new(handle: Handle, concurrency: usize, ordered: bool) -> Self { + assert!(concurrency > 0); + Self { + handle, + concurrency, + ordered, + } + } + + /// Prepares `builder` where compute runs, composes its planners, and streams their batches. + /// Preparation errors are the stream's first item. + pub fn run(self, builder: ScanBuilder) -> BoxStream<'static, VortexResult> { + let handle = self.handle.clone(); + stream::once(async move { + handle + .spawn_cpu(move || { + let scan = Arc::new(builder.prepare()?); + VortexResult::Ok(self.execute(scan)) + }) + .await + }) + .try_flatten() + .boxed() + } + + fn execute(self, scan: Arc) -> BoxStream<'static, VortexResult> { + let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); + let inline = in_flight <= 1; + let scope = WorkScope { + file_ordinal: 0, + rows: scan.rows().clone(), + }; + let io = scan.io(); + let root = compose(&scan); + let (chunk_handle, chunk_io) = (self.handle.clone(), Arc::clone(&io)); + let batches = run_stream(self.handle, root, scope, io.session(), inline) + .map_ok(move |yielded| match yielded { + Yielded::Batch(start, batch) => { + stream::once(future::ready(Ok((start, batch)))).boxed() + } + Yielded::Spawn(scope, chunk) => { + chunk_stream(&chunk_handle, chunk, scope, chunk_io.session()) + } + }) + .try_flatten_unordered(in_flight); + // The prepared scan holds the file's shared cache: keep it until the stream ends, so + // scans of the file that start meanwhile find it. + let batches = batches.chain(stream::poll_fn(move |_| { + let _ = &scan; + Poll::Ready(None) + })); + if !self.ordered { + return batches.map_ok(|(_, batch)| batch).boxed(); + } + stream::once(async move { + let mut held: Vec<(u64, ArrayRef)> = batches.try_collect().await?; + // Morsels finish in whatever order their reads arrive. + held.sort_by_key(|(start, _)| *start); + VortexResult::Ok(stream::iter(held.into_iter().map(|(_, batch)| Ok(batch)))) + }) + .try_flatten() + .boxed() + } +} + +/// The planners of a scan, composed: with a filter, its proofs first over the whole scan, +/// continuing into the chunks of the rows they keep; without, the chunks directly. +fn compose(scan: &Arc) -> Box { + let rows = scan.rows().clone(); + let chunks = { + let scan = Arc::clone(scan); + let rows = rows.clone(); + move |kept: Option| -> Box { + let selection = scan.selection().clone(); + Box::new(ChunkPlanner::new(scan, rows, selection, kept)) + } + }; + if scan.filter().is_some() { + Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(chunks))) + } else { + chunks(None) + } +} + +/// What the root run's stream yields: a batch with the row it starts at, or a chunk a planner of +/// the run spawned. +enum Yielded { + Batch(u64, ArrayRef), + Spawn(WorkScope, Box), +} + +/// One CPU step of a run, taken on a `spawn_cpu` task with the run. +enum Step { + Batch(u64, ArrayRef), + Spawn(WorkScope, Box), + RootDone, + Idle, + Waiting, +} + +/// The state of the root run's stream. +struct RunState { + handle: Handle, + /// Taken while a CPU step runs on its task and put back with its result. + run: Option, + io: Arc, + /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. + inline: bool, +} + +/// Runs `root` over `scope` on a planning run of its own, reading through `io`, each CPU step +/// on a `spawn_cpu` task, and yields its batches and, unless `inline`, the chunks it spawns. +fn run_stream( + handle: Handle, + root: Box, + scope: WorkScope, + io: Arc, + inline: bool, +) -> BoxStream<'static, VortexResult> { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + let state = RunState { + handle, + run: Some(run), + io, + inline, + }; + stream::try_unfold(state, |mut state| async move { + loop { + let run = state + .run + .take() + .ok_or_else(|| vortex_err!("the scan's run is missing"))?; + let (run, step) = state + .handle + .spawn_cpu(move || { + let mut run = run; + let step = run.advance().map(|progress| match progress { + Progress::Batch(batch) => Step::Batch(batch.scope.rows.start, batch.array), + Progress::Spawn(scope, child) => Step::Spawn(scope, child), + Progress::RootDone(_) => Step::RootDone, + Progress::Idle => Step::Idle, + Progress::Waiting => Step::Waiting, + }); + (run, step) + }) + .await; + let run = state.run.insert(run); + match step? { + Step::Batch(start, batch) => { + return Ok(Some((Yielded::Batch(start, batch), state))); + } + Step::Spawn(scope, child) if state.inline => { + run.admit(child, scope, Arc::clone(&state.io)); + } + Step::Spawn(scope, child) => { + return Ok(Some((Yielded::Spawn(scope, child), state))); + } + Step::RootDone if run.live_roots() > 0 => {} + Step::RootDone | Step::Idle => return Ok(None), + Step::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }) + .boxed() +} + +/// Runs a spawned chunk's `root` over `scope` on a planning run of its own, on one task that +/// drives the run to its end and awaits only its reads, and yields its batches with the row each +/// starts at. A planner the chunk spawns in turn runs as another root of the same run. +fn chunk_stream( + handle: &Handle, + root: Box, + scope: WorkScope, + io: Arc, +) -> BoxStream<'static, VortexResult<(u64, ArrayRef)>> { + 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)); + loop { + match run.advance()? { + Progress::Batch(batch) => { + if send + .send(Ok((batch.scope.rows.start, batch.array))) + .await + .is_err() + { + return Ok(()); + } + } + Progress::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&io)); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => return Ok(()), + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }); + let errors = stream::once(task) + .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); + stream::select(recv, errors).boxed() +} From 95424b8e0c71493ec0de21555eef70dd1f8df6a5 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 15:15:21 +0100 Subject: [PATCH 04/22] Compose a V2 scan's planners in the DataFusion opener The opener prepares the scan where compute runs, composes its planners itself, a prune planner continuing into a chunk planner when there is a filter and the chunk planner alone when there is not, and hands the driver the first planner. The driver only drives. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-datafusion/src/persistent/opener.rs | 36 +++++++++++- vortex-datafusion/src/persistent/scan_v2.rs | 63 ++++----------------- 2 files changed, 44 insertions(+), 55 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index f85482a31ba..c7e31a65f94 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -56,13 +56,17 @@ use vortex::io::InstrumentedReadAt; use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::layout::scan::v2::ChunkPlanner; +use vortex::layout::scan::v2::PrunePlanner; use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; use vortex::layout::scan::v2::ScanCache; use vortex::layout::scan::v2::WeakScanCache; use vortex::layout::scan::v2::partition_row_range; use vortex::layout::segments::SegmentSource; +use vortex::mask::Mask; use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; +use vortex::scan::planning::planner::Planner; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; use vortex_utils::aliases::dash_map::DashMap; @@ -463,9 +467,35 @@ impl FileOpener for VortexOpener { .div_ceil(partitions) .max(1) }); - let driver = - ScanDriver::new(vxf.session().handle(), concurrency, has_output_ordering); - Ok(driver.run(builder, map_batch)) + let handle = vxf.session().handle(); + // Lowering and planning a wide file is real CPU: it runs where compute runs. + let scan = handle + .spawn_cpu(move || builder.prepare()) + .await + .map_err(|e| exec_datafusion_err!("Failed to plan Vortex scan: {e}"))?; + let scan = Arc::new(scan); + + // The planners, composed: with a filter, its zone proofs first over the whole + // partition, continuing into the chunks of the rows they keep; without a + // filter, the chunks directly. A chunk plans its conjuncts, then its + // projection. + let chunks = { + let scan = Arc::clone(&scan); + move |kept: Option| -> Box { + let rows = scan.rows().clone(); + let selection = scan.selection().clone(); + Box::new(ChunkPlanner::new(scan, rows, selection, kept)) + } + }; + let root: Box = if scan.filter().is_some() { + let rows = scan.rows().clone(); + Box::new(PrunePlanner::new(Arc::clone(&scan), rows, Box::new(chunks))) + } else { + chunks(None) + }; + + let driver = ScanDriver::new(handle, concurrency, has_output_ordering); + Ok(driver.run(scan, root, map_batch)) } else { // Share V1 layout readers across partitions of the file. let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { diff --git a/vortex-datafusion/src/persistent/scan_v2.rs b/vortex-datafusion/src/persistent/scan_v2.rs index 10b36559f56..10ea7b969ea 100644 --- a/vortex-datafusion/src/persistent/scan_v2.rs +++ b/vortex-datafusion/src/persistent/scan_v2.rs @@ -3,13 +3,13 @@ //! Drives a V2 scan for DataFusion. //! -//! A planner says what it needs: compute, or reads. This driver decides where each runs. It -//! composes the planners over a [`PreparedScan`], the zone proofs first when the scan has a -//! filter and the chunks of the rows they keep after, and runs them on planning runs it owns: -//! the root run stepped on `spawn_cpu` tasks as the stream is polled, each spawned chunk on a -//! task of its own, or on the root run when the engine's partitions already fill the cores. -//! Every read is a `spawn_io` task of the file's segment source. The DataFusion thread polling -//! the stream does no scan work itself. +//! A planner says what it needs: compute, or reads. This driver decides where each runs. The +//! opener composes the planners over a [`PreparedScan`] and hands the driver the first; the +//! driver runs it and everything it emits on planning runs it owns: the root run stepped on +//! `spawn_cpu` tasks as the stream is polled, each spawned chunk on a task of its own, or on the +//! root run when the engine's partitions already fill the cores. Every read is a `spawn_io` task +//! of the file's segment source. The DataFusion thread polling the stream does no scan work +//! itself. use std::future; use std::sync::Arc; @@ -27,11 +27,7 @@ use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::io::request::IoSource; use vortex::io::runtime::Handle; -use vortex::layout::scan::v2::ChunkPlanner; use vortex::layout::scan::v2::PreparedScan; -use vortex::layout::scan::v2::PrunePlanner; -use vortex::layout::scan::v2::ScanBuilder; -use vortex::mask::Mask; use vortex::scan::planning::driver::Progress; use vortex::scan::planning::driver::Run; use vortex::scan::planning::planner::Planner; @@ -59,32 +55,15 @@ impl ScanDriver { } } - /// Prepares `builder` where compute runs, composes its planners, and streams their batches - /// mapped by `map`. Preparation errors are the stream's first item. + /// Runs `root`, the first planner composed over `scan`, and everything it emits, and + /// streams their batches mapped by `map`. pub(crate) fn run( self, - builder: ScanBuilder, + scan: Arc, + root: Box, map: impl Fn(ArrayRef) -> VortexResult + Send + Sync + 'static, ) -> BoxStream<'static, VortexResult> { - let handle = self.handle.clone(); let map: MapBatch = Arc::new(map); - stream::once(async move { - handle - .spawn_cpu(move || { - let scan = Arc::new(builder.prepare()?); - VortexResult::Ok(self.execute(scan, map)) - }) - .await - }) - .try_flatten() - .boxed() - } - - fn execute( - self, - scan: Arc, - map: MapBatch, - ) -> BoxStream<'static, VortexResult> { let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); let inline = in_flight <= 1; let scope = WorkScope { @@ -92,7 +71,6 @@ impl ScanDriver { rows: scan.rows().clone(), }; let io = scan.io(); - let root = compose(&scan); 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) @@ -129,25 +107,6 @@ impl ScanDriver { } } -/// The planners of a scan, composed: with a filter, its proofs first over the whole scan, -/// continuing into the chunks of the rows they keep; without, the chunks directly. -fn compose(scan: &Arc) -> Box { - let rows = scan.rows().clone(); - let chunks = { - let scan = Arc::clone(scan); - let rows = rows.clone(); - move |kept: Option| -> Box { - let selection = scan.selection().clone(); - Box::new(ChunkPlanner::new(scan, rows, selection, kept)) - } - }; - if scan.filter().is_some() { - Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(chunks))) - } else { - chunks(None) - } -} - /// What the root run's stream yields: a mapped batch with the row it starts at, or a chunk a /// planner of the run spawned. enum Yielded { From 04de7678ce4ab60399a01b55a7618af7ea6dcc2b Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 15:45:26 +0100 Subject: [PATCH 05/22] Drop is_proved: a proved proof lists no segment The prune planner skipped prefetching the zone maps of proofs an earlier scan of the file had proved. The segment list now comes from compiling the proof, and a proved proof compiles to its kept result without a segment scan, so the check repeated what compile already knew. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- .../src/plan/physical/plans/zoned.rs | 9 ------ vortex-layout/src/scan/execution/prune.rs | 29 ++----------------- 2 files changed, 3 insertions(+), 35 deletions(-) diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index 504e2aff126..45a678cc873 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -241,15 +241,6 @@ impl ZonedPlan { self.child_required(index) } - /// Whether an earlier execution already proved this pruning plan's expression over its zones, - /// so compiling it emits the kept result without reading the zone table. - pub(crate) fn is_proved(&self, session: &VortexSession) -> VortexResult { - match (self.pruned_zones(), self.pruning_expression()) { - (Some(pruned), Some(expression)) => Ok(pruned.get(expression, session)?.is_some()), - _ => Ok(false), - } - } - /// Returns whether this plan represents a zone-backed pruning proof. pub fn is_pruning(&self) -> bool { self.data().pruning.is_some() diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 9389d9b2d72..199c888242e 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -20,7 +20,6 @@ use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; use vortex_scan::planning::planner::WorkScope; -use vortex_session::VortexSession; use super::PreparedScan; use super::adaptor::AdaptorStep; @@ -29,25 +28,9 @@ use super::filter::hint; use super::filter::piece_mask; use super::filter::segments_read; use crate::plan::physical::PlanRef; -use crate::plan::physical::Zoned; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; -/// Whether every zone proof in `plan` has already been proved, so compiling it reads no zones. -fn is_proved(plan: &PlanRef, session: &VortexSession) -> VortexResult { - if let Some(zoned) = plan.as_opt::() - && zoned.is_pruning() - { - return zoned.is_proved(session); - } - for child in plan.children().iter() { - if !is_proved(&child?, session)? { - return Ok(false); - } - } - Ok(true) -} - /// The planner a prune planner continues into, given the rows the proofs kept, or `None` when /// no proof ran. pub type PruneNext = Box) -> Box + Send>; @@ -124,17 +107,11 @@ impl Planner for PrunePlanner { } if !self.announced { self.announced = true; - // A proof whose zones an earlier scan of the file proved reads nothing: a shared - // plan keeps its proved zones, so only unproved proofs need their zone maps. - let mut unproved = Vec::new(); - for proof in self.proofs.as_slice() { - if !is_proved(proof, &self.scan.session)? { - unproved.push(proof); - } - } + // A proof whose zones an earlier scan of the file proved compiles to its kept + // result and lists no segment, so only unproved proofs prefetch their zone maps. let segments = segments_read( &self.scan.session, - unproved.into_iter(), + self.proofs.as_slice().iter(), self.rows.clone(), &self.mask, )?; From 429e00510e28c86aacf5c7310bab3efd8f3086e3 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 17:32:50 +0100 Subject: [PATCH 06/22] Call the zone-map plans pruning plans, not proofs The scan code named the filter's zone-map plans proofs. They are the filter's pruning expressions evaluated over the file's zone map, so the layout, file and DataFusion code now says pruning, matching prune_zones, PrunedZones and the pruning field they already used. Behaviour is unchanged. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Joe Isaacs --- vortex-datafusion/src/persistent/opener.rs | 4 +- vortex-file/src/v2/scan.rs | 4 +- .../src/plan/physical/plans/legacy_stats.rs | 6 +-- .../src/plan/physical/plans/zoned.rs | 10 ++-- vortex-layout/src/plan/physical/tests.rs | 2 +- vortex-layout/src/scan/execution/chunk.rs | 2 +- vortex-layout/src/scan/execution/filter.rs | 12 +++-- vortex-layout/src/scan/execution/mod.rs | 4 +- vortex-layout/src/scan/execution/prune.rs | 54 +++++++++---------- vortex-layout/src/scan/v2/cache.rs | 8 +-- vortex-layout/src/scan/v2/mod.rs | 2 +- vortex-layout/src/scan/v2/tests.rs | 12 ++--- 12 files changed, 61 insertions(+), 59 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index c7e31a65f94..208691f0b0e 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -122,7 +122,7 @@ pub(crate) struct VortexOpener { /// shared layout reader. pub segment_sources: Arc>>, /// The V2 scan cache each open file's partitions share: its lowered layout, filter and - /// projection plans, proved zones and dictionary values, built once per file. + /// projection plans, pruned zones and dictionary values, built once per file. pub scan_caches: Arc>, /// Shared V1 full-file natural splits keyed by file path. pub natural_splits: Arc>>, @@ -475,7 +475,7 @@ impl FileOpener for VortexOpener { .map_err(|e| exec_datafusion_err!("Failed to plan Vortex scan: {e}"))?; let scan = Arc::new(scan); - // The planners, composed: with a filter, its zone proofs first over the whole + // The planners, composed: with a filter, its pruning plans first over the whole // partition, continuing into the chunks of the rows they keep; without a // filter, the chunks directly. A chunk plans its conjuncts, then its // projection. diff --git a/vortex-file/src/v2/scan.rs b/vortex-file/src/v2/scan.rs index 7af6c934f1a..14753e58777 100644 --- a/vortex-file/src/v2/scan.rs +++ b/vortex-file/src/v2/scan.rs @@ -4,7 +4,7 @@ //! Drives a V2 scan of a file for callers without an engine of their own. //! //! A planner says what it needs: compute, or reads. This driver decides where each runs. It -//! composes the planners over a [`PreparedScan`], the zone proofs first when the scan has a +//! composes the planners over a [`PreparedScan`], the pruning plans first when the scan has a //! filter and the chunks of the rows they keep after, and runs them on planning runs it owns: //! the root run stepped on `spawn_cpu` tasks as the stream is polled, each spawned chunk on a //! task of its own, up to the concurrency, or on the root run at a concurrency of one. Every @@ -112,7 +112,7 @@ impl ScanDriver { } } -/// The planners of a scan, composed: with a filter, its proofs first over the whole scan, +/// The planners of a scan, composed: with a filter, its pruning plans first over the whole scan, /// continuing into the chunks of the rows they keep; without, the chunks directly. fn compose(scan: &Arc) -> Box { let rows = scan.rows().clone(); diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index 37d077a6f2c..8aa0e897da7 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/vortex-layout/src/plan/physical/plans/legacy_stats.rs @@ -31,7 +31,7 @@ use crate::plan::physical::optimizer::PlanParentReduceRule; use crate::plan::physical::optimizer::PlanReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::plans::zoned::is_zone_proof; +use crate::plan::physical::plans::zoned::is_zone_pruning; const DATA: usize = 0; const ZONES: usize = 1; @@ -181,7 +181,7 @@ impl PlanReduceRule for LegacyStatsRule { } } -/// Converts the legacy stats beneath a zone proof in place, so the proof meets a [`Zoned`] plan +/// Converts the legacy stats beneath a zone pruning plan in place, so the plan meets a [`Zoned`] plan /// before the data path optimizes the stats away. #[derive(Debug)] pub(crate) struct ExpressionLegacyStatsRule; @@ -196,7 +196,7 @@ impl PlanParentReduceRule for ExpressionLegacyStatsRule { _child_idx: usize, session: &VortexSession, ) -> VortexResult> { - if !is_zone_proof(parent)? { + if !is_zone_pruning(parent)? { return Ok(None); } let Some(converted) = LegacyStatsRule.reduce(child, session)? else { diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index 45a678cc873..15bf8cd7dfd 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -456,7 +456,7 @@ impl PlanParentReduceRule for ExpressionZonedRule { _child_idx: usize, _session: &VortexSession, ) -> VortexResult> { - if !is_zone_proof(parent)? { + if !is_zone_pruning(parent)? { return Ok(None); } Ok(child @@ -466,8 +466,8 @@ impl PlanParentReduceRule for ExpressionZonedRule { } /// Whether `eval` is a boolean expression over abstract statistics alone, which a zone table can -/// prove without reading the data it summarises. -pub(crate) fn is_zone_proof(eval: &Plan) -> VortexResult { +/// prune without reading the data it summarises. +pub(crate) fn is_zone_pruning(eval: &Plan) -> VortexResult { let mut contains_stat = false; let mut contains_root = false; eval.expression().clone().transform_down(|expression| { @@ -489,8 +489,8 @@ pub(crate) fn is_zone_proof(eval: &Plan) -> VortexResult { } /// Replaces a zoned data plan by its data: a scan reads the data through it and never its zone -/// table. Pruning proofs are derived from the filter over the lowered plan before this runs, so -/// an optimized plan carries a zone table only where a proof reads it. +/// table. Pruning plans are derived from the filter over the lowered plan before this runs, so +/// an optimized plan carries a zone table only where a pruning plan reads it. #[derive(Debug)] pub(crate) struct ZonedDataRule; diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 642ea16e612..c86edcf4d30 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -1602,7 +1602,7 @@ fn zoned_over_legacy_stats_read_as_data_keeps_only_the_data() -> VortexResult<() } /// Legacy stats over a zoned plan are dropped: read as data only the data remains, and a zone -/// proof is proved over the aggregate zones, never the legacy table. +/// pruning plan is evaluated over the aggregate zones, never the legacy table. #[test] fn legacy_stats_over_zoned_are_dropped() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 7ef7d24bcba..382425d9bb0 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -66,7 +66,7 @@ pub struct ChunkPlanner { rows: Range, chunks: IntoIter>, selection: Selection, - /// The rows zone pruning kept, over `rows`, or `None` when nothing was proved. + /// The rows zone pruning kept, over `rows`, or `None` when nothing was pruned. kept: Option, } diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 1873e5af62a..d3d7b640892 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -48,7 +48,7 @@ use crate::segments::SegmentId; pub struct FilterPlans { expression: FilterExpr, predicates: Vec, - /// A metadata-only proof per conjunct, where the conjunct has one. + /// A metadata-only pruning plan per conjunct, where the conjunct has one. pruning: Vec>, } @@ -75,7 +75,9 @@ impl FilterPlans { .iter() .map(|expr| { expr.falsify(session)? - .map(|proof| optimize(plan_row_idx_expression(proof, root.clone())?, session)) + .map(|pruning| { + optimize(plan_row_idx_expression(pruning, root.clone())?, session) + }) .transpose() }) .map(|plan| { @@ -93,13 +95,13 @@ impl FilterPlans { }) } - /// The pruning proofs, run once over a scan's rows before its chunks are planned. - pub(crate) fn proofs(&self) -> impl Iterator { + /// The pruning plans, run once over a scan's rows before its chunks are planned. + pub(crate) fn pruning(&self) -> impl Iterator { self.pruning.iter().flatten() } } -/// A pruning proof must not fetch data when its requested statistics are unavailable. +/// A pruning plan must not fetch data when its requested statistics are unavailable. fn metadata_only(plan: &PlanRef) -> VortexResult { if plan .as_opt::() diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 33200450e26..0185f61d6a9 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -4,7 +4,7 @@ //! The planners of a V2 scan, and the prepared plans they run. //! //! A [`PreparedScan`] holds a scan's plans over one file. An engine composes the planners over it -//! and drives them: a [`PrunePlanner`] proves the filter's zone maps over the scan's rows and +//! and drives them: a [`PrunePlanner`] evaluates the filter's pruning plans over the scan's rows and //! emits the planner it was built to continue into, with the rows they keep; a [`ChunkPlanner`] //! cuts those rows into chunks and spawns a [`FilterPlanner`] or a [`ProjectionPlanner`] per //! chunk; a filter planner runs its chunk's conjuncts and emits a projection planner for the rows @@ -91,7 +91,7 @@ pub struct PreparedScan { pub(super) locations: Arc<[SegmentLocation]>, /// The optimized projection plan, over the scan's root row domain. pub(super) projection: PlanRef, - /// The filter's conjuncts, their plans and their proofs. + /// The filter's conjuncts, their plans and their pruning plans. pub(super) filter: Option>, /// The global row index of the root row domain's first row. pub(super) row_offset: u64, diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 199c888242e..325f5f897d7 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Proves a filter's zone maps once over a scan's rows, before any chunk is planned. +//! Evaluates a filter's pruning plans once over a scan's rows, before any chunk is planned. use std::ops::Range; use std::sync::Arc; @@ -31,24 +31,24 @@ use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; -/// The planner a prune planner continues into, given the rows the proofs kept, or `None` when -/// no proof ran. +/// The planner a prune planner continues into, given the rows pruning kept, or `None` when +/// no pruning plan ran. pub type PruneNext = Box) -> Box + Send>; -/// Runs the filter's zone-map proofs over the scan's rows and emits the planner it was built -/// to continue into, with the rows no proof excludes. The zone maps are the only reads, -/// started together. Rows every proof excludes end the scan here: nothing is emitted. A filter -/// without proofs emits the next planner at once. +/// Evaluates the filter's pruning plans over the scan's rows and emits the planner it was +/// built to continue into, with the rows no plan pruned. The zone maps are the only reads, +/// started together. Rows every plan prunes end the scan here: nothing is emitted. A filter +/// without pruning plans emits the next planner at once. pub struct PrunePlanner { scan: Arc, rows: Range, - /// What to emit once every proof has run, taken when it is. + /// What to emit once every pruning plan has run, taken when it is. next: Option, - /// Whether any proof ran, so the kept rows mean something. - proved: bool, - /// The rows no proof has excluded yet, over `rows`. + /// Whether any pruning plan ran, so the kept rows mean something. + pruned: bool, + /// The rows no pruning plan has excluded yet, over `rows`. mask: Mask, - proofs: IntoIter, + pruning: IntoIter, active: Option, pieces: Vec, announced: bool, @@ -58,21 +58,21 @@ pub struct PrunePlanner { } impl PrunePlanner { - /// The planner of the proofs of `scan`'s filter over `rows`, continuing into `next`. + /// The planner of `scan`'s filter's pruning plans over `rows`, continuing into `next`. pub fn new(scan: Arc, rows: Range, next: PruneNext) -> Self { - let proofs: Vec<_> = scan + let pruning: Vec<_> = scan .filter .iter() - .flat_map(|filter| filter.proofs().cloned()) + .flat_map(|filter| filter.pruning().cloned()) .collect(); let len = usize::try_from(rows.end - rows.start).unwrap_or(usize::MAX); Self { scan, rows, next: Some(next), - proved: false, + pruned: false, mask: Mask::new_true(len), - proofs: proofs.into_iter(), + pruning: pruning.into_iter(), active: None, pieces: Vec::new(), announced: false, @@ -107,11 +107,11 @@ impl Planner for PrunePlanner { } if !self.announced { self.announced = true; - // A proof whose zones an earlier scan of the file proved compiles to its kept - // result and lists no segment, so only unproved proofs prefetch their zone maps. + // A pruning plan whose zones an earlier scan of the file pruned compiles to its kept + // result and lists no segment, so only unpruned plans prefetch their zone maps. let segments = segments_read( &self.scan.session, - self.proofs.as_slice().iter(), + self.pruning.as_slice().iter(), self.rows.clone(), &self.mask, )?; @@ -149,20 +149,20 @@ impl Planner for PrunePlanner { row_end = self.rows.end, input_rows = input, output_rows = self.mask.true_count(), - "v2 zone proof evaluated" + "v2 pruning plan evaluated" ); self.active = None; self.pieces.clear(); - self.proved = true; + self.pruned = true; } if self.mask.all_false() { self.done = true; return Ok(PlannerOutput::Done); } - if let Some(proof) = self.proofs.next() { + if let Some(plan) = self.pruning.next() { let scan = Scan::try_new( self.scan.session.clone(), - proof, + plan, vec![Split { rows: self.rows.clone(), mask: self.mask.clone(), @@ -173,20 +173,20 @@ impl Planner for PrunePlanner { self.active = Some(ScanAdaptor::new(scan, Arc::clone(&self.scan.locations))); return Ok(PlannerOutput::Continue); } - // Every proof has run: the planner this one continues into gets the kept rows. + // Every pruning plan has run: the planner this one continues into gets the kept rows. self.done = true; tracing::debug!( row_start = self.rows.start, row_end = self.rows.end, kept_rows = self.mask.true_count(), - "v2 zone proofs finished" + "v2 pruning finished" ); let next = self .next .take() .ok_or_else(|| vortex_err!("the prune planner has nothing to continue into"))?; let kept = self - .proved + .pruned .then(|| std::mem::replace(&mut self.mask, Mask::new_true(0))); Ok(PlannerOutput::Planner( WorkScope { diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 5a26e28ec97..c621bf19289 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -23,8 +23,8 @@ use crate::scan::execution::FilterPlans; /// over it, indexes every segment's byte range, and decodes the dictionary values it needs. A scan /// that finds an entry here reuses it, and one that misses builds the entry and keeps it for later /// scans. The lowered plan keeps the children it has initialized, each distinct projection keeps -/// its optimized plan, each distinct filter keeps its conjunct and zone-proof plans along with the -/// zones those proofs have already excluded, and dictionary values stay keyed by the cached plan +/// its optimized plan, each distinct filter keeps its conjunct and pruning plans along with the +/// zones those plans have already excluded, and dictionary values stay keyed by the cached plan /// that produced them. /// /// A cache describes one file's layout, segment locations, and segment source. Use a new cache @@ -40,8 +40,8 @@ struct Inner { /// them, so a projection must reuse its plan for its values to be found again. projections: Mutex>, ranges: Mutex>, - /// Plans of filters, keyed by the exact filter expression. A proof plan remembers the zones it - /// excluded, so sharing the plans proves a file's zones once for every scan of that filter. + /// Plans of filters, keyed by the exact filter expression. A pruning plan remembers the zones it + /// excluded, so sharing the plans prunes a file's zones once for every scan of that filter. filters: Mutex>>, values: ValueCache, } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 77c7c1fb29b..1f1566ed175 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -17,7 +17,7 @@ //! layout's chunk boundaries. Full-file projections use smaller 8,192-row multiples when //! serialized bytes per row are high. //! -//! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and proved +//! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and pruned //! zones, segment index, and dictionary values with every other scan of that file, like a cached //! layout reader does for V1 scans. diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 822b0a2cafc..eff68dfe996 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -197,7 +197,7 @@ fn builder( )) } -/// Composes a scan's planners the way an engine does: the proofs first when there is a filter, +/// Composes a scan's planners the way an engine does: the pruning plans first when there is a filter, /// continuing into the chunks of the rows they keep. fn compose(scan: &Arc, rows: Range) -> Box { let chunks = { @@ -805,7 +805,7 @@ fn partition_rows_tile_the_file_on_row_blocks() { ); } -/// Zone proofs run once before any chunk is planned: rows every proof excludes are never +/// Pruning runs once before any chunk is planned: rows every pruning plan excludes are never /// given a chunk, and a scan whose zones all fail plans nothing. #[rstest] #[case::tail(2_999, Some(3_000..4_000))] @@ -835,10 +835,10 @@ async fn pruned_rows_get_no_chunk( Ok(()) } -/// Scans of one file that share a cache share their filter plans, so the zone proofs run once: +/// Scans of one file that share a cache share their filter plans, so the pruning plans run once: /// the second scan with the same filter re-reads the data it needs but never the zone maps. #[tokio::test] -async fn shared_cache_proves_zones_once() -> VortexResult<()> { +async fn shared_cache_prunes_zones_once() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, zoned_table().as_ref()).await?; let counted = Arc::new(CountingSegments { @@ -866,8 +866,8 @@ async fn shared_cache_proves_zones_once() -> VortexResult<()> { let plans = cache .filter_plans(&filter) .expect("the first scan kept its filter plans"); - assert_eq!(plans.proofs().count(), 1); - // Each scan reads the one chunk the proof keeps; only the first reads the zone map. + assert_eq!(plans.pruning().count(), 1); + // Each scan reads the one chunk pruning keeps; only the first reads the zone map. assert_eq!(reads, vec![2, 1]); Ok(()) } From 1db58aae169756e57a4ce32d2f443b7a04f9e119 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 15:33:52 +0000 Subject: [PATCH 07/22] Add a planner-driven V2 scan path to the DuckDB extension DuckDB scans can now run on the V2 scan executor, like DataFusion's scan_v2 option. Each file is cut into one block-aligned row range per core, and every range is a V2 scan stream that shares the file's ScanCache. The extension enables it with VORTEX_SCAN_V2=1 or vortex_duckdb::set_scan_v2. duckdb-bench gains --scan-v2, which defaults to true for TPC-H, matching datafusion-bench. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- benchmarks/duckdb-bench/src/main.rs | 10 ++ .../src/e2e_test/vortex_scan_test.rs | 47 +++++++ vortex-duckdb/src/file_reader.rs | 130 +++++++++++++++--- vortex-duckdb/src/lib.rs | 16 +++ vortex-duckdb/src/table_function.rs | 12 +- 5 files changed, 192 insertions(+), 23 deletions(-) diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index 352d07bd5a2..a047e275b24 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -102,6 +102,10 @@ struct Args { to keep all work on the same threads" )] reuse: bool, + + /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H). + #[arg(long, action = clap::ArgAction::Set)] + scan_v2: Option, } #[global_allocator] @@ -112,6 +116,12 @@ 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") + }); + vortex_duckdb::set_scan_v2(scan_v2); + setup_logging_and_tracing(args.verbose, args.tracing)?; let benchmark = create_benchmark(args.benchmark, &opts)?; diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 628088f36d2..97b310de84f 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -1390,3 +1390,50 @@ fn test_duplicate_struct_fields() { .is_err() ); } + +/// Runs `query` on the file at `path` and collects its single BIGINT output column. +fn query_i64_column(path: &str, query: &str) -> Result> { + let conn = database_connection(); + let result = conn.query(&query.replace('?', &format!("'{path}'")))?; + let mut values = Vec::new(); + for chunk in result { + let vec = chunk.get_vector(0); + values.extend_from_slice(vec.as_slice_with_len::(chunk.len().as_())); + } + Ok(values) +} + +#[rstest] +#[case::filtered_count("SELECT count(*) FROM ? WHERE b < 100")] +#[case::filtered_max("SELECT max(a) FROM ? WHERE b BETWEEN 300 AND 310")] +#[case::filtered_projection("SELECT a FROM ? WHERE b = 42 ORDER BY a")] +#[case::expression("SELECT a + b FROM ? WHERE a >= 50000 ORDER BY 1")] +#[case::full_scan("SELECT a FROM ? ORDER BY a")] +#[case::empty("SELECT a FROM ? WHERE b > 5000")] +fn test_vortex_scan_v2_matches_v1(#[case] query: &str) -> Result<()> { + // Many row blocks, so the file is cut into several V2 splits. + let rows = 100_000i64; + let file = RUNTIME.block_on(write_vortex_file( + [ + ("a", (0..rows).collect::().into_array()), + ( + "b", + (0..rows) + .map(|i| (i * 37) % 1000) + .collect::() + .into_array(), + ), + ] + .into_iter(), + )); + let path = file.path().to_string_lossy(); + + crate::set_scan_v2(false); + let v1 = query_i64_column(&path, query)?; + crate::set_scan_v2(true); + let v2 = query_i64_column(&path, query); + crate::set_scan_v2(false); + + assert_eq!(v2?, v1); + Ok(()) +} diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index 1cd7ff87a3d..c18df0fec4c 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -1,13 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use futures::FutureExt; +use futures::StreamExt; use object_store::registry::ObjectStoreRegistry; use url::Url; +use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::struct_::StructArrayExt as _; use vortex::dtype::FieldPath; @@ -26,8 +29,12 @@ 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::v2::ScanBuilder as V2ScanBuilder; +use vortex::layout::scan::v2::ScanFile; +use vortex::layout::scan::v2::partition_row_range; use vortex::mask::Mask; use vortex::session::SessionExt as _; +use vortex_utils::parallelism::get_available_parallelism; use crate::REGISTRY; use crate::RUNTIME; @@ -114,6 +121,9 @@ pub struct OpenFileReader { /// The state of an open file that the threads scanning it share. pub struct FileScanState { reader: LayoutReaderRef, + /// The file as the V2 executor reads it, with a scan cache that every split of the file + /// shares. + scan_file: ScanFile, cache: ConversionCache, } @@ -122,9 +132,12 @@ impl OpenFileReader { let (fs, fs_path) = resolve_filesystem(&parse_uri_or_path(&path)?)?; let source = fs.open_read(&fs_path).await?; let file = open_cached(&SESSION, Some(&path), source, None, &|options| options).await?; + let mut scan_file = file.scan_file(); + scan_file.cache.get_or_insert_default(); Ok(OpenFileReader { scan_state: Some(Arc::new(FileScanState { reader: file.layout_reader()?, + scan_file, cache: ConversionCache::default(), })), splits: vec![], @@ -132,12 +145,14 @@ impl OpenFileReader { }) } - fn reader(&self) -> &LayoutReaderRef { - &self - .scan_state + fn scan_state(&self) -> &FileScanState { + self.scan_state .as_ref() .vortex_expect("file used after its scan state was released") - .reader + } + + fn reader(&self) -> &LayoutReaderRef { + &self.scan_state().reader } fn can_skip(&self, filter: &Filter) -> VortexResult { @@ -197,14 +212,21 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor // Getting splits is non-trivial work so we prefer doing it here under file // lock and not in reader_try_initialize_scan under global lock. - let reader = Arc::clone(file.reader()); - let filter = &global.filter; - let builder = ScanBuilder::new(SESSION.clone(), reader) - .with_projection(global.projection.clone()) - .with_some_filter(filter.filter.clone()) - .with_selection(filter.row_selection.clone()); - let scan = builder.prepare()?; - let mut splits = scan.execute(filter.row_range.clone())?; + let mut splits = if global.scan_v2 { + v2_splits(file.scan_state(), global)? + } else { + let reader = Arc::clone(file.reader()); + let filter = &global.filter; + let builder = ScanBuilder::new(SESSION.clone(), reader) + .with_projection(global.projection.clone()) + .with_some_filter(filter.filter.clone()) + .with_selection(filter.row_selection.clone()); + let scan = builder.prepare()?; + scan.execute(filter.row_range.clone())? + .into_iter() + .map(Split::V1) + .collect() + }; // threads take last element of file.splits so we need to reverse splits.reverse(); @@ -213,6 +235,72 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor Ok(false) } +/// Splits the file's scanned rows into one V2 scan per available core, each over a contiguous +/// row range on the executor's row blocks. The splits share the file's scan cache, so the file is +/// planned once. A split's chunks run beside each other only as far as the engine's threads leave +/// cores idle, like DataFusion's partitions. +fn v2_splits(state: &FileScanState, global: &GlobalState) -> VortexResult> { + let filter = &global.filter; + let row_count = state.reader.row_count(); + let rows = filter.row_range.as_ref().map_or(0..row_count, |range| { + range.start.min(row_count)..range.end.min(row_count) + }); + let parts = get_available_parallelism().unwrap_or(1) as u64; + split_row_ranges(rows, parts) + .map(|rows| { + let stream = V2ScanBuilder::new(SESSION.clone(), state.scan_file.clone()) + .with_projection(global.projection.clone()) + .with_some_filter(filter.filter.clone()) + .with_selection(filter.row_selection.clone()) + .with_row_range(rows) + .with_ordered(false) + .with_concurrency(1) + .into_stream()?; + Ok(Split::V2(stream)) + }) + .collect() +} + +/// Cuts `rows` into up to `parts` non-empty, contiguous ranges that tile it. +fn split_row_ranges(rows: Range, parts: u64) -> impl Iterator> { + let len = rows.end.saturating_sub(rows.start); + let parts = parts.max(1); + (0..parts).filter_map(move |part| { + let split = partition_row_range(part..part + 1, parts, len); + (!split.is_empty()).then(|| rows.start + split.start..rows.start + split.end) + }) +} + +/// What a thread's split yields next. +enum SplitItem { + /// The split has no more rows. + Done, + /// The split produced no rows this time but may produce more. + Empty, + Array(ArrayRef), +} + +/// Takes the next array of the thread's split, keeping the split while it has more to yield. +fn next_split_item(local: &mut LocalState) -> VortexResult { + match local.split.take() { + None => Ok(SplitItem::Done), + Some(Split::V1(task)) => Ok(RUNTIME + .block_on(task)? + .map_or(SplitItem::Empty, SplitItem::Array)), + Some(Split::V2(mut stream)) => { + let Some(array) = RUNTIME.block_on(stream.next()).transpose()? else { + return Ok(SplitItem::Done); + }; + local.split = Some(Split::V2(stream)); + Ok(if array.is_empty() { + SplitItem::Empty + } else { + SplitItem::Array(array) + }) + } + } +} + /// Called by all threads under global lock. If this function returns true, /// thread calls reader_scan on this file. If this function returns false, /// duckdb thinks file is exhausted, closes the file, and the first thread to @@ -262,12 +350,11 @@ pub fn reader_scan( } if local.exporter.is_none() { - let Some(split) = local.split.take() else { - return Ok(false); - }; - let Some(array) = RUNTIME.block_on(split)? else { + let array = match next_split_item(local)? { + SplitItem::Done => return Ok(false), // split is filtered - return Ok(true); + SplitItem::Empty => return Ok(true), + SplitItem::Array(array) => array, }; let mut ctx = SESSION.create_execution_ctx(); let array = convert_result(array, &mut ctx)?; @@ -284,12 +371,11 @@ pub fn reader_scan( } fn reader_scan_aggregate(global: &GlobalState, local: &mut LocalState) -> VortexResult { - let Some(split) = local.split.take() else { - return Ok(false); - }; - let Some(array) = RUNTIME.block_on(split)? else { + let array = match next_split_item(local)? { + SplitItem::Done => return Ok(false), // split is filtered - return Ok(true); + SplitItem::Empty => return Ok(true), + SplitItem::Array(array) => array, }; let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index 0201079ba40..d434a78206c 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -7,6 +7,8 @@ use std::ffi::c_char; use std::ffi::c_void; use std::sync::LazyLock; use std::sync::OnceLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use vortex::VortexSessionDefault; use vortex::cloud::Registry; @@ -55,6 +57,20 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +/// Whether scans run on the planner-driven V2 executor. Set `VORTEX_SCAN_V2=1` to enable it. +static SCAN_V2: LazyLock = + LazyLock::new(|| AtomicBool::new(std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "1"))); + +/// Selects the planner-driven V2 scan executor for scans initialized from now on. +pub fn set_scan_v2(enabled: bool) { + SCAN_V2.store(enabled, Ordering::Relaxed); +} + +/// Whether scans initialized now use the planner-driven V2 scan executor. +pub fn scan_v2() -> bool { + SCAN_V2.load(Ordering::Relaxed) +} + // Duckdb's logger requires a *Context as first argument which // would be hard to integrate with tracing::. We use logging for // debugging only anyway, so that's good enough. diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 7c7d172b38c..675b9d3ed51 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -11,6 +11,7 @@ use std::sync::atomic::Ordering; use custom_labels::CURRENT_LABELSET; use futures::future::BoxFuture; +use futures::stream::BoxStream; use itertools::Itertools; use num_traits::AsPrimitive; use parking_lot::Mutex; @@ -140,6 +141,8 @@ pub struct GlobalState { pub projection: BoundExpression, pub filter: Filter, pub file_row_number_column_pos: Option, + /// Whether the scan runs on the planner-driven V2 executor. + pub scan_v2: bool, // Following fields are used only in aggregate scans. pub aggregate_state: Mutex, @@ -156,7 +159,13 @@ pub struct AggregateState { } assert_impl_all!(GlobalState: Send, Sync); -pub type Split = BoxFuture<'static, VortexResult>>; +/// A unit of scan work that one thread takes at a time. +pub enum Split { + /// A V1 scan task yielding one array, or `None` if the filter removed every row. + V1(BoxFuture<'static, VortexResult>>), + /// A V2 scan over a row range of the file, yielding arrays until it ends. + V2(BoxStream<'static, VortexResult>), +} /// field position, accumulator pub type Partials = Vec<(usize, Box)>; @@ -304,6 +313,7 @@ pub fn init_global(init_input: &TableInitInput) -> VortexResult { aggregate_state: Mutex::new(AggregateState::default()), row_count: AtomicU64::new(0), file_row_number_column_pos, + scan_v2: crate::scan_v2(), }) } From fd7a173382a091483dd4ee2178a5f0d0a3afb5f8 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 16:08:58 +0000 Subject: [PATCH 08/22] End DuckDB V2 splits on layout chunk boundaries Cutting a file into one equal row range per core put split ends inside the layout's 65,536-row chunks. On many-core machines a split covered only part of a chunk, so every chunk was read and decoded by up to eight splits; CI's TPC-H SF1 DuckDB run regressed by 15%. Splits now end on the natural split boundaries the V1 scan uses, spread evenly over the cores, so each chunk belongs to one split. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/file_reader.rs | 72 ++++++++++++++++++++++++++------ 1 file changed, 60 insertions(+), 12 deletions(-) diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index c18df0fec4c..f3a1fbd1b2c 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -31,7 +31,6 @@ use vortex::layout::LayoutReaderRef; use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; use vortex::layout::scan::v2::ScanFile; -use vortex::layout::scan::v2::partition_row_range; use vortex::mask::Mask; use vortex::session::SessionExt as _; use vortex_utils::parallelism::get_available_parallelism; @@ -236,17 +235,23 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor } /// Splits the file's scanned rows into one V2 scan per available core, each over a contiguous -/// row range on the executor's row blocks. The splits share the file's scan cache, so the file is -/// planned once. A split's chunks run beside each other only as far as the engine's threads leave -/// cores idle, like DataFusion's partitions. +/// run of the layout's chunks. A split that ended inside a chunk would read and decode the whole +/// chunk for a part of its rows, and so would its neighbour. The splits share the file's scan +/// cache, so the file is planned once. A split's chunks run beside each other only as far as the +/// engine's threads leave cores idle, like DataFusion's partitions. fn v2_splits(state: &FileScanState, global: &GlobalState) -> VortexResult> { let filter = &global.filter; let row_count = state.reader.row_count(); let rows = filter.row_range.as_ref().map_or(0..row_count, |range| { range.start.min(row_count)..range.end.min(row_count) }); + let boundaries = ScanBuilder::new(SESSION.clone(), Arc::clone(&state.reader)) + .with_projection(global.projection.clone()) + .with_some_filter(filter.filter.clone()) + .full_file_splits()?; let parts = get_available_parallelism().unwrap_or(1) as u64; - split_row_ranges(rows, parts) + split_row_ranges(rows, &boundaries, parts) + .into_iter() .map(|rows| { let stream = V2ScanBuilder::new(SESSION.clone(), state.scan_file.clone()) .with_projection(global.projection.clone()) @@ -261,14 +266,31 @@ fn v2_splits(state: &FileScanState, global: &GlobalState) -> VortexResult, parts: u64) -> impl Iterator> { - let len = rows.end.saturating_sub(rows.start); +/// Cuts `rows` into up to `parts` non-empty, contiguous ranges that tile it and end only on the +/// sorted row `boundaries` or at the end of `rows`. Each range ends on the first boundary at or +/// past its even share of the rows. +fn split_row_ranges(rows: Range, boundaries: &[u64], parts: u64) -> Vec> { + let len = u128::from(rows.end.saturating_sub(rows.start)); let parts = parts.max(1); - (0..parts).filter_map(move |part| { - let split = partition_row_range(part..part + 1, parts, len); - (!split.is_empty()).then(|| rows.start + split.start..rows.start + split.end) - }) + let mut ranges = Vec::new(); + let mut start = rows.start; + for part in 1..parts { + let share = u64::try_from(len * u128::from(part) / u128::from(parts)) + .vortex_expect("a share of the rows fits u64"); + let target = rows.start + share; + let next = boundaries.partition_point(|&boundary| boundary < target.max(start + 1)); + match boundaries.get(next) { + Some(&end) if end < rows.end => { + ranges.push(start..end); + start = end; + } + _ => break, + } + } + if start < rows.end { + ranges.push(start..rows.end); + } + ranges } /// What a thread's split yields next. @@ -464,3 +486,29 @@ pub fn footer_get_statistics(footer: &Footer, column: &str) -> Option vortex_panic!(e), } } + +#[cfg(test)] +mod tests { + use std::ops::Range; + + use rstest::rstest; + + use super::split_row_ranges; + + #[rstest] + #[case::even(0..400, &[0, 100, 200, 300, 400], 2, vec![0..200, 200..400])] + #[case::more_parts_than_chunks(0..250, &[0, 100, 200, 250], 8, vec![0..100, 100..200, 200..250])] + #[case::one_part(0..400, &[0, 100, 200, 300, 400], 1, vec![0..400])] + #[case::uneven_chunks(0..1000, &[0, 900, 1000], 4, vec![0..900, 900..1000])] + #[case::row_range_inside_chunks(150..350, &[0, 100, 200, 300, 400], 2, vec![150..300, 300..350])] + #[case::row_range_in_one_chunk(120..180, &[0, 100, 200], 4, vec![120..180])] + #[case::empty(50..50, &[0, 100], 4, vec![])] + fn split_row_ranges_end_on_boundaries( + #[case] rows: Range, + #[case] boundaries: &[u64], + #[case] parts: u64, + #[case] expected: Vec>, + ) { + assert_eq!(split_row_ranges(rows, boundaries, parts), expected); + } +} From 58940c06ec1ee4b2f64bed93e6928c15c39d11cc Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 16:26:33 +0000 Subject: [PATCH 09/22] Drive DuckDB V2 scans on DuckDB's threads with a custom driver Like the DataFusion opener, the DuckDB extension now prepares a V2 scan and composes its planners itself, and drives them with its own driver. Each file is prepared once; the root planner's run is shared by every thread scanning the file. A thread takes the next part the root spawns and runs its chunks inline on a run of its own, waiting only for its reads, so there are no per-split scans, task hops or channels. The root cuts the rows into a few parts per thread on whole 65,536-row blocks, so filtered scans, whose chunks are up to 1,048,576 rows, still spread over every thread and no chunk spans two parts. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/file_reader.rs | 157 +++++--------- vortex-duckdb/src/lib.rs | 1 + vortex-duckdb/src/scan_v2.rs | 313 ++++++++++++++++++++++++++++ vortex-duckdb/src/table_function.rs | 9 +- 4 files changed, 375 insertions(+), 105 deletions(-) create mode 100644 vortex-duckdb/src/scan_v2.rs diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index f3a1fbd1b2c..9f402f0e3ca 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -1,13 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ops::Range; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use futures::FutureExt; -use futures::StreamExt; use object_store::registry::ObjectStoreRegistry; use url::Url; use vortex::array::ArrayRef; @@ -33,7 +32,6 @@ use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; use vortex::layout::scan::v2::ScanFile; use vortex::mask::Mask; use vortex::session::SessionExt as _; -use vortex_utils::parallelism::get_available_parallelism; use crate::REGISTRY; use crate::RUNTIME; @@ -46,6 +44,7 @@ use crate::exporter::ArrayExporter; use crate::exporter::ConversionCache; use crate::projection::Filter; use crate::projection::extract_schema_from_dtype; +use crate::scan_v2::FileScan; use crate::table_function::BindState; use crate::table_function::GlobalState; use crate::table_function::LocalState; @@ -120,9 +119,11 @@ pub struct OpenFileReader { /// The state of an open file that the threads scanning it share. pub struct FileScanState { reader: LayoutReaderRef, - /// The file as the V2 executor reads it, with a scan cache that every split of the file - /// shares. + /// The file as the V2 executor reads it. scan_file: ScanFile, + /// The file's V2 scan, which every thread scanning the file drives. Set when the file is + /// initialized for a V2 scan. + v2: OnceLock, cache: ConversionCache, } @@ -137,6 +138,7 @@ impl OpenFileReader { scan_state: Some(Arc::new(FileScanState { reader: file.layout_reader()?, scan_file, + v2: OnceLock::new(), cache: ConversionCache::default(), })), splits: vec![], @@ -209,23 +211,30 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor return Ok(true); } - // Getting splits is non-trivial work so we prefer doing it here under file - // lock and not in reader_try_initialize_scan under global lock. - let mut splits = if global.scan_v2 { - v2_splits(file.scan_state(), global)? - } else { - let reader = Arc::clone(file.reader()); - let filter = &global.filter; - let builder = ScanBuilder::new(SESSION.clone(), reader) - .with_projection(global.projection.clone()) - .with_some_filter(filter.filter.clone()) - .with_selection(filter.row_selection.clone()); - let scan = builder.prepare()?; - scan.execute(filter.row_range.clone())? - .into_iter() - .map(Split::V1) - .collect() - }; + // Getting splits, or planning a V2 scan, is non-trivial work so we prefer + // doing it here under file lock and not in reader_try_initialize_scan under + // global lock. + if global.scan_v2 { + let state = file.scan_state(); + let scan = FileScan::try_new(v2_builder(state, global))?; + if state.v2.set(scan).is_err() { + vortex_panic!("file initialized twice"); + } + return Ok(false); + } + + let reader = Arc::clone(file.reader()); + let filter = &global.filter; + let builder = ScanBuilder::new(SESSION.clone(), reader) + .with_projection(global.projection.clone()) + .with_some_filter(filter.filter.clone()) + .with_selection(filter.row_selection.clone()); + let scan = builder.prepare()?; + let mut splits: Vec = scan + .execute(filter.row_range.clone())? + .into_iter() + .map(Split::V1) + .collect(); // threads take last element of file.splits so we need to reverse splits.reverse(); @@ -234,63 +243,18 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor Ok(false) } -/// Splits the file's scanned rows into one V2 scan per available core, each over a contiguous -/// run of the layout's chunks. A split that ended inside a chunk would read and decode the whole -/// chunk for a part of its rows, and so would its neighbour. The splits share the file's scan -/// cache, so the file is planned once. A split's chunks run beside each other only as far as the -/// engine's threads leave cores idle, like DataFusion's partitions. -fn v2_splits(state: &FileScanState, global: &GlobalState) -> VortexResult> { +/// The V2 scan of the file's rows that the scan's row range and selection keep. +fn v2_builder(state: &FileScanState, global: &GlobalState) -> V2ScanBuilder { let filter = &global.filter; let row_count = state.reader.row_count(); let rows = filter.row_range.as_ref().map_or(0..row_count, |range| { range.start.min(row_count)..range.end.min(row_count) }); - let boundaries = ScanBuilder::new(SESSION.clone(), Arc::clone(&state.reader)) + V2ScanBuilder::new(SESSION.clone(), state.scan_file.clone()) .with_projection(global.projection.clone()) .with_some_filter(filter.filter.clone()) - .full_file_splits()?; - let parts = get_available_parallelism().unwrap_or(1) as u64; - split_row_ranges(rows, &boundaries, parts) - .into_iter() - .map(|rows| { - let stream = V2ScanBuilder::new(SESSION.clone(), state.scan_file.clone()) - .with_projection(global.projection.clone()) - .with_some_filter(filter.filter.clone()) - .with_selection(filter.row_selection.clone()) - .with_row_range(rows) - .with_ordered(false) - .with_concurrency(1) - .into_stream()?; - Ok(Split::V2(stream)) - }) - .collect() -} - -/// Cuts `rows` into up to `parts` non-empty, contiguous ranges that tile it and end only on the -/// sorted row `boundaries` or at the end of `rows`. Each range ends on the first boundary at or -/// past its even share of the rows. -fn split_row_ranges(rows: Range, boundaries: &[u64], parts: u64) -> Vec> { - let len = u128::from(rows.end.saturating_sub(rows.start)); - let parts = parts.max(1); - let mut ranges = Vec::new(); - let mut start = rows.start; - for part in 1..parts { - let share = u64::try_from(len * u128::from(part) / u128::from(parts)) - .vortex_expect("a share of the rows fits u64"); - let target = rows.start + share; - let next = boundaries.partition_point(|&boundary| boundary < target.max(start + 1)); - match boundaries.get(next) { - Some(&end) if end < rows.end => { - ranges.push(start..end); - start = end; - } - _ => break, - } - } - if start < rows.end { - ranges.push(start..rows.end); - } - ranges + .with_selection(filter.row_selection.clone()) + .with_row_range(rows) } /// What a thread's split yields next. @@ -309,11 +273,16 @@ fn next_split_item(local: &mut LocalState) -> VortexResult { Some(Split::V1(task)) => Ok(RUNTIME .block_on(task)? .map_or(SplitItem::Empty, SplitItem::Array)), - Some(Split::V2(mut stream)) => { - let Some(array) = RUNTIME.block_on(stream.next()).transpose()? else { + Some(Split::V2) => { + let file = local.file.as_ref().vortex_expect("split without a file"); + let scan = file + .v2 + .get() + .vortex_expect("V2 split of a file without a V2 scan"); + let Some(array) = scan.next_array(&mut local.run)? else { return Ok(SplitItem::Done); }; - local.split = Some(Split::V2(stream)); + local.split = Some(Split::V2); Ok(if array.is_empty() { SplitItem::Empty } else { @@ -332,7 +301,14 @@ fn next_split_item(local: &mut LocalState) -> VortexResult { /// the global lock. It only moves references between the file and the local /// state, and reader_scan drops the retired ones. pub fn reader_try_initialize_scan(file: &mut OpenFileReader, local: &mut LocalState) -> bool { - let Some(split) = file.splits.pop() else { + // A V2 scan's threads take chunks from the file's scan until it has spawned + // them all, so the file has work for another thread until then. + let v2 = file + .scan_state + .as_ref() + .and_then(|state| state.v2.get()) + .map(|scan| (!scan.is_exhausted()).then_some(Split::V2)); + let Some(split) = v2.unwrap_or_else(|| file.splits.pop()) else { // DuckDB releases the file after this call, under the global lock. // Give the reference of the file to this thread so it drops outside. local.retired_files.extend(file.scan_state.take()); @@ -445,6 +421,9 @@ pub fn reader_get_statistics( /// Called from a separate thread (not related to threads for Vortex /// table function) under global lock. pub fn reader_get_progress_in_file(file: &OpenFileReader) -> f64 { + if let Some(scan) = file.scan_state.as_ref().and_then(|state| state.v2.get()) { + return scan.progress(); + } let total = file.total_splits; let left = file.splits.len(); let denom = total + (total == 0) as usize; @@ -486,29 +465,3 @@ pub fn footer_get_statistics(footer: &Footer, column: &str) -> Option vortex_panic!(e), } } - -#[cfg(test)] -mod tests { - use std::ops::Range; - - use rstest::rstest; - - use super::split_row_ranges; - - #[rstest] - #[case::even(0..400, &[0, 100, 200, 300, 400], 2, vec![0..200, 200..400])] - #[case::more_parts_than_chunks(0..250, &[0, 100, 200, 250], 8, vec![0..100, 100..200, 200..250])] - #[case::one_part(0..400, &[0, 100, 200, 300, 400], 1, vec![0..400])] - #[case::uneven_chunks(0..1000, &[0, 900, 1000], 4, vec![0..900, 900..1000])] - #[case::row_range_inside_chunks(150..350, &[0, 100, 200, 300, 400], 2, vec![150..300, 300..350])] - #[case::row_range_in_one_chunk(120..180, &[0, 100, 200], 4, vec![120..180])] - #[case::empty(50..50, &[0, 100], 4, vec![])] - fn split_row_ranges_end_on_boundaries( - #[case] rows: Range, - #[case] boundaries: &[u64], - #[case] parts: u64, - #[case] expected: Vec>, - ) { - assert_eq!(split_row_ranges(rows, boundaries, parts), expected); - } -} diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index d434a78206c..2330a240000 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -31,6 +31,7 @@ mod exporter; mod ffi; mod file_reader; mod projection; +mod scan_v2; mod table_function; #[rustfmt::skip] diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs new file mode 100644 index 00000000000..cee78cdc441 --- /dev/null +++ b/vortex-duckdb/src/scan_v2.rs @@ -0,0 +1,313 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Drives a V2 scan for DuckDB. +//! +//! A planner says what it needs: compute, or reads. This driver decides where each runs, and +//! DuckDB's own threads run all of it. Each file is prepared once and its planners are composed +//! over the prepared scan: with a filter, its pruning plans over the scanned rows first, +//! continuing into the chunks of the rows they keep; without, the chunks directly. The root +//! planner's run is shared by every thread scanning the file. A thread takes the next chunk the +//! root spawns and runs it on a run of its own, inline, waiting only for its reads. Threads keep +//! taking chunks until the root is done, so the file's work spreads over as many threads as +//! DuckDB gives the scan, however its rows divide. +//! +//! The root cuts the scanned rows into parts of whole 65,536-row blocks, a few per thread, and +//! spawns a chunk planner per part. A filtered scan's chunks are up to 1,048,576 rows, so without +//! parts a file of a few million rows would keep only a few threads busy. + +use std::future; +use std::ops::Range; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use parking_lot::Mutex; +use vortex::array::ArrayRef; +use vortex::error::VortexExpect; +use vortex::error::VortexResult; +use vortex::io::request::IoConsumer; +use vortex::io::request::IoRequestId; +use vortex::io::request::IoResult; +use vortex::io::request::IoService; +use vortex::io::runtime::BlockingRuntime as _; +use vortex::layout::scan::v2::ChunkPlanner; +use vortex::layout::scan::v2::PreparedScan; +use vortex::layout::scan::v2::PrunePlanner; +use vortex::layout::scan::v2::ScanBuilder; +use vortex::mask::Mask; +use vortex::scan::planning::driver::Progress; +use vortex::scan::planning::driver::Run; +use vortex::scan::planning::planner::Planner; +use vortex::scan::planning::planner::PlannerOutput; +use vortex::scan::planning::planner::State; +use vortex::scan::planning::planner::WorkScope; +use vortex::scan::selection::Selection; +use vortex_utils::parallelism::get_available_parallelism; + +use crate::RUNTIME; + +/// The chunks a thread runs beside each other on its run, so one chunk's reads overlap another's +/// compute. +const CHUNKS_PER_THREAD: usize = 1; +/// The parts of a file's rows per thread, so threads that finish early find more work. +const PARTS_PER_THREAD: u64 = 4; +/// The row blocks parts are cut on: the rows of a projection chunk, and of the writer's chunks. +const PART_ROW_BLOCK: u64 = 65_536; + +/// One file's V2 scan, shared by the threads scanning it. +pub(crate) struct FileScan { + scan: Arc, + io: Arc, + /// The run of the root planner, locked while a thread takes the chunk it spawns next. + root: Mutex, + /// Set once the root has spawned its last chunk. + exhausted: AtomicBool, + /// The rows of the chunks taken so far, for DuckDB's progress bar. + rows_taken: AtomicU64, +} + +/// What a thread takes from the root next. +enum Work { + /// A chunk to run on the thread's run. + Chunk(WorkScope, Box), + /// A batch the root produced itself. + Batch(ArrayRef), +} + +impl FileScan { + /// Prepares `builder` and composes its planners. Lowering and planning touch no segment. + pub(crate) fn try_new(builder: ScanBuilder) -> VortexResult { + let scan = Arc::new(builder.prepare()?); + let io = scan.io(); + let rows = scan.rows().clone(); + let mut root = Run::new(); + root.admit( + compose(&scan), + WorkScope { + file_ordinal: 0, + rows, + }, + io.session(), + ); + Ok(Self { + scan, + io, + root: Mutex::new(root), + exhausted: AtomicBool::new(false), + rows_taken: AtomicU64::new(0), + }) + } + + /// Whether the root has spawned its last chunk, so a thread joining the scan finds no work. + pub(crate) fn is_exhausted(&self) -> bool { + self.exhausted.load(Ordering::Acquire) + } + + /// The percentage of the scanned rows whose chunks threads have taken. + pub(crate) fn progress(&self) -> f64 { + let rows = self.scan.rows(); + let total = rows.end - rows.start; + if total == 0 || self.exhausted.load(Ordering::Relaxed) { + return 100.0; + } + 100.0 * self.rows_taken.load(Ordering::Relaxed) as f64 / total as f64 + } + + /// The next array of this file for the thread whose run is `run`, or `None` once the file has + /// no more work for the thread. Keeps the thread's run fed with chunks from the root, and + /// waits for reads on the calling thread. + pub(crate) fn next_array(&self, run: &mut Run) -> VortexResult> { + loop { + while run.live_roots() < CHUNKS_PER_THREAD { + match self.next_work()? { + Some(Work::Chunk(scope, chunk)) => { + run.admit(chunk, scope, self.io.session()); + } + Some(Work::Batch(array)) => return Ok(Some(array)), + None => break, + } + } + match run.advance()? { + Progress::Batch(batch) => return Ok(Some(batch.array)), + Progress::Spawn(scope, child) => { + run.admit(child, scope, self.io.session()); + } + Progress::RootDone(_) => {} + Progress::Idle => return Ok(None), + Progress::Waiting => { + RUNTIME.block_on(future::poll_fn(|cx| run.poll_completion(cx)))?; + } + } + } + } + + /// Advances the root until it spawns a chunk or produces a batch, or `None` once it is done. + fn next_work(&self) -> VortexResult> { + if self.exhausted.load(Ordering::Acquire) { + return Ok(None); + } + let mut root = self.root.lock(); + loop { + if self.exhausted.load(Ordering::Acquire) { + return Ok(None); + } + match root.advance()? { + Progress::Batch(batch) => return Ok(Some(Work::Batch(batch.array))), + Progress::Spawn(scope, chunk) => { + self.rows_taken + .fetch_add(scope.rows.end - scope.rows.start, Ordering::Relaxed); + return Ok(Some(Work::Chunk(scope, chunk))); + } + Progress::RootDone(_) if root.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => { + self.exhausted.store(true, Ordering::Release); + return Ok(None); + } + // The root waits for the pruning plans' reads. The threads waiting on the lock + // would have nothing to run until they arrive. + Progress::Waiting => { + RUNTIME.block_on(future::poll_fn(|cx| root.poll_completion(cx)))?; + } + } + } + } +} + +/// The planners of a scan, composed: with a filter, its pruning plans first over the scan's rows, +/// continuing into the parts of the rows they keep; without, the parts directly. +fn compose(scan: &Arc) -> Box { + let rows = scan.rows().clone(); + let parts = { + let scan = Arc::clone(scan); + let rows = rows.clone(); + move |kept: Option| -> Box { + let parallelism = get_available_parallelism().unwrap_or(1) as u64; + Box::new(PartPlanner::new( + scan, + rows, + kept, + parallelism * PARTS_PER_THREAD, + )) + } + }; + if scan.filter().is_some() { + Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(parts))) + } else { + parts(None) + } +} + +/// Spawns a chunk planner per part of a scan's rows, in row order. Parts end on whole +/// [`PART_ROW_BLOCK`]s, so no chunk of the scan spans two parts. +struct PartPlanner { + scan: Arc, + selection: Selection, + /// The first row of `kept`. + first_row: u64, + /// The rows zone pruning kept, from `first_row`, or `None` when nothing was pruned. + kept: Option, + parts: std::vec::IntoIter>, +} + +impl PartPlanner { + fn new(scan: Arc, rows: Range, kept: Option, parts: u64) -> Self { + Self { + selection: scan.selection().clone(), + scan, + first_row: rows.start, + kept, + parts: part_ranges(rows, parts).into_iter(), + } + } +} + +impl IoConsumer for PartPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for PartPlanner { + fn state(&self) -> State { + if self.parts.as_slice().is_empty() { + State::Done + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + for part in self.parts.by_ref() { + let kept = match &self.kept { + Some(kept) => { + let offset = |row: u64| { + usize::try_from(row - self.first_row).vortex_expect("a part fits usize") + }; + let kept = kept.slice(offset(part.start)..offset(part.end)); + if kept.all_false() { + continue; + } + Some(kept) + } + None => None, + }; + let chunks = ChunkPlanner::new( + Arc::clone(&self.scan), + part.clone(), + self.selection.clone(), + kept, + ); + let scope = WorkScope { + file_ordinal: 0, + rows: part, + }; + return Ok(PlannerOutput::Spawn(scope, Box::new(chunks))); + } + Ok(PlannerOutput::Done) + } +} + +/// Cuts `rows` into up to `parts` non-empty ranges of whole [`PART_ROW_BLOCK`]s of the file's +/// row domain that tile it, clipped to `rows`. +fn part_ranges(rows: Range, parts: u64) -> Vec> { + if rows.is_empty() { + return Vec::new(); + } + let first_block = rows.start / PART_ROW_BLOCK; + let blocks = rows.end.div_ceil(PART_ROW_BLOCK) - first_block; + let blocks_per_part = blocks.div_ceil(parts.max(1)); + (0..blocks.div_ceil(blocks_per_part)) + .map(|part| { + let start = (first_block + part * blocks_per_part) * PART_ROW_BLOCK; + let end = start + blocks_per_part * PART_ROW_BLOCK; + start.max(rows.start)..end.min(rows.end) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use std::ops::Range; + + use rstest::rstest; + + use super::PART_ROW_BLOCK; + use super::part_ranges; + + const B: u64 = PART_ROW_BLOCK; + + #[rstest] + #[case::even(0..4 * B, 2, vec![0..2 * B, 2 * B..4 * B])] + #[case::ragged_end(0..3 * B + 10, 2, vec![0..2 * B, 2 * B..3 * B + 10])] + #[case::more_parts_than_blocks(0..2 * B, 8, vec![0..B, B..2 * B])] + #[case::small(0..100, 4, vec![0..100])] + #[case::unaligned_start(B / 2..3 * B, 2, vec![B / 2..2 * B, 2 * B..3 * B])] + #[case::empty(5..5, 4, vec![])] + fn parts_end_on_blocks( + #[case] rows: Range, + #[case] parts: u64, + #[case] expected: Vec>, + ) { + assert_eq!(part_ranges(rows, parts), expected); + } +} diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 675b9d3ed51..df670ccd558 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -11,7 +11,6 @@ use std::sync::atomic::Ordering; use custom_labels::CURRENT_LABELSET; use futures::future::BoxFuture; -use futures::stream::BoxStream; use itertools::Itertools; use num_traits::AsPrimitive; use parking_lot::Mutex; @@ -33,6 +32,7 @@ use vortex::expr::BoundExpression; use vortex::expr::Expression; use vortex::metrics::tracing::get_global_labels; use vortex::scalar::Scalar; +use vortex::scan::planning::driver::Run; use vortex_utils::aliases::hash_map::HashMap; use crate::convert::PushedAggregate; @@ -163,8 +163,8 @@ assert_impl_all!(GlobalState: Send, Sync); pub enum Split { /// A V1 scan task yielding one array, or `None` if the filter removed every row. V1(BoxFuture<'static, VortexResult>>), - /// A V2 scan over a row range of the file, yielding arrays until it ends. - V2(BoxStream<'static, VortexResult>), + /// Chunks of the file's V2 scan, taken by the thread until the scan has none left. + V2, } /// field position, accumulator @@ -179,6 +179,8 @@ pub struct LocalState { /// Scan states that this thread no longer uses. They are dropped outside /// DuckDB's global multi-file lock; see `OpenFileReader`. pub(crate) retired_files: Vec>, + /// The planning run on which this thread runs the chunks of V2 scans. + pub(crate) run: Run, /// Empty for non-aggregate scan pub partials: Partials, @@ -380,6 +382,7 @@ pub fn init_local(bind_data: &BindState, global: &GlobalState) -> LocalState { split: None, file: None, retired_files: Vec::new(), + run: Run::new(), finished: false, } } From 43671cc9325df3fd722347ef200ff71a21d24899 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 18:27:46 +0000 Subject: [PATCH 10/22] Default duckdb-bench to the V2 scan for TPC-DS and ClickBench The CI SQL benchmarks then compare DuckDB's V2 scan with the V1 baseline on those suites too, as for TPC-H. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- benchmarks/duckdb-bench/src/main.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index a047e275b24..4b5e2cfbb8c 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -103,7 +103,7 @@ struct Args { )] reuse: bool, - /// 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, } @@ -117,8 +117,10 @@ fn main() -> anyhow::Result<()> { 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") }); vortex_duckdb::set_scan_v2(scan_v2); From 8a4cefe9c60161bcb12f485c7de1ba431cf00f42 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 19:14:18 +0000 Subject: [PATCH 11/22] Cut DuckDB V2 scan chunks on the file's chunk boundaries The DuckDB V2 driver cut parts on 64K-row blocks and the chunk planner cut chunks on a fixed grid. Files whose chunks are not 64K rows, such as ClickBench's 87K and 95K-row chunks, had nearly every chunk span two scan chunks, so it was decoded twice. Take the boundaries V1 splits on (the referenced columns' chunk boundaries, with wide spans subdivided) and cut chunks only there, merging spans up to the scan's chunk size. Parts group whole chunks. ChunkPlanner gains a with_chunks constructor for drivers that cut the rows themselves. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/file_reader.rs | 14 ++- vortex-duckdb/src/scan_v2.rs | 144 +++++++++++++++------- vortex-layout/src/scan/execution/chunk.rs | 22 +++- vortex-layout/src/scan/execution/mod.rs | 5 + 4 files changed, 136 insertions(+), 49 deletions(-) diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index 9f402f0e3ca..96c62cc43a7 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -28,6 +28,8 @@ 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::scan_builder::referenced_field_masks; +use vortex::layout::scan::split_by::SplitBy; use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; use vortex::layout::scan::v2::ScanFile; use vortex::mask::Mask; @@ -216,7 +218,7 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor // global lock. if global.scan_v2 { let state = file.scan_state(); - let scan = FileScan::try_new(v2_builder(state, global))?; + let scan = FileScan::try_new(v2_builder(state, global), &v2_boundaries(state, global)?)?; if state.v2.set(scan).is_err() { vortex_panic!("file initialized twice"); } @@ -257,6 +259,16 @@ fn v2_builder(state: &FileScanState, global: &GlobalState) -> V2ScanBuilder { .with_row_range(rows) } +/// The rows the file's chunks of the columns the scan reads end at, the boundaries V1 splits on. +fn v2_boundaries(state: &FileScanState, global: &GlobalState) -> VortexResult> { + let field_mask = referenced_field_masks(&global.projection, global.filter.filter.as_ref())?; + SplitBy::default().splits( + state.reader.as_ref(), + &(0..state.reader.row_count()), + &field_mask, + ) +} + /// What a thread's split yields next. enum SplitItem { /// The split has no more rows. diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index cee78cdc441..79d021ee91d 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -12,9 +12,11 @@ //! taking chunks until the root is done, so the file's work spreads over as many threads as //! DuckDB gives the scan, however its rows divide. //! -//! The root cuts the scanned rows into parts of whole 65,536-row blocks, a few per thread, and -//! spawns a chunk planner per part. A filtered scan's chunks are up to 1,048,576 rows, so without -//! parts a file of a few million rows would keep only a few threads busy. +//! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk +//! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up +//! to the scan's chunk size. The root groups the chunks into parts, a few per thread, and spawns a +//! chunk planner per part. A filtered scan's chunks are up to 1,048,576 rows, so without parts a +//! file of a few million rows would keep only a few threads busy. use std::future; use std::ops::Range; @@ -52,9 +54,7 @@ use crate::RUNTIME; /// compute. const CHUNKS_PER_THREAD: usize = 1; /// The parts of a file's rows per thread, so threads that finish early find more work. -const PARTS_PER_THREAD: u64 = 4; -/// The row blocks parts are cut on: the rows of a projection chunk, and of the writer's chunks. -const PART_ROW_BLOCK: u64 = 65_536; +const PARTS_PER_THREAD: usize = 4; /// One file's V2 scan, shared by the threads scanning it. pub(crate) struct FileScan { @@ -77,14 +77,17 @@ enum Work { } impl FileScan { - /// Prepares `builder` and composes its planners. Lowering and planning touch no segment. - pub(crate) fn try_new(builder: ScanBuilder) -> VortexResult { + /// Prepares `builder` and composes its planners, cutting chunks on `boundaries`, the sorted + /// rows the file's chunks of the scanned columns end at. Lowering and planning touch no + /// segment. + pub(crate) fn try_new(builder: ScanBuilder, boundaries: &[u64]) -> VortexResult { let scan = Arc::new(builder.prepare()?); let io = scan.io(); let rows = scan.rows().clone(); + let chunks = chunk_ranges(rows.clone(), boundaries, scan.chunk_rows() as u64); let mut root = Run::new(); root.admit( - compose(&scan), + compose(&scan, chunks), WorkScope { file_ordinal: 0, rows, @@ -177,18 +180,18 @@ impl FileScan { /// The planners of a scan, composed: with a filter, its pruning plans first over the scan's rows, /// continuing into the parts of the rows they keep; without, the parts directly. -fn compose(scan: &Arc) -> Box { +fn compose(scan: &Arc, chunks: Vec>) -> Box { let rows = scan.rows().clone(); let parts = { let scan = Arc::clone(scan); let rows = rows.clone(); move |kept: Option| -> Box { - let parallelism = get_available_parallelism().unwrap_or(1) as u64; + let parallelism = get_available_parallelism().unwrap_or(1); Box::new(PartPlanner::new( scan, rows, kept, - parallelism * PARTS_PER_THREAD, + parts(chunks, parallelism * PARTS_PER_THREAD), )) } }; @@ -199,8 +202,8 @@ fn compose(scan: &Arc) -> Box { } } -/// Spawns a chunk planner per part of a scan's rows, in row order. Parts end on whole -/// [`PART_ROW_BLOCK`]s, so no chunk of the scan spans two parts. +/// Spawns a chunk planner per part of a scan's rows, in row order. A part is a run of whole +/// chunks. struct PartPlanner { scan: Arc, selection: Selection, @@ -208,17 +211,23 @@ struct PartPlanner { first_row: u64, /// The rows zone pruning kept, from `first_row`, or `None` when nothing was pruned. kept: Option, - parts: std::vec::IntoIter>, + /// The chunks of each part, which tile the scan's rows. + parts: std::vec::IntoIter>>, } impl PartPlanner { - fn new(scan: Arc, rows: Range, kept: Option, parts: u64) -> Self { + fn new( + scan: Arc, + rows: Range, + kept: Option, + parts: Vec>>, + ) -> Self { Self { selection: scan.selection().clone(), scan, first_row: rows.start, kept, - parts: part_ranges(rows, parts).into_iter(), + parts: parts.into_iter(), } } } @@ -237,7 +246,11 @@ impl Planner for PartPlanner { } fn compute(&mut self) -> VortexResult { - for part in self.parts.by_ref() { + for chunks in self.parts.by_ref() { + let (Some(first), Some(last)) = (chunks.first(), chunks.last()) else { + continue; + }; + let part = first.start..last.end; let kept = match &self.kept { Some(kept) => { let offset = |row: u64| { @@ -251,9 +264,10 @@ impl Planner for PartPlanner { } None => None, }; - let chunks = ChunkPlanner::new( + let chunks = ChunkPlanner::with_chunks( Arc::clone(&self.scan), part.clone(), + chunks, self.selection.clone(), kept, ); @@ -267,22 +281,44 @@ impl Planner for PartPlanner { } } -/// Cuts `rows` into up to `parts` non-empty ranges of whole [`PART_ROW_BLOCK`]s of the file's -/// row domain that tile it, clipped to `rows`. -fn part_ranges(rows: Range, parts: u64) -> Vec> { - if rows.is_empty() { +/// Cuts `rows` into chunks that end only on `boundaries` or at the end of `rows`, merging the +/// spans between boundaries into chunks of at most `max_rows` rows. A span wider than `max_rows` +/// stays one chunk, so no chunk of the file is split. +fn chunk_ranges(rows: Range, boundaries: &[u64], max_rows: u64) -> Vec> { + let mut chunks = Vec::new(); + let mut start = rows.start; + let mut end = rows.start; + let ends = boundaries + .iter() + .copied() + .filter(|&row| row > rows.start && row < rows.end) + .chain((!rows.is_empty()).then_some(rows.end)); + for row in ends { + if row - start > max_rows && end > start { + chunks.push(start..end); + start = end; + } + end = row; + } + if end > start { + chunks.push(start..end); + } + chunks +} + +/// Groups `chunks` into up to `parts` runs of consecutive chunks, of as even a count as they +/// divide into. +fn parts(chunks: Vec>, parts: usize) -> Vec>> { + if chunks.is_empty() { return Vec::new(); } - let first_block = rows.start / PART_ROW_BLOCK; - let blocks = rows.end.div_ceil(PART_ROW_BLOCK) - first_block; - let blocks_per_part = blocks.div_ceil(parts.max(1)); - (0..blocks.div_ceil(blocks_per_part)) - .map(|part| { - let start = (first_block + part * blocks_per_part) * PART_ROW_BLOCK; - let end = start + blocks_per_part * PART_ROW_BLOCK; - start.max(rows.start)..end.min(rows.end) - }) - .collect() + let per_part = chunks.len().div_ceil(parts.max(1)); + let mut chunks = chunks.into_iter().peekable(); + let mut grouped = Vec::new(); + while chunks.peek().is_some() { + grouped.push(chunks.by_ref().take(per_part).collect()); + } + grouped } #[cfg(test)] @@ -291,23 +327,39 @@ mod tests { use rstest::rstest; - use super::PART_ROW_BLOCK; - use super::part_ranges; - - const B: u64 = PART_ROW_BLOCK; + use super::chunk_ranges; + use super::parts; #[rstest] - #[case::even(0..4 * B, 2, vec![0..2 * B, 2 * B..4 * B])] - #[case::ragged_end(0..3 * B + 10, 2, vec![0..2 * B, 2 * B..3 * B + 10])] - #[case::more_parts_than_blocks(0..2 * B, 8, vec![0..B, B..2 * B])] - #[case::small(0..100, 4, vec![0..100])] - #[case::unaligned_start(B / 2..3 * B, 2, vec![B / 2..2 * B, 2 * B..3 * B])] - #[case::empty(5..5, 4, vec![])] - fn parts_end_on_blocks( + #[case::merged(0..400, &[100, 200, 300, 400], 200, vec![0..200, 200..400])] + #[case::wide_spans(0..300, &[0, 87, 182, 300], 64, vec![0..87, 87..182, 182..300])] + #[case::clipped(50..250, &[100, 200, 300], 100, vec![50..100, 100..200, 200..250])] + #[case::no_boundaries(0..10, &[], 4, vec![0..10])] + #[case::empty(5..5, &[5], 4, vec![])] + fn chunks_end_on_boundaries( #[case] rows: Range, - #[case] parts: u64, + #[case] boundaries: &[u64], + #[case] max_rows: u64, #[case] expected: Vec>, ) { - assert_eq!(part_ranges(rows, parts), expected); + assert_eq!(chunk_ranges(rows, boundaries, max_rows), expected); + } + + #[rstest] + #[case::even(4, 2, vec![2, 2])] + #[case::ragged(5, 2, vec![3, 2])] + #[case::more_parts_than_chunks(2, 8, vec![1, 1])] + #[case::empty(0, 4, vec![])] + fn parts_group_whole_chunks( + #[case] chunks: u64, + #[case] count: usize, + #[case] expected: Vec, + ) { + let grouped = parts((0..chunks).map(|i| i..i + 1).collect(), count); + assert_eq!(grouped.iter().map(Vec::len).collect::>(), expected); + assert_eq!( + grouped.concat(), + (0..chunks).map(|i| i..i + 1).collect::>() + ); } } diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 382425d9bb0..d8c423d7eb3 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -79,11 +79,29 @@ impl ChunkPlanner { selection: Selection, kept: Option, ) -> Self { - let chunks = row_chunks(rows.clone(), scan.chunk_rows).into_iter(); + let chunks = row_chunks(rows.clone(), scan.chunk_rows); + Self::with_chunks(scan, rows, chunks, selection, kept) + } + + /// The planner of `chunks`, which tile `rows` in row order, for a driver that cuts the rows + /// itself, such as on the file's own chunk boundaries. + pub fn with_chunks( + scan: Arc, + rows: Range, + chunks: Vec>, + selection: Selection, + kept: Option, + ) -> Self { + debug_assert!( + chunks.first().is_none_or(|chunk| chunk.start == rows.start) + && chunks.last().is_none_or(|chunk| chunk.end == rows.end) + && chunks.windows(2).all(|pair| pair[0].end == pair[1].start), + "chunks must tile the rows in order" + ); Self { scan, rows, - chunks, + chunks: chunks.into_iter(), selection, kept, } diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 0185f61d6a9..497e2684904 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -118,6 +118,11 @@ impl PreparedScan { self.filter.as_deref() } + /// The most rows of one chunk a [`ChunkPlanner`] cuts the scan's rows into. + pub fn chunk_rows(&self) -> usize { + self.chunk_rows + } + /// The service a driver takes IO sessions from: every read a planner of the scan requests is /// served from it by segment. pub fn io(&self) -> Arc { From 8a464a97c2259ca5af843842fde42b9b7db71c45 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 19:28:52 +0000 Subject: [PATCH 12/22] Keep a few V2 chunks per thread when merging file chunks Merging the spans between file chunk boundaries up to the scan's chunk size gave filtered scans 1M-row chunks, so a 6M-row TPC-H lineitem file became six chunks and kept six threads busy (q13 4.7x, q01 2.9x slower on SF1). Merge no further than leaves parallelism * PARTS_PER_THREAD chunks; a file chunk is still never split. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/scan_v2.rs | 41 ++++++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index 79d021ee91d..2ff1668fa3a 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -14,9 +14,10 @@ //! //! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk //! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up -//! to the scan's chunk size. The root groups the chunks into parts, a few per thread, and spawns a -//! chunk planner per part. A filtered scan's chunks are up to 1,048,576 rows, so without parts a -//! file of a few million rows would keep only a few threads busy. +//! to the scan's chunk size, but no further than leaves a few chunks per thread: a filtered scan's +//! chunks are up to 1,048,576 rows, so a file of a few million rows would otherwise keep only a few +//! threads busy. The root groups the chunks into parts, a few per thread, and spawns a chunk +//! planner per part. use std::future; use std::ops::Range; @@ -84,10 +85,12 @@ impl FileScan { let scan = Arc::new(builder.prepare()?); let io = scan.io(); let rows = scan.rows().clone(); - let chunks = chunk_ranges(rows.clone(), boundaries, scan.chunk_rows() as u64); + let parts = get_available_parallelism().unwrap_or(1) * PARTS_PER_THREAD; + let max_rows = (scan.chunk_rows() as u64).min(merged_rows(&rows, parts)); + let chunks = chunk_ranges(rows.clone(), boundaries, max_rows); let mut root = Run::new(); root.admit( - compose(&scan, chunks), + compose(&scan, chunks, parts), WorkScope { file_ordinal: 0, rows, @@ -180,19 +183,13 @@ impl FileScan { /// The planners of a scan, composed: with a filter, its pruning plans first over the scan's rows, /// continuing into the parts of the rows they keep; without, the parts directly. -fn compose(scan: &Arc, chunks: Vec>) -> Box { +fn compose(scan: &Arc, chunks: Vec>, count: usize) -> Box { let rows = scan.rows().clone(); let parts = { let scan = Arc::clone(scan); let rows = rows.clone(); move |kept: Option| -> Box { - let parallelism = get_available_parallelism().unwrap_or(1); - Box::new(PartPlanner::new( - scan, - rows, - kept, - parts(chunks, parallelism * PARTS_PER_THREAD), - )) + Box::new(PartPlanner::new(scan, rows, kept, parts(chunks, count))) } }; if scan.filter().is_some() { @@ -281,6 +278,11 @@ impl Planner for PartPlanner { } } +/// The most rows spans merge into while `rows` still divides into `parts` chunks. +fn merged_rows(rows: &Range, parts: usize) -> u64 { + (rows.end - rows.start).div_ceil(parts.max(1) as u64).max(1) +} + /// Cuts `rows` into chunks that end only on `boundaries` or at the end of `rows`, merging the /// spans between boundaries into chunks of at most `max_rows` rows. A span wider than `max_rows` /// stays one chunk, so no chunk of the file is split. @@ -328,6 +330,7 @@ mod tests { use rstest::rstest; use super::chunk_ranges; + use super::merged_rows; use super::parts; #[rstest] @@ -345,6 +348,18 @@ mod tests { assert_eq!(chunk_ranges(rows, boundaries, max_rows), expected); } + #[rstest] + #[case::few_parts(0..1_000, 4, 250)] + #[case::many_parts(0..10, 384, 1)] + #[case::empty(5..5, 4, 1)] + fn merging_leaves_chunks_for_every_part( + #[case] rows: Range, + #[case] parts: usize, + #[case] expected: u64, + ) { + assert_eq!(merged_rows(&rows, parts), expected); + } + #[rstest] #[case::even(4, 2, vec![2, 2])] #[case::ragged(5, 2, vec![3, 2])] From 527684bdb3c3675403240e294cd8504d4e083fe1 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 19:46:22 +0000 Subject: [PATCH 13/22] Plan each DuckDB V2 file fully at initialize and queue its chunks The root planner ran under a lock shared by the file's threads, and DuckDB handed the file to new threads until the root was done. While one thread waited on the zone-map reads of a filtered scan, every other thread queued on that lock instead of scanning other files, so selective ClickBench queries ran files nearly one at a time (q01 2.4x, q07 2.2x, q26 2.1x slower in CI). Run the root to the end, pruning included, on the thread that initializes the file, as V1 builds its splits there, and queue the chunks it spawns. DuckDB gets one thread per queued chunk at most. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/file_reader.rs | 6 +-- vortex-duckdb/src/scan_v2.rs | 93 +++++++++++++++++--------------- 2 files changed, 52 insertions(+), 47 deletions(-) diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index 96c62cc43a7..bc19ed19a1c 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -313,13 +313,13 @@ fn next_split_item(local: &mut LocalState) -> VortexResult { /// the global lock. It only moves references between the file and the local /// state, and reader_scan drops the retired ones. pub fn reader_try_initialize_scan(file: &mut OpenFileReader, local: &mut LocalState) -> bool { - // A V2 scan's threads take chunks from the file's scan until it has spawned - // them all, so the file has work for another thread until then. + // A V2 scan's threads take chunks from the file's queue until it is empty, + // so the file has work for one more thread per queued chunk at most. let v2 = file .scan_state .as_ref() .and_then(|state| state.v2.get()) - .map(|scan| (!scan.is_exhausted()).then_some(Split::V2)); + .map(|scan| scan.claim().then_some(Split::V2)); let Some(split) = v2.unwrap_or_else(|| file.splits.pop()) else { // DuckDB releases the file after this call, under the global lock. // Give the reference of the file to this thread so it drops outside. diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index 2ff1668fa3a..84cbb85cf96 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -6,11 +6,14 @@ //! A planner says what it needs: compute, or reads. This driver decides where each runs, and //! DuckDB's own threads run all of it. Each file is prepared once and its planners are composed //! over the prepared scan: with a filter, its pruning plans over the scanned rows first, -//! continuing into the chunks of the rows they keep; without, the chunks directly. The root -//! planner's run is shared by every thread scanning the file. A thread takes the next chunk the -//! root spawns and runs it on a run of its own, inline, waiting only for its reads. Threads keep -//! taking chunks until the root is done, so the file's work spreads over as many threads as -//! DuckDB gives the scan, however its rows divide. +//! continuing into the chunks of the rows they keep; without, the chunks directly. +//! +//! The thread initializing the file runs the root planner to the end, pruning included, and +//! queues the chunks it spawns. DuckDB initializes a file on one thread while its others scan +//! other files, so no thread waits on another file's pruning reads. DuckDB then gives the file +//! one thread per queued chunk at most. A thread takes the next chunk and runs it on a run of its +//! own, inline, waiting only for its reads, and keeps taking chunks until the queue is empty, so +//! the file's work spreads over as many threads as it has chunks, however its rows divide. //! //! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk //! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up @@ -19,11 +22,12 @@ //! threads busy. The root groups the chunks into parts, a few per thread, and spawns a chunk //! planner per part. +use std::collections::VecDeque; use std::future; use std::ops::Range; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use parking_lot::Mutex; @@ -61,10 +65,12 @@ const PARTS_PER_THREAD: usize = 4; pub(crate) struct FileScan { scan: Arc, io: Arc, - /// The run of the root planner, locked while a thread takes the chunk it spawns next. - root: Mutex, - /// Set once the root has spawned its last chunk. - exhausted: AtomicBool, + /// The work the root planned, in row order, that no thread has taken yet. + work: Mutex>, + /// The work the root planned, taken or not. + planned: usize, + /// The threads DuckDB has given the file so far. + claims: AtomicUsize, /// The rows of the chunks taken so far, for DuckDB's progress bar. rows_taken: AtomicU64, } @@ -78,9 +84,9 @@ enum Work { } impl FileScan { - /// Prepares `builder` and composes its planners, cutting chunks on `boundaries`, the sorted - /// rows the file's chunks of the scanned columns end at. Lowering and planning touch no - /// segment. + /// Prepares `builder`, composes its planners, cutting chunks on `boundaries`, the sorted rows + /// the file's chunks of the scanned columns end at, and runs the root planner to the end on + /// the calling thread, waiting for its pruning reads. pub(crate) fn try_new(builder: ScanBuilder, boundaries: &[u64]) -> VortexResult { let scan = Arc::new(builder.prepare()?); let io = scan.io(); @@ -97,25 +103,28 @@ impl FileScan { }, io.session(), ); + let work = plan(&mut root)?; Ok(Self { scan, io, - root: Mutex::new(root), - exhausted: AtomicBool::new(false), + planned: work.len(), + work: Mutex::new(work), + claims: AtomicUsize::new(0), rows_taken: AtomicU64::new(0), }) } - /// Whether the root has spawned its last chunk, so a thread joining the scan finds no work. - pub(crate) fn is_exhausted(&self) -> bool { - self.exhausted.load(Ordering::Acquire) + /// Whether the file has work for one more thread: DuckDB gives it one thread per piece of + /// planned work at most. + pub(crate) fn claim(&self) -> bool { + self.claims.fetch_add(1, Ordering::Relaxed) < self.planned } /// The percentage of the scanned rows whose chunks threads have taken. pub(crate) fn progress(&self) -> f64 { let rows = self.scan.rows(); let total = rows.end - rows.start; - if total == 0 || self.exhausted.load(Ordering::Relaxed) { + if total == 0 || self.work.lock().is_empty() { return 100.0; } 100.0 * self.rows_taken.load(Ordering::Relaxed) as f64 / total as f64 @@ -149,33 +158,29 @@ impl FileScan { } } - /// Advances the root until it spawns a chunk or produces a batch, or `None` once it is done. + /// The next work no thread has taken, or `None` once all is taken. fn next_work(&self) -> VortexResult> { - if self.exhausted.load(Ordering::Acquire) { - return Ok(None); + let work = self.work.lock().pop_front(); + if let Some(Work::Chunk(scope, _)) = &work { + self.rows_taken + .fetch_add(scope.rows.end - scope.rows.start, Ordering::Relaxed); } - let mut root = self.root.lock(); - loop { - if self.exhausted.load(Ordering::Acquire) { - return Ok(None); - } - match root.advance()? { - Progress::Batch(batch) => return Ok(Some(Work::Batch(batch.array))), - Progress::Spawn(scope, chunk) => { - self.rows_taken - .fetch_add(scope.rows.end - scope.rows.start, Ordering::Relaxed); - return Ok(Some(Work::Chunk(scope, chunk))); - } - Progress::RootDone(_) if root.live_roots() > 0 => {} - Progress::RootDone(_) | Progress::Idle => { - self.exhausted.store(true, Ordering::Release); - return Ok(None); - } - // The root waits for the pruning plans' reads. The threads waiting on the lock - // would have nothing to run until they arrive. - Progress::Waiting => { - RUNTIME.block_on(future::poll_fn(|cx| root.poll_completion(cx)))?; - } + Ok(work) + } +} + +/// Runs `root` to the end, waiting for its reads, and returns the chunks it spawns and the +/// batches it produces, in order. +fn plan(root: &mut Run) -> VortexResult> { + let mut work = VecDeque::new(); + loop { + match root.advance()? { + Progress::Batch(batch) => work.push_back(Work::Batch(batch.array)), + Progress::Spawn(scope, chunk) => work.push_back(Work::Chunk(scope, chunk)), + Progress::RootDone(_) if root.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => return Ok(work), + Progress::Waiting => { + RUNTIME.block_on(future::poll_fn(|cx| root.poll_completion(cx)))?; } } } From 1ed88fe308a31f9a90e6ccc3d5fbf35eb6139894 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 21:08:01 +0000 Subject: [PATCH 14/22] Turn on the DuckDB V2 scan for TPC-DS and ClickBench in vx-bench Pass --scan-v2 true from the benchmark orchestrator for DuckDB's TPC-DS and ClickBench runs instead of changing duckdb-bench's default, which goes back to V2 for TPC-H only. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- .../bench_orchestrator/runner/executor.py | 6 ++++++ bench-orchestrator/tests/test_executor.py | 21 +++++++++++++++++++ benchmarks/duckdb-bench/src/main.rs | 8 +++---- 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/bench-orchestrator/bench_orchestrator/runner/executor.py b/bench-orchestrator/bench_orchestrator/runner/executor.py index 1369c164f00..49ba000f65d 100644 --- a/bench-orchestrator/bench_orchestrator/runner/executor.py +++ b/bench-orchestrator/bench_orchestrator/runner/executor.py @@ -17,6 +17,10 @@ console = Console() +# DuckDB benchmarks that run on the planner-driven V2 Vortex scan. TPC-H turns it on in +# duckdb-bench itself. +DUCKDB_SCAN_V2_BENCHMARKS = frozenset({Benchmark.TPCDS, Benchmark.CLICKBENCH}) + @final class BenchmarkExecutor: @@ -56,6 +60,8 @@ def build_command( cmd.extend(["--formats", ",".join(fmt.value for fmt in formats)]) if self.backend == Engine.DUCKDB: cmd.append("--delete-duckdb-database") + if benchmark in DUCKDB_SCAN_V2_BENCHMARKS: + cmd.extend(["--scan-v2", "true"]) if query is not None: cmd.extend(["--queries", str(query)]) diff --git a/bench-orchestrator/tests/test_executor.py b/bench-orchestrator/tests/test_executor.py index b8afb066e75..0a29fbb9f2d 100644 --- a/bench-orchestrator/tests/test_executor.py +++ b/bench-orchestrator/tests/test_executor.py @@ -5,6 +5,7 @@ import textwrap from pathlib import Path +import pytest from bench_orchestrator.config import Benchmark, Engine, Format from bench_orchestrator.runner.executor import BenchmarkExecutor @@ -33,6 +34,26 @@ def test_build_command_adds_duckdb_cleanup_flag() -> None: assert cmd[cmd.index("--opt") + 1] == "scale-factor=1.0" +@pytest.mark.parametrize( + ("engine", "benchmark", "expected"), + [ + (Engine.DUCKDB, Benchmark.TPCDS, True), + (Engine.DUCKDB, Benchmark.CLICKBENCH, True), + (Engine.DUCKDB, Benchmark.TPCH, False), + (Engine.DATAFUSION, Benchmark.CLICKBENCH, False), + ], +) +def test_build_command_turns_on_duckdb_scan_v2(engine: Engine, benchmark: Benchmark, expected: bool) -> None: + executor = BenchmarkExecutor(Path("/tmp/bench"), engine) + + cmd = executor.build_command(benchmark=benchmark, formats=[Format.VORTEX]) + + if expected: + assert cmd[cmd.index("--scan-v2") + 1] == "true" + else: + assert "--scan-v2" not in cmd + + def test_build_command_serializes_vortex_spatial_native_format() -> None: executor = BenchmarkExecutor(Path("/tmp/duckdb-bench"), Engine.DUCKDB) diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index 4b5e2cfbb8c..a047e275b24 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -103,7 +103,7 @@ struct Args { )] reuse: bool, - /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H, TPC-DS and ClickBench). + /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H). #[arg(long, action = clap::ArgAction::Set)] scan_v2: Option, } @@ -117,10 +117,8 @@ fn main() -> anyhow::Result<()> { let opts = Opts::from(args.options); let scan_v2 = args.scan_v2.unwrap_or_else(|| { - matches!( - args.benchmark, - BenchmarkArg::TpcH | BenchmarkArg::TpcDS | BenchmarkArg::ClickBench - ) && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0") + matches!(args.benchmark, BenchmarkArg::TpcH) + && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0") }); vortex_duckdb::set_scan_v2(scan_v2); From 431717c46f30197355fb05a5b956b85c9c32a4ce Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 21:52:49 +0000 Subject: [PATCH 15/22] Run a DuckDB V2 part's chunks inline so they share reads Each chunk a part spawned was admitted as its own root with a fresh IO session, and a session caches only its own reads. For a wide projection the chunks are cut on every column's boundaries, so a column chunk spanning several scan chunks was read once per scan chunk. DuckDB's late-materialization fetch of ClickBench q23 (SELECT * ... ORDER BY ... LIMIT 10) read many segments two or three times. ChunkPlanner::inline emits the chunks as children of the part's root, on its run and session, instead of spawning them. A thread runs one part at a time, so this gives up no parallelism. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/scan_v2.rs | 11 ++++++++--- vortex-layout/src/scan/execution/chunk.rs | 17 ++++++++++++++++- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index 84cbb85cf96..1efdd87a59b 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -15,6 +15,10 @@ //! own, inline, waiting only for its reads, and keeps taking chunks until the queue is empty, so //! the file's work spreads over as many threads as it has chunks, however its rows divide. //! +//! A part runs its chunks as children of one root, on the thread that takes it, so they share an +//! IO session: a segment that spans several chunks, such as a column chunk longer than the +//! scan's chunks, is read once. +//! //! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk //! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up //! to the scan's chunk size, but no further than leaves a few chunks per thread: a filtered scan's @@ -204,8 +208,8 @@ fn compose(scan: &Arc, chunks: Vec>, count: usize) -> B } } -/// Spawns a chunk planner per part of a scan's rows, in row order. A part is a run of whole -/// chunks. +/// Spawns a chunk planner per part of a scan's rows, in row order, which runs the part's chunks +/// inline. A part is a run of whole chunks. struct PartPlanner { scan: Arc, selection: Selection, @@ -272,7 +276,8 @@ impl Planner for PartPlanner { chunks, self.selection.clone(), kept, - ); + ) + .inline(); let scope = WorkScope { file_ordinal: 0, rows: part, diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index d8c423d7eb3..0a94e18d7eb 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -68,6 +68,8 @@ pub struct ChunkPlanner { selection: Selection, /// The rows zone pruning kept, over `rows`, or `None` when nothing was pruned. kept: Option, + /// Whether the chunks run as children of this planner instead of being spawned. + inline: bool, } impl ChunkPlanner { @@ -104,8 +106,17 @@ impl ChunkPlanner { chunks: chunks.into_iter(), selection, kept, + inline: false, } } + + /// Runs the chunks as children of this planner, on its run and IO session, instead of + /// spawning them for the driver to place. Chunks that read the same segment, such as a + /// column chunk longer than the scan's chunks, then share the read. + pub fn inline(mut self) -> Self { + self.inline = true; + self + } } impl IoConsumer for ChunkPlanner { @@ -130,7 +141,11 @@ impl Planner for ChunkPlanner { }; if let Some((scope, planner)) = plan_selection(&self.scan, scope, &self.selection, kept) { - return Ok(PlannerOutput::Spawn(scope, planner)); + return Ok(if self.inline { + PlannerOutput::Planner(scope, planner) + } else { + PlannerOutput::Spawn(scope, planner) + }); } } Ok(PlannerOutput::Done) From 746f52b6a2b250814866244b77491bc197505382 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 22:01:41 +0000 Subject: [PATCH 16/22] Choose the V2 scan per benchmark in vx-bench only datafusion-bench and duckdb-bench no longer pick the V2 Vortex scan by benchmark or read VORTEX_SCAN_V2: --scan-v2 defaults to false. The benchmark orchestrator passes --scan-v2 true for DataFusion TPC-H and for DuckDB TPC-H, TPC-DS and ClickBench. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- .../bench_orchestrator/runner/executor.py | 12 +++++++----- bench-orchestrator/tests/test_executor.py | 7 +++++-- benchmarks/datafusion-bench/src/main.rs | 11 ++++------- benchmarks/duckdb-bench/src/main.rs | 14 ++++++-------- 4 files changed, 22 insertions(+), 22 deletions(-) diff --git a/bench-orchestrator/bench_orchestrator/runner/executor.py b/bench-orchestrator/bench_orchestrator/runner/executor.py index 49ba000f65d..2814aabd9f2 100644 --- a/bench-orchestrator/bench_orchestrator/runner/executor.py +++ b/bench-orchestrator/bench_orchestrator/runner/executor.py @@ -17,9 +17,11 @@ console = Console() -# DuckDB benchmarks that run on the planner-driven V2 Vortex scan. TPC-H turns it on in -# duckdb-bench itself. -DUCKDB_SCAN_V2_BENCHMARKS = frozenset({Benchmark.TPCDS, Benchmark.CLICKBENCH}) +# The benchmarks each engine runs on the planner-driven V2 Vortex scan. +SCAN_V2_BENCHMARKS: dict[Engine, frozenset[Benchmark]] = { + Engine.DATAFUSION: frozenset({Benchmark.TPCH}), + Engine.DUCKDB: frozenset({Benchmark.TPCH, Benchmark.TPCDS, Benchmark.CLICKBENCH}), +} @final @@ -60,8 +62,8 @@ def build_command( cmd.extend(["--formats", ",".join(fmt.value for fmt in formats)]) if self.backend == Engine.DUCKDB: cmd.append("--delete-duckdb-database") - if benchmark in DUCKDB_SCAN_V2_BENCHMARKS: - cmd.extend(["--scan-v2", "true"]) + if benchmark in SCAN_V2_BENCHMARKS.get(self.backend, frozenset()): + cmd.extend(["--scan-v2", "true"]) if query is not None: cmd.extend(["--queries", str(query)]) diff --git a/bench-orchestrator/tests/test_executor.py b/bench-orchestrator/tests/test_executor.py index 0a29fbb9f2d..36ef737e766 100644 --- a/bench-orchestrator/tests/test_executor.py +++ b/bench-orchestrator/tests/test_executor.py @@ -39,11 +39,14 @@ def test_build_command_adds_duckdb_cleanup_flag() -> None: [ (Engine.DUCKDB, Benchmark.TPCDS, True), (Engine.DUCKDB, Benchmark.CLICKBENCH, True), - (Engine.DUCKDB, Benchmark.TPCH, False), + (Engine.DUCKDB, Benchmark.TPCH, True), + (Engine.DUCKDB, Benchmark.FINEWEB, False), + (Engine.DATAFUSION, Benchmark.TPCH, True), (Engine.DATAFUSION, Benchmark.CLICKBENCH, False), + (Engine.LANCE, Benchmark.TPCH, False), ], ) -def test_build_command_turns_on_duckdb_scan_v2(engine: Engine, benchmark: Benchmark, expected: bool) -> None: +def test_build_command_turns_on_scan_v2(engine: Engine, benchmark: Benchmark, expected: bool) -> None: executor = BenchmarkExecutor(Path("/tmp/bench"), engine) cmd = executor.build_command(benchmark=benchmark, formats=[Format.VORTEX]) diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 1a83d374fcb..765b34fafcf 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -120,9 +120,9 @@ struct Args { #[arg(long, value_delimiter = ',', value_parser = value_parser!(Format))] formats: Vec, - /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H). - #[arg(long, action = clap::ArgAction::Set)] - scan_v2: Option, + /// Use the planner-driven V2 Vortex scan. + #[arg(long, default_value_t = false, action = clap::ArgAction::Set)] + scan_v2: bool, #[arg(long = "opt", value_delimiter = ',', value_parser = value_parser!(Opt))] options: Vec, @@ -136,10 +136,7 @@ async fn main() -> anyhow::Result<()> { vortex_bench::configure_allocator(); 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") - }); + let scan_v2 = args.scan_v2; set_join_set_tracer(get_static_tracer())?; setup_logging_and_tracing(args.verbose, args.tracing)?; diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index a047e275b24..f0b2cede63d 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -103,9 +103,9 @@ struct Args { )] reuse: bool, - /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H). - #[arg(long, action = clap::ArgAction::Set)] - scan_v2: Option, + /// Use the planner-driven V2 Vortex scan. + #[arg(long, default_value_t = false, action = clap::ArgAction::Set)] + scan_v2: bool, } #[global_allocator] @@ -116,11 +116,9 @@ 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") - }); - vortex_duckdb::set_scan_v2(scan_v2); + if args.scan_v2 { + vortex_duckdb::set_scan_v2(true); + } setup_logging_and_tracing(args.verbose, args.tracing)?; From 4e9cede0ae109eb2f4f4ef603b2b80a6e2020bed Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 22:21:45 +0000 Subject: [PATCH 17/22] Prune each DuckDB V2 part again under a dynamic filter A file is planned, pruning included, when DuckDB initializes it. With 100 files and 96 threads every file is planned before DuckDB's top-N bound has tightened, so a dynamic comparison prunes nothing and every chunk is read. V1 prunes each split when it runs and skips the splits the tightened bound excludes: on ClickBench q24 V2 read 44 segments (8.5 MB) V1 never touched. When the filter has a dynamic comparison, run a part behind its own pruning when a thread takes it, intersected with the file's static pruning. The zone maps were read when the file was planned, so this reads nothing new. Adds top-N cases to the V2-matches-V1 test. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- .../src/e2e_test/vortex_scan_test.rs | 2 + vortex-duckdb/src/file_reader.rs | 14 +++- vortex-duckdb/src/scan_v2.rs | 69 ++++++++++++++----- 3 files changed, 68 insertions(+), 17 deletions(-) diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 97b310de84f..2be6adf2c59 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -1410,6 +1410,8 @@ fn query_i64_column(path: &str, query: &str) -> Result> { #[case::expression("SELECT a + b FROM ? WHERE a >= 50000 ORDER BY 1")] #[case::full_scan("SELECT a FROM ? ORDER BY a")] #[case::empty("SELECT a FROM ? WHERE b > 5000")] +#[case::top_n("SELECT a FROM ? ORDER BY b DESC, a LIMIT 7")] +#[case::filtered_top_n("SELECT a FROM ? WHERE b < 500 ORDER BY a DESC LIMIT 10")] fn test_vortex_scan_v2_matches_v1(#[case] query: &str) -> Result<()> { // Many row blocks, so the file is cut into several V2 splits. let rows = 100_000i64; diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index bc19ed19a1c..d8dc5f23072 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -12,6 +12,7 @@ use url::Url; use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::struct_::StructArrayExt as _; +use vortex::array::scalar_fn::fns::dynamic::DynamicComparison; use vortex::dtype::FieldPath; use vortex::error::VortexExpect; use vortex::error::VortexResult; @@ -218,7 +219,18 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor // global lock. if global.scan_v2 { let state = file.scan_state(); - let scan = FileScan::try_new(v2_builder(state, global), &v2_boundaries(state, global)?)?; + let dynamic = global + .filter + .filter + .as_ref() + .map(|filter| filter.contains::()) + .transpose()? + .unwrap_or(false); + let scan = FileScan::try_new( + v2_builder(state, global), + &v2_boundaries(state, global)?, + dynamic, + )?; if state.v2.set(scan).is_err() { vortex_panic!("file initialized twice"); } diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index 1efdd87a59b..3124e2c07b1 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -19,6 +19,11 @@ //! IO session: a segment that spans several chunks, such as a column chunk longer than the //! scan's chunks, is read once. //! +//! A filter with dynamic comparisons, such as the top-N bound DuckDB tightens during the scan, is +//! pruned again over each part when a thread takes it, as V1 prunes each split when it runs: a +//! part taken once the bound has tightened skips the zones it excludes. The zone maps were read +//! when the file was planned, so pruning a part again reads nothing. +//! //! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk //! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up //! to the scan's chunk size, but no further than leaves a few chunks per thread: a filtered scan's @@ -90,8 +95,13 @@ enum Work { impl FileScan { /// Prepares `builder`, composes its planners, cutting chunks on `boundaries`, the sorted rows /// the file's chunks of the scanned columns end at, and runs the root planner to the end on - /// the calling thread, waiting for its pruning reads. - pub(crate) fn try_new(builder: ScanBuilder, boundaries: &[u64]) -> VortexResult { + /// the calling thread, waiting for its pruning reads. With `dynamic`, each part is pruned + /// again when a thread takes it. + pub(crate) fn try_new( + builder: ScanBuilder, + boundaries: &[u64], + dynamic: bool, + ) -> VortexResult { let scan = Arc::new(builder.prepare()?); let io = scan.io(); let rows = scan.rows().clone(); @@ -100,7 +110,7 @@ impl FileScan { let chunks = chunk_ranges(rows.clone(), boundaries, max_rows); let mut root = Run::new(); root.admit( - compose(&scan, chunks, parts), + compose(&scan, chunks, parts, dynamic), WorkScope { file_ordinal: 0, rows, @@ -192,13 +202,24 @@ fn plan(root: &mut Run) -> VortexResult> { /// The planners of a scan, composed: with a filter, its pruning plans first over the scan's rows, /// continuing into the parts of the rows they keep; without, the parts directly. -fn compose(scan: &Arc, chunks: Vec>, count: usize) -> Box { +fn compose( + scan: &Arc, + chunks: Vec>, + count: usize, + dynamic: bool, +) -> Box { let rows = scan.rows().clone(); let parts = { let scan = Arc::clone(scan); let rows = rows.clone(); move |kept: Option| -> Box { - Box::new(PartPlanner::new(scan, rows, kept, parts(chunks, count))) + Box::new(PartPlanner::new( + scan, + rows, + kept, + parts(chunks, count), + dynamic, + )) } }; if scan.filter().is_some() { @@ -209,7 +230,8 @@ fn compose(scan: &Arc, chunks: Vec>, count: usize) -> B } /// Spawns a chunk planner per part of a scan's rows, in row order, which runs the part's chunks -/// inline. A part is a run of whole chunks. +/// inline, behind the part's own pruning when the filter is dynamic. A part is a run of whole +/// chunks. struct PartPlanner { scan: Arc, selection: Selection, @@ -219,6 +241,8 @@ struct PartPlanner { kept: Option, /// The chunks of each part, which tile the scan's rows. parts: std::vec::IntoIter>>, + /// Whether each part is pruned again when it runs. + dynamic: bool, } impl PartPlanner { @@ -227,6 +251,7 @@ impl PartPlanner { rows: Range, kept: Option, parts: Vec>>, + dynamic: bool, ) -> Self { Self { selection: scan.selection().clone(), @@ -234,6 +259,7 @@ impl PartPlanner { first_row: rows.start, kept, parts: parts.into_iter(), + dynamic, } } } @@ -270,19 +296,30 @@ impl Planner for PartPlanner { } None => None, }; - let chunks = ChunkPlanner::with_chunks( - Arc::clone(&self.scan), - part.clone(), - chunks, - self.selection.clone(), - kept, - ) - .inline(); let scope = WorkScope { file_ordinal: 0, - rows: part, + rows: part.clone(), + }; + let scan = Arc::clone(&self.scan); + let selection = self.selection.clone(); + if !self.dynamic { + let chunks = + ChunkPlanner::with_chunks(scan, part, chunks, selection, kept).inline(); + return Ok(PlannerOutput::Spawn(scope, Box::new(chunks))); + } + let next = move |pruned: Option| -> Box { + let kept = match (kept, pruned) { + (Some(kept), Some(pruned)) => Some(kept & &pruned), + (kept, pruned) => kept.or(pruned), + }; + Box::new( + ChunkPlanner::with_chunks(Arc::clone(&scan), part, chunks, selection, kept) + .inline(), + ) }; - return Ok(PlannerOutput::Spawn(scope, Box::new(chunks))); + let pruning = + PrunePlanner::new(Arc::clone(&self.scan), scope.rows.clone(), Box::new(next)); + return Ok(PlannerOutput::Spawn(scope, Box::new(pruning))); } Ok(PlannerOutput::Done) } From e4eb6f2d2ed4243669390c9c79630ab8fd099d35 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 22:41:54 +0000 Subject: [PATCH 18/22] Prune a DuckDB V2 file only per part under a dynamic filter Planning a file runs its pruning plans before DuckDB's top-N bound exists, so under a dynamic filter they prove little and only delay the file's first part behind the zone map reads, which for a long string column such as ClickBench's URL are large. Each part is pruned when a thread takes it anyway, so skip the file-level pruning when the filter is dynamic and let the parts prune, in parallel, against the latest bound. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/scan_v2.rs | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index 3124e2c07b1..1fa3a06c8e3 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -20,9 +20,10 @@ //! scan's chunks, is read once. //! //! A filter with dynamic comparisons, such as the top-N bound DuckDB tightens during the scan, is -//! pruned again over each part when a thread takes it, as V1 prunes each split when it runs: a -//! part taken once the bound has tightened skips the zones it excludes. The zone maps were read -//! when the file was planned, so pruning a part again reads nothing. +//! pruned over each part when a thread takes it instead of over the file when it is planned, as +//! V1 prunes each split when it runs: a part taken once the bound has tightened skips the zones +//! it excludes. A file is planned before the bound exists, so pruning it then would only delay +//! its first part behind the zone map reads. //! //! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk //! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up @@ -95,8 +96,8 @@ enum Work { impl FileScan { /// Prepares `builder`, composes its planners, cutting chunks on `boundaries`, the sorted rows /// the file's chunks of the scanned columns end at, and runs the root planner to the end on - /// the calling thread, waiting for its pruning reads. With `dynamic`, each part is pruned - /// again when a thread takes it. + /// the calling thread, waiting for its pruning reads. With `dynamic`, each part is pruned when + /// a thread takes it instead. pub(crate) fn try_new( builder: ScanBuilder, boundaries: &[u64], @@ -200,8 +201,9 @@ fn plan(root: &mut Run) -> VortexResult> { } } -/// The planners of a scan, composed: with a filter, its pruning plans first over the scan's rows, -/// continuing into the parts of the rows they keep; without, the parts directly. +/// The planners of a scan, composed: with a static filter, its pruning plans first over the scan's +/// rows, continuing into the parts of the rows they keep; otherwise, the parts directly, which +/// prune themselves under a `dynamic` filter. fn compose( scan: &Arc, chunks: Vec>, @@ -222,7 +224,7 @@ fn compose( )) } }; - if scan.filter().is_some() { + if scan.filter().is_some() && !dynamic { Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(parts))) } else { parts(None) @@ -241,7 +243,7 @@ struct PartPlanner { kept: Option, /// The chunks of each part, which tile the scan's rows. parts: std::vec::IntoIter>>, - /// Whether each part is pruned again when it runs. + /// Whether each part is pruned when it runs. dynamic: bool, } @@ -307,11 +309,8 @@ impl Planner for PartPlanner { ChunkPlanner::with_chunks(scan, part, chunks, selection, kept).inline(); return Ok(PlannerOutput::Spawn(scope, Box::new(chunks))); } - let next = move |pruned: Option| -> Box { - let kept = match (kept, pruned) { - (Some(kept), Some(pruned)) => Some(kept & &pruned), - (kept, pruned) => kept.or(pruned), - }; + // A dynamic filter is not pruned over the file, so `kept` is `None` here. + let next = move |kept: Option| -> Box { Box::new( ChunkPlanner::with_chunks(Arc::clone(&scan), part, chunks, selection, kept) .inline(), From 9b8a90efaf9a06944ef7ad77e0d9c9f773370539 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 22:58:59 +0000 Subject: [PATCH 19/22] Share a segment's read across a V2 scan's live IO sessions Every part of a V2 scan reads through its own IO session, so parts whose chunks cut into the same column segment each read it. ClickBench q23's row-id fetch, which reads all 105 columns for a few rows, re-requested 430 segments (46 MB) that V1's layout readers read once. Keep a scan-wide map of each segment's read, held weakly, so a session shares a read another session still holds. A read lives no longer than before: the map only keeps it while some session does. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-layout/src/scan/v2/io.rs | 116 ++++++++++++++++++++++++++------ 1 file changed, 94 insertions(+), 22 deletions(-) diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 0e690ebfde7..4569b24f5c7 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -12,6 +12,7 @@ use futures::StreamExt; use futures::TryFutureExt; use futures::future::BoxFuture; use futures::future::Shared; +use futures::future::WeakShared; use futures::stream::FuturesUnordered; use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; @@ -46,26 +47,40 @@ pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { } /// Serves splits from any [`SegmentSource`], by segment id. +/// +/// Its sessions share a segment's read while any of them still holds it, so splits that run at +/// once and need the same segment, such as chunks that each cut into one column segment, read +/// it once. pub(crate) struct SegmentScanIo { segments: Arc, ranges: SegmentRanges, + reads: SharedReads, } impl SegmentScanIo { pub(crate) fn new(segments: Arc, ranges: SegmentRanges) -> Self { - Self { segments, ranges } + Self { + segments, + ranges, + reads: SharedReads::default(), + } } } impl IoService for SegmentScanIo { fn session(&self) -> Arc { - Arc::new(SegmentIoSource::new( + Arc::new(SegmentIoSource::with_shared_reads( Arc::clone(&self.segments), Arc::clone(&self.ranges), + Arc::clone(&self.reads), )) } } +/// The reads of a scan's sessions, by segment, held weakly: a read lives only while a session +/// holds it, so sharing it keeps no segment in memory longer. +type SharedReads = Arc>>>; + /// Serves a split's range requests from a segment source, without blocking. /// /// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once, @@ -85,13 +100,27 @@ pub(super) struct SegmentIoSource { prefetched: Mutex>, /// Announced segments, requested and never polled. announced: Mutex>>, + /// The reads of the scan's other sessions. + shared: SharedReads, } +/// The read of a segment. +type SegmentRead = BoxFuture<'static, Result>>; + /// A read started ahead of the fetch that needs it. -type Prefetched = Shared>>>; +type Prefetched = Shared; impl SegmentIoSource { + #[cfg(test)] pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self::with_shared_reads(segments, ranges, SharedReads::default()) + } + + fn with_shared_reads( + segments: Arc, + ranges: SegmentRanges, + shared: SharedReads, + ) -> Self { Self { segments, ranges, @@ -99,8 +128,31 @@ impl SegmentIoSource { ready: Mutex::default(), prefetched: Mutex::default(), announced: Mutex::default(), + shared, } } + + /// The read of `segment`: one another session of the scan still holds, else the read of an + /// announcement of it, else a new one. + fn read(&self, segment: SegmentId) -> Prefetched { + let mut shared = self.shared.lock(); + if let Some(read) = shared.get(&segment).and_then(WeakShared::upgrade) { + return read; + } + let read = self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared(); + if let Some(weak) = read.downgrade() { + shared.insert(segment, weak); + } + read + } } impl IoSource for SegmentIoSource { @@ -129,15 +181,7 @@ impl IoSource for SegmentIoSource { } if request.intent == IoIntent::Prefetch { prefetched.entry(segment).or_insert_with(|| { - let read = self - .announced - .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) - .map_err(Arc::new) - .boxed() - .shared(); + let read = self.read(segment); // Polling registers the read as wanted, so the source starts it now. drop( read.clone() @@ -152,16 +196,7 @@ impl IoSource for SegmentIoSource { // first has completed. let mut bytes = prefetched .entry(segment) - .or_insert_with(|| { - self.announced - .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) - .map_err(Arc::new) - .boxed() - .shared() - }) + .or_insert_with(|| self.read(segment)) .clone() .map_err(move |err| vortex_err!("read of segment {segment} failed: {err}")) .boxed(); @@ -249,10 +284,12 @@ mod tests { use vortex_io::request::IoOwnerId; use vortex_io::request::IoRequest; use vortex_io::request::IoRequestId; + use vortex_io::request::IoService; use vortex_io::request::IoSource; use vortex_io::request::IoTarget; use super::SegmentIoSource; + use super::SegmentScanIo; use super::segment_ranges; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; @@ -370,6 +407,41 @@ mod tests { Ok(()) } + /// Sessions of one scan share a segment's read while one of them holds it, and read it again + /// once none does. + #[tokio::test] + async fn sessions_share_a_held_read() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let scan = SegmentScanIo::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let first = scan.session(); + let second = scan.session(); + for (io, id) in [(&first, 0), (&second, 1)] { + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(id)); + assert!(completion.result.is_ok()); + } + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + + first.clear(); + second.clear(); + let third = scan.session(); + third.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 2)])?; + assert!( + std::future::poll_fn(|cx| third.poll_completion(cx)) + .await? + .result + .is_ok() + ); + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + Ok(()) + } + #[tokio::test] async fn announcements_keep_independent_consumers() -> VortexResult<()> { let segments = Arc::new(CountingSegments::default()); From 3b2d95dd895af6f3495f63e0f10c8ffee61405a3 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 22:59:50 +0000 Subject: [PATCH 20/22] Prune a DuckDB V2 file again before its parts under a dynamic filter This reverts e4eb6f2. CI measured ClickBench q24 at +7.6% against V1 without the file-level pruning, where it was -8.8% with it, and no query gained, so prune the file at initialize and each part again when a thread takes it. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-duckdb/src/scan_v2.rs | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index 1fa3a06c8e3..3124e2c07b1 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -20,10 +20,9 @@ //! scan's chunks, is read once. //! //! A filter with dynamic comparisons, such as the top-N bound DuckDB tightens during the scan, is -//! pruned over each part when a thread takes it instead of over the file when it is planned, as -//! V1 prunes each split when it runs: a part taken once the bound has tightened skips the zones -//! it excludes. A file is planned before the bound exists, so pruning it then would only delay -//! its first part behind the zone map reads. +//! pruned again over each part when a thread takes it, as V1 prunes each split when it runs: a +//! part taken once the bound has tightened skips the zones it excludes. The zone maps were read +//! when the file was planned, so pruning a part again reads nothing. //! //! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk //! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up @@ -96,8 +95,8 @@ enum Work { impl FileScan { /// Prepares `builder`, composes its planners, cutting chunks on `boundaries`, the sorted rows /// the file's chunks of the scanned columns end at, and runs the root planner to the end on - /// the calling thread, waiting for its pruning reads. With `dynamic`, each part is pruned when - /// a thread takes it instead. + /// the calling thread, waiting for its pruning reads. With `dynamic`, each part is pruned + /// again when a thread takes it. pub(crate) fn try_new( builder: ScanBuilder, boundaries: &[u64], @@ -201,9 +200,8 @@ fn plan(root: &mut Run) -> VortexResult> { } } -/// The planners of a scan, composed: with a static filter, its pruning plans first over the scan's -/// rows, continuing into the parts of the rows they keep; otherwise, the parts directly, which -/// prune themselves under a `dynamic` filter. +/// The planners of a scan, composed: with a filter, its pruning plans first over the scan's rows, +/// continuing into the parts of the rows they keep; without, the parts directly. fn compose( scan: &Arc, chunks: Vec>, @@ -224,7 +222,7 @@ fn compose( )) } }; - if scan.filter().is_some() && !dynamic { + if scan.filter().is_some() { Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(parts))) } else { parts(None) @@ -243,7 +241,7 @@ struct PartPlanner { kept: Option, /// The chunks of each part, which tile the scan's rows. parts: std::vec::IntoIter>>, - /// Whether each part is pruned when it runs. + /// Whether each part is pruned again when it runs. dynamic: bool, } @@ -309,8 +307,11 @@ impl Planner for PartPlanner { ChunkPlanner::with_chunks(scan, part, chunks, selection, kept).inline(); return Ok(PlannerOutput::Spawn(scope, Box::new(chunks))); } - // A dynamic filter is not pruned over the file, so `kept` is `None` here. - let next = move |kept: Option| -> Box { + let next = move |pruned: Option| -> Box { + let kept = match (kept, pruned) { + (Some(kept), Some(pruned)) => Some(kept & &pruned), + (kept, pruned) => kept.or(pruned), + }; Box::new( ChunkPlanner::with_chunks(Arc::clone(&scan), part, chunks, selection, kept) .inline(), From c867e73ed91db08691959f397a8365a7190fb689 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 00:51:22 +0000 Subject: [PATCH 21/22] Plan a V2 projection of fields as one pack of column plans A projection that packs fields by path, such as DuckDB's row-id fetch of every column beside its row ids, kept an Eval of the whole pack above a Pack plan. Each chunk then built the struct, renamed its fields and packed it again, comparing the 105-field ClickBench dtypes deeply at each step, and rebuilt a row index for every row of the chunk before taking the few selected ones. EvalPackRule now lifts a pack whose fields are all field paths into a non-nullable Pack child into one PackPlan, stepping each path straight into its column's plan, and ExpressionPackRule steps a field path into its field without annotating every field. RowIdx compiles a lazy sequence, as the V1 row index reader does, so a sparse mask produces only its rows. On ClickBench q23, single threaded under callgrind, the V2 scan now runs 3.21G instructions, against 3.56G before and 3.32G for V1, and plans each file in a third of the time. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- vortex-layout/src/layouts/row_idx/mod.rs | 4 +- vortex-layout/src/plan/physical/plans/eval.rs | 54 ++++++++++++++++--- vortex-layout/src/plan/physical/plans/pack.rs | 30 +++++++++++ .../src/plan/physical/plans/row_idx.rs | 13 ++--- vortex-layout/src/plan/physical/tests.rs | 36 +++++++++++++ 5 files changed, 119 insertions(+), 18 deletions(-) diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index f9641cbd6cb..389cad6059d 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -293,8 +293,8 @@ fn row_idx_dtype() -> DType { DType::Primitive(PType::U64, NonNullable) } -// Returns a SequenceArray representing the row indices for the given row range, -fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { +/// A [`SequenceArray`] of the row indices of `row_range`, offset by `row_offset`. +pub(crate) fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { Sequence::try_new( PValue::U64(row_offset + row_range.start), PValue::U64(1), diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index 043f0676198..d3ed93771bb 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -12,6 +12,7 @@ use vortex_array::expr::BoundExpression; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::pack::Pack; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -20,6 +21,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::plan::physical::Pack as PackVTable; use crate::plan::physical::PackPlan; use crate::plan::physical::Plan; use crate::plan::physical::PlanChildren; @@ -194,7 +196,12 @@ impl PlanParentReduceRule for EvalEvalRule { } } -/// Eval(pack({0: expr})) -> PackPlan(expr) for a non-nullable pack +/// Eval(pack({0: expr})) -> PackPlan(expr) for a non-nullable pack, and +/// Eval(pack({0: path0, 1: path1, ..}), Pack) -> PackPlan(path0, path1, ..) when every field is a +/// field path into a non-nullable [`Pack`] child. +/// +/// Lifting a pack of field paths into the plan lets each field reach its column's plan, so a scan +/// assembles each chunk's struct once instead of packing, renaming and repacking it per chunk. #[derive(Debug)] pub(crate) struct EvalPackRule; impl PlanReduceRule for EvalPackRule { @@ -205,21 +212,52 @@ impl PlanReduceRule for EvalPackRule { if expr.nullability != Nullability::NonNullable { return Ok(None); } - let [child_expr] = plan.expression.children() else { - return Ok(None); - }; - let fields = StructFields::new(expr.names.clone(), vec![child_expr.dtype().clone()]); - + let child_exprs = plan.expression.children(); let child_plan = plan.child(0)?.vortex_expect("no child plan"); - let child_plan = EvalPlan::try_new(child_expr.clone(), child_plan)?.into_plan(); + let child_plans = match child_exprs { + [child_expr] => vec![EvalPlan::try_new(child_expr.clone(), child_plan)?.into_plan()], + // A child read by several fields is shared only where each field steps into it by + // name, so that each field reads the plan of its own column. + _ => { + let Some(pack) = child_plan.as_opt::() else { + return Ok(None); + }; + if child_exprs.is_empty() + || pack.dtype().is_nullable() + || !child_exprs.iter().all(is_field_path) + { + return Ok(None); + } + child_exprs + .iter() + .map(|child| pack.field_path_plan(child)) + .collect::>>()? + } + }; + let fields = StructFields::new( + expr.names.clone(), + child_exprs + .iter() + .map(|child| child.dtype().clone()) + .collect(), + ); let plan = PackPlan::try_new( fields, Nullability::NonNullable, plan.row_count(), - vec![child_plan], + child_plans, None, )?; Ok(Some(plan.into_plan())) } } + +/// Whether `expression` is the root or a chain of field accesses from it, such as `$.a.b`. +pub(super) fn is_field_path(expression: &BoundExpression) -> bool { + if expression.is_root() { + return true; + } + expression.as_opt::().is_some() + && matches!(expression.children(), [child] if is_field_path(child)) +} diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index d3886de19ad..4b82bc6d0f9 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -32,6 +32,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use super::eval::is_field_path; use crate::plan::physical::Eval; use crate::plan::physical::EvalPlan; use crate::plan::physical::Plan; @@ -291,6 +292,30 @@ fn validate_validity_child(expected_row_count: u64, child: &PlanRef) -> VortexRe } impl PackPlan { + /// The plan of `path`, a chain of field accesses such as `$.a.b` over this non-nullable + /// pack, stepped into the plan of its first field. + pub(crate) fn field_path_plan(&self, path: &BoundExpression) -> VortexResult { + let mut first = path; + while let [child] = first.children() + && !child.is_root() + { + first = child; + } + let Some(name) = first.as_opt::() else { + return Ok(EvalPlan::try_new(path.clone(), self.to_plan())?.into_plan()); + }; + let index = self + .fields() + .find(name) + .ok_or_else(|| vortex_err!("Struct expression references unknown field '{name}'"))?; + let field = field_plan(self, index)?; + let lowered = step_into_struct_field(path.clone(), name, field.dtype().clone())?; + if lowered.is_root() { + return Ok(field); + } + Ok(EvalPlan::try_new(lowered, field)?.into_plan()) + } + /// Rebuilds this plan with only `fields`, which must be a subset of the current fields. /// /// Pruning is only sound for a non-nullable struct: dropping a field of a nullable struct @@ -338,6 +363,11 @@ impl PlanParentReduceRule for ExpressionPackRule { } let expression = parent.expression(); + // A field path steps straight into its field's plan, without the annotation of every + // field that a general expression needs. + if !expression.is_root() && is_field_path(expression) { + return Ok(Some(child.field_path_plan(expression)?)); + } let fields = child.fields(); let referenced_fields = descendent_bound_annotations(expression, make_bound_free_field_annotator(fields)) diff --git a/vortex-layout/src/plan/physical/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs index 5c45d2cae4f..e300506b866 100644 --- a/vortex-layout/src/plan/physical/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -17,7 +17,6 @@ use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar_fn::fns::pack::Pack as PackFn; -use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -26,6 +25,7 @@ use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::layouts::row_idx::RowIdx as RowIdxFn; +use crate::layouts::row_idx::idx_array; use crate::plan::physical::EvalPlan; use crate::plan::physical::PackPlan; use crate::plan::physical::Plan; @@ -104,13 +104,10 @@ impl PlanVTable for RowIdx { if mask.all_false() { return Ok(None); } - let offset = compiler.row_offset(); - let indices = Buffer::from_iter(rows.start + offset..rows.end + offset).into_array(); - let indices = if mask.all_true() { - indices - } else { - indices.filter(mask.clone())? - }; + // The sequence and its filter are lazy, so only the rows the mask keeps are produced. + let indices = idx_array(compiler.row_offset(), &rows) + .into_array() + .filter(mask.clone())?; Ok(Some(Chain::new(OnceSource::new(indices)))) } } diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index c86edcf4d30..6791a0003c2 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -630,6 +630,42 @@ fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { Ok(()) } +/// A projection of row ids and renamed columns, as an engine fetching rows by id plans it, +/// reads each column's plan directly beside the row indices. +#[test] +fn field_projection_with_row_idx_packs_column_plans() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), + Nullability::NonNullable, + ), + vec![flat(3, value_dtype.clone(), 0), flat(3, value_dtype, 1)], + ) + .into_layout(); + let projection = pack( + [ + ("id", row_idx()), + ("y", get_item("b", root())), + ("x", get_item("a", root())), + ], + Nullability::NonNullable, + ); + let plan = make_row_idx_plan(projection, make_plan(layout)?)?; + + let optimized = optimize(plan, &SESSION)?; + insta::assert_snapshot!(optimized.display_tree(), @" + root: vortex.plan.pack({id=u64, y=i32, x=i32}, rows=3) + id: vortex.plan.row_idx(u64, rows=3) + y: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) + x: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) + "); + Ok(()) +} + #[test] fn empty_projection_prunes_row_idx_child_fields() -> VortexResult<()> { let value_dtype = primitive(PType::I32, Nullability::NonNullable); From 0ad75b121412f58cff2d383fb0450eecb339e0c7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 00:51:22 +0000 Subject: [PATCH 22/22] Retain a sparse V2 fetch's reads and plan only its selected chunks A row-id fetch, such as DuckDB's late materialization in ClickBench q23, selects a few rows of each file, but each wide column segment it touches is needed by every chunk that selects a row of it. The chunks run one after another, so sharing a read only while another session holds it read those segments again: 3905 requests for 2312 segments, 849 MB read for 321 MB of segments. A scan whose selection keeps at most one row in 1024 now retains its reads for its lifetime, as V1's layout readers do. The DuckDB driver also planned every chunk of such a file as a part, most of them empty, and offered DuckDB a thread for each. It now keeps only the chunks with a selected row, each its own part. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M8NsCJMx4C45Wf2Df24qzZ --- .../src/e2e_test/vortex_scan_test.rs | 2 + vortex-duckdb/src/scan_v2.rs | 12 +- vortex-layout/src/scan/execution/mod.rs | 18 +++ vortex-layout/src/scan/v2/io.rs | 103 ++++++++++++++++-- 4 files changed, 122 insertions(+), 13 deletions(-) diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 2be6adf2c59..af005b4a5d4 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -1412,6 +1412,8 @@ fn query_i64_column(path: &str, query: &str) -> Result> { #[case::empty("SELECT a FROM ? WHERE b > 5000")] #[case::top_n("SELECT a FROM ? ORDER BY b DESC, a LIMIT 7")] #[case::filtered_top_n("SELECT a FROM ? WHERE b < 500 ORDER BY a DESC LIMIT 10")] +// DuckDB fetches the columns beyond the sort key by row id, as in ClickBench q23. +#[case::late_materialized_top_n("SELECT * FROM ? WHERE b < 700 ORDER BY a DESC LIMIT 5")] fn test_vortex_scan_v2_matches_v1(#[case] query: &str) -> Result<()> { // Many row blocks, so the file is cut into several V2 splits. let rows = 100_000i64; diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index 3124e2c07b1..b74ad56b1e2 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -105,9 +105,17 @@ impl FileScan { let scan = Arc::new(builder.prepare()?); let io = scan.io(); let rows = scan.rows().clone(); - let parts = get_available_parallelism().unwrap_or(1) * PARTS_PER_THREAD; + let mut parts = get_available_parallelism().unwrap_or(1) * PARTS_PER_THREAD; let max_rows = (scan.chunk_rows() as u64).min(merged_rows(&rows, parts)); - let chunks = chunk_ranges(rows.clone(), boundaries, max_rows); + let mut chunks = chunk_ranges(rows.clone(), boundaries, max_rows); + if scan.sparse_selection() { + // A sparse selection, such as an engine's fetch of a few rows by id, selects rows in + // few chunks: each of them is its own part, and the rest are not planned at all, so + // DuckDB gives the file a thread per chunk that has work, as V1 does with its splits. + let selection = scan.selection(); + chunks.retain(|chunk| selection.row_count_in_range(chunk) > 0); + parts = chunks.len(); + } let mut root = Run::new(); root.admit( compose(&scan, chunks, parts, dynamic), diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 497e2684904..61a785c9568 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -71,6 +71,9 @@ impl SegmentLocation { } } +/// A selection of at most one row in this many of a scan's rows is sparse. +const SPARSE_SELECTION: u64 = 1024; + /// A scan's plans over one file, built once from the file's footer and read by every planner of /// the scan. Nothing here touches a segment, and no planner plans again. pub struct PreparedScan { @@ -123,12 +126,27 @@ impl PreparedScan { self.chunk_rows } + /// Whether the scan selects at most one row in 1024 of its rows, as an + /// engine fetching a few rows by id does. + pub fn sparse_selection(&self) -> bool { + !matches!(self.selection, Selection::All) + && self + .selection + .row_count_in_range(&self.rows) + .saturating_mul(SPARSE_SELECTION) + <= self.rows.end - self.rows.start + } + /// The service a driver takes IO sessions from: every read a planner of the scan requests is /// served from it by segment. + /// + /// A [sparse selection](Self::sparse_selection) retains its reads for the scan's lifetime, + /// as the chunks that select its few rows read the same wide segments one after another. pub fn io(&self) -> Arc { Arc::new(SegmentScanIo::new( Arc::clone(&self.segments), Arc::clone(&self.ranges), + self.sparse_selection(), )) } diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 4569b24f5c7..1b77b5f578d 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -50,7 +50,8 @@ pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { /// /// Its sessions share a segment's read while any of them still holds it, so splits that run at /// once and need the same segment, such as chunks that each cut into one column segment, read -/// it once. +/// it once. A scan that retains its reads keeps every read for its lifetime, so splits that run +/// one after another share them too. pub(crate) struct SegmentScanIo { segments: Arc, ranges: SegmentRanges, @@ -58,11 +59,23 @@ pub(crate) struct SegmentScanIo { } impl SegmentScanIo { - pub(crate) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + /// A scan's IO, which retains its reads when `retain` is set. + /// + /// Retaining suits a sparse selection, such as an engine fetching a few rows by id: its + /// reads are few, but each wide column segment it touches is needed by every chunk that + /// selects a row of it. + pub(crate) fn new( + segments: Arc, + ranges: SegmentRanges, + retain: bool, + ) -> Self { Self { segments, ranges, - reads: SharedReads::default(), + reads: Arc::new(Mutex::new(ScanReads { + retain, + reads: HashMap::default(), + })), } } } @@ -77,9 +90,45 @@ impl IoService for SegmentScanIo { } } -/// The reads of a scan's sessions, by segment, held weakly: a read lives only while a session -/// holds it, so sharing it keeps no segment in memory longer. -type SharedReads = Arc>>>; +/// The reads of a scan's sessions, by segment. +type SharedReads = Arc>; + +/// The reads of a scan's sessions, by segment. +#[derive(Default)] +struct ScanReads { + /// Whether reads are held for the scan's lifetime rather than while a session holds them. + retain: bool, + reads: HashMap, +} + +/// A read of a scan, held weakly so that sharing it keeps no segment in memory longer, or +/// retained. +enum ScanRead { + Weak(WeakShared), + Retained(Prefetched), +} + +impl ScanReads { + /// The read of `segment` the scan still has. + fn get(&self, segment: SegmentId) -> Option { + match self.reads.get(&segment)? { + ScanRead::Weak(read) => read.upgrade(), + ScanRead::Retained(read) => Some(read.clone()), + } + } + + fn insert(&mut self, segment: SegmentId, read: &Prefetched) { + let read = if self.retain { + ScanRead::Retained(read.clone()) + } else { + let Some(read) = read.downgrade() else { + return; + }; + ScanRead::Weak(read) + }; + self.reads.insert(segment, read); + } +} /// Serves a split's range requests from a segment source, without blocking. /// @@ -136,7 +185,7 @@ impl SegmentIoSource { /// announcement of it, else a new one. fn read(&self, segment: SegmentId) -> Prefetched { let mut shared = self.shared.lock(); - if let Some(read) = shared.get(&segment).and_then(WeakShared::upgrade) { + if let Some(read) = shared.get(segment) { return read; } let read = self @@ -148,9 +197,7 @@ impl SegmentIoSource { .map_err(Arc::new) .boxed() .shared(); - if let Some(weak) = read.downgrade() { - shared.insert(segment, weak); - } + shared.insert(segment, &read); read } } @@ -417,7 +464,11 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let scan = SegmentScanIo::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let scan = SegmentScanIo::new( + Arc::clone(&segments) as _, + segment_ranges(&[location]), + false, + ); let first = scan.session(); let second = scan.session(); for (io, id) in [(&first, 0), (&second, 1)] { @@ -442,6 +493,36 @@ mod tests { Ok(()) } + /// A scan that retains its reads shares a read with a later session after the earlier one + /// let it go. + #[tokio::test] + async fn retained_reads_outlive_their_sessions() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let scan = SegmentScanIo::new( + Arc::clone(&segments) as _, + segment_ranges(&[location]), + true, + ); + for id in 0..2 { + let io = scan.session(); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + assert!( + std::future::poll_fn(|cx| io.poll_completion(cx)) + .await? + .result + .is_ok() + ); + io.clear(); + } + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + Ok(()) + } + #[tokio::test] async fn announcements_keep_independent_consumers() -> VortexResult<()> { let segments = Arc::new(CountingSegments::default());