From c56ec58c71a5f99983134960d6306f2396af9e03 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 12:28:04 +0100 Subject: [PATCH 01/23] Plan a V2 scan as a chain of planners from one prepared scan A `PreparedScan` holds a scan's plans over one file, built once from the footer, and makes its planners; nothing in it touches a segment and no planner plans again. The planners form a chain, each emitting the next through the planning protocol: - `PrunePlanner` proves the filter's zone maps over the scan's rows, prefetching the zone tables together, and emits one chunk planner for the rows no proof excludes. Rows a proof excludes are never planned. - `ChunkPlanner` cuts those rows into chunks and emits a filter or a projection planner per chunk, skipping chunks nothing is selected or kept of. The run drives an emitted chunk ahead of the next, so chunks waiting on reads overlap. - `FilterPlanner` prefetches its chunk's conjunct segments and announces the projection's at its first compute, runs the conjuncts in the adaptive order over the surviving rows, and emits the projection. - `ProjectionPlanner` emits one morsel per piece of the rows. A stage with nothing to do is never made: a filter without proofs starts at the chunk planner, a scan without a filter at the projection. A scan runs as one task owning one planning run, instead of one task and run per split plus a private run for the proofs; `with_concurrency` sets nothing. The IO session keeps every read it starts, so a later stage or morsel of a chunk that reads the segment again shares the read. The segments a stage prefetches come from compiling its plans. DataFusion attaches each partition's rows from the cached footer when it repartitions a file, so the V2 opener reads exactly those rows instead of translating a byte range through the layout. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-layout/src/scan/execution/chunk.rs | 120 +++++++++++ vortex-layout/src/scan/execution/filter.rs | 182 ++++++++++------- vortex-layout/src/scan/execution/mod.rs | 92 +++++++-- .../src/scan/execution/projection.rs | 34 ++-- vortex-layout/src/scan/execution/prune.rs | 177 ++++++++++++++++ vortex-layout/src/scan/execution/selection.rs | 38 ++-- vortex-layout/src/scan/v2/io.rs | 53 +++-- vortex-layout/src/scan/v2/mod.rs | 192 +++++++----------- vortex-layout/src/scan/v2/split.rs | 74 ------- vortex-layout/src/scan/v2/tests.rs | 133 +++++++++--- 10 files changed, 739 insertions(+), 356 deletions(-) create mode 100644 vortex-layout/src/scan/execution/chunk.rs create mode 100644 vortex-layout/src/scan/execution/prune.rs delete mode 100644 vortex-layout/src/scan/v2/split.rs diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs new file mode 100644 index 00000000000..268841b0d2f --- /dev/null +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Cuts a scan's rows into chunks and plans each one. + +use std::ops::Range; +use std::sync::Arc; +use std::vec::IntoIter; + +use vortex_error::VortexResult; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +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 super::PreparedScan; +use super::SegmentLocation; +use super::selection::plan_selection; + +/// The row grid every chunk boundary lies on. +pub(crate) const ROW_BLOCK_SIZE: usize = 8192; +/// The rows of a projection chunk. +pub(crate) const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; +/// The rows of a filtered chunk: large, so a conjunct's reads coalesce and its mask is cheap. +pub(crate) const FILTERED_ROWS_PER_CHUNK: usize = 16 * MAX_ROWS_PER_CHUNK; +const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; + +/// Limit estimated bytes per chunk without changing the writer-aligned row granularity. +pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { + if row_count <= MAX_ROWS_PER_CHUNK as u64 { + return MAX_ROWS_PER_CHUNK; + } + let bytes: u64 = locations + .iter() + .map(|location| u64::from(location.length)) + .sum(); + let bytes_per_row = bytes.div_ceil(row_count).max(1); + let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; + blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE +} + +/// Fixed boundaries in the file's row domain, clipped to the requested range. +pub(crate) fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { + if rows.is_empty() { + return Vec::new(); + } + let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64; + (first..rows.end) + .step_by(rows_per_chunk) + .map(|start| { + start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end) + }) + .collect() +} + +/// Emits one planner per chunk of a scan's rows, in row order. A chunk none of whose rows is +/// selected and kept is never planned. +/// +/// The run drives an emitted chunk ahead of this planner's next compute, so a chunk waiting on +/// its reads lets the next chunk start its own, and the reads of the chunks in flight overlap. +pub(in crate::scan) struct ChunkPlanner { + scan: Arc, + rows: Range, + chunks: IntoIter>, + selection: Selection, + /// The rows zone pruning kept, over `rows`, or `None` when nothing was proved. + kept: Option, +} + +impl ChunkPlanner { + pub(super) fn new( + scan: Arc, + rows: Range, + selection: Selection, + kept: Option, + ) -> Self { + let chunks = row_chunks(rows.clone(), scan.chunk_rows).into_iter(); + Self { + scan, + rows, + chunks, + selection, + kept, + } + } +} + +impl IoConsumer for ChunkPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for ChunkPlanner { + fn state(&self) -> State { + if self.chunks.as_slice().is_empty() { + State::Done + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + let kept = self.kept.as_ref().map(|kept| (kept, self.rows.start)); + for chunk in self.chunks.by_ref() { + let scope = WorkScope { + file_ordinal: 0, + rows: chunk, + }; + if let Some((scope, planner)) = plan_selection(&self.scan, scope, &self.selection, kept) + { + return Ok(PlannerOutput::Planner(scope, planner)); + } + } + Ok(PlannerOutput::Done) + } +} diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 2ce5f588f02..3eb9c291b56 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use bit_vec::BitVec; @@ -11,6 +12,7 @@ use vortex_array::expr::transform::find_between; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; +use vortex_io::request::IoBatch; use vortex_io::request::IoConsumer; use vortex_io::request::IoIntent; use vortex_io::request::IoRequest; @@ -22,9 +24,11 @@ use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; use vortex_session::VortexSession; -use super::ScanPlans; +use super::PreparedScan; +use super::SegmentLocation; use super::SelectedRows; use super::adaptor::AdaptorStep; +use super::adaptor::Piece; use super::adaptor::ScanAdaptor; use super::projection::ProjectionPlanner; use crate::plan::physical::FilterPlan; @@ -36,11 +40,13 @@ use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; use crate::plan::physical::plan_row_idx_expression; use crate::scan::filter::FilterExpr; +use crate::segments::SegmentId; -/// Plans and adaptive conjunct ordering shared by all splits of a scan. +/// Plans and adaptive conjunct ordering shared by every chunk of a scan. pub(in crate::scan) struct FilterPlans { expression: FilterExpr, predicates: Vec, + /// A metadata-only proof per conjunct, where the conjunct has one. pruning: Vec>, } @@ -81,6 +87,11 @@ impl FilterPlans { pruning, }) } + + /// The pruning proofs, run once over a scan's rows before its chunks are planned. + pub(in crate::scan) fn proofs(&self) -> impl Iterator { + self.pruning.iter().flatten() + } } /// A pruning proof must not fetch data when its requested statistics are unavailable. @@ -102,37 +113,87 @@ fn metadata_only(plan: &PlanRef) -> VortexResult { Ok(true) } -/// Proves zone exclusions, evaluates conjuncts over survivors, then hands the mask to projection. +/// Evaluates a chunk's conjuncts over its selected rows, then hands the mask to projection. +/// +/// The first compute starts the reads of every conjunct's segments for the selected rows and +/// announces the projection's, so the conjuncts run back to back and the projection's reads +/// coalesce with them. Conjuncts run in the order the shared [`FilterExpr`] picks. pub(super) struct FilterPlanner { - plans: ScanPlans, + scan: Arc, selected: Option, filter: Arc, remaining: BitVec, - prune_index: usize, - active: Option<(usize, bool, ScanAdaptor)>, + active: Option<(usize, ScanAdaptor)>, pieces: Vec, announced: bool, + /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. + next_hint: u32, +} + +/// The segments `plans` read over `rows` under `mask`, each once, in id order. Compiling finds +/// them without IO. +pub(super) fn segments_read<'a>( + session: &VortexSession, + plans: impl Iterator, + rows: Range, + mask: &Mask, +) -> VortexResult> { + let mut segments = Vec::new(); + for plan in plans { + segments.extend(Scan::segments(session, plan, rows.clone(), mask)?); + } + segments.sort_unstable(); + segments.dedup(); + Ok(segments) } impl FilterPlanner { - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { - let filter = Arc::clone(plans.filter.as_ref().vortex_expect("filter plans exist")); + pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { + let filter = Arc::clone(scan.filter.as_ref().vortex_expect("filter plans exist")); Self { remaining: BitVec::from_elem(filter.predicates.len(), true), - plans, + scan, selected: Some(selected), filter, - prune_index: 0, active: None, pieces: Vec::new(), announced: false, + next_hint: u32::MAX, } } } +/// Publishes a hint for each of `segments` not already in `batch`. +pub(super) fn hint( + locations: &[SegmentLocation], + next_hint: &mut u32, + intent: IoIntent, + segments: &[SegmentId], + batch: &mut IoBatch, +) { + for segment in segments { + let target = locations[**segment as usize].target(); + if batch.iter().any(|request| request.target == target) { + continue; + } + batch.push(IoRequest { + intent, + request: IoRequestId(*next_hint), + target, + }); + *next_hint -= 1; + } +} + +/// Executes a scan's boolean piece to a mask, nulls as false. +pub(super) fn piece_mask(session: &VortexSession, piece: Piece) -> VortexResult { + let mut ctx = session.create_execution_ctx(); + piece.array.fill_null(false)?.execute::(&mut ctx) +} + impl IoConsumer for FilterPlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { - if let Some((_, _, adaptor)) = &mut self.active { + if let Some((_, adaptor)) = &mut self.active { adaptor.set_io_result(request, result); } } @@ -144,7 +205,7 @@ impl Planner for FilterPlanner { return State::Done; } match &self.active { - Some((_, _, adaptor)) if adaptor.state() != State::Done => adaptor.state(), + Some((_, adaptor)) if adaptor.state() != State::Done => adaptor.state(), _ => State::NeedsCompute, } } @@ -155,42 +216,47 @@ impl Planner for FilterPlanner { }; if !self.announced { self.announced = true; + let session = &self.scan.session; + let rows = selected.scope.rows.clone(); + let predicates = segments_read( + session, + self.filter.predicates.iter(), + rows.clone(), + &selected.mask, + )?; + let projection = segments_read( + session, + std::iter::once(&self.scan.projection), + rows, + &selected.mask, + )?; let mut batch = Vec::new(); - for plan in self - .filter - .pruning - .iter() - .flatten() - .chain(self.filter.predicates.iter()) - .chain(std::iter::once(&self.plans.projection)) - { - let segments = Scan::segments( - &self.plans.session, - plan, - selected.scope.rows.clone(), - &selected.mask, - )?; - for segment in segments { - batch.push(IoRequest { - intent: IoIntent::Announce, - request: IoRequestId(u32::MAX - u32::try_from(batch.len())?), - target: self.plans.locations[*segment as usize].target(), - }); - } - } + let locations = &self.scan.locations; + hint( + locations, + &mut self.next_hint, + IoIntent::Prefetch, + &predicates, + &mut batch, + ); + hint( + locations, + &mut self.next_hint, + IoIntent::Announce, + &projection, + &mut batch, + ); if !batch.is_empty() { return Ok(PlannerOutput::NeedsIO(batch)); } } - if let Some((index, pruning, adaptor)) = &mut self.active { + if let Some((index, adaptor)) = &mut self.active { if adaptor.state() != State::Done { return Ok(match adaptor.compute()? { AdaptorStep::Yield => PlannerOutput::Continue, AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), AdaptorStep::Piece(piece) => { - let mut ctx = self.plans.session.create_execution_ctx(); - let mask = piece.array.fill_null(false)?.execute::(&mut ctx)?; - self.pieces.push(mask); + self.pieces.push(piece_mask(&self.scan.session, piece)?); PlannerOutput::Continue } }); @@ -198,22 +264,18 @@ impl Planner for FilterPlanner { let input = selected.mask.true_count(); let result = Mask::concat(self.pieces.iter())?; vortex_ensure_eq!(result.len(), input); - let result = if *pruning { !result } else { result }; selected.mask = selected.mask.intersect_by_rank(&result); tracing::debug!( row_start = selected.scope.rows.start, row_end = selected.scope.rows.end, input_rows = input, output_rows = selected.mask.true_count(), - pruning = *pruning, conjunct = *index, "v2 filter conjunct evaluated" ); - if !*pruning { - self.filter - .expression - .report_selectivity(*index, selected.mask.true_count() as f64 / input as f64); - } + self.filter + .expression + .report_selectivity(*index, selected.mask.true_count() as f64 / input as f64); self.active = None; self.pieces.clear(); } @@ -221,44 +283,28 @@ impl Planner for FilterPlanner { self.selected = None; return Ok(PlannerOutput::Done); } - - let mut next = None; - while self.prune_index < self.filter.pruning.len() { - let index = self.prune_index; - self.prune_index += 1; - if let Some(plan) = &self.filter.pruning[index] { - next = Some((index, true, plan.clone())); - break; - } - } - if next.is_none() - && let Some(index) = self.filter.expression.next_conjunct(&self.remaining) - { + if let Some(index) = self.filter.expression.next_conjunct(&self.remaining) { self.remaining.set(index, false); - next = Some((index, false, self.filter.predicates[index].clone())); - } - if let Some((index, pruning, plan)) = next { let scan = Scan::try_new( - self.plans.session.clone(), - plan, + self.scan.session.clone(), + self.filter.predicates[index].clone(), vec![Split { rows: selected.scope.rows.clone(), mask: selected.mask.clone(), }], )? - .with_row_offset(self.plans.row_offset) - .with_value_cache(self.plans.values.clone()); + .with_row_offset(self.scan.row_offset) + .with_value_cache(self.scan.values.clone()); self.active = Some(( index, - pruning, - ScanAdaptor::new(scan, Arc::clone(&self.plans.locations)), + ScanAdaptor::new(scan, Arc::clone(&self.scan.locations)), )); return Ok(PlannerOutput::Continue); } let selected = self.selected.take().vortex_expect("selected rows exist"); Ok(PlannerOutput::Planner( selected.scope.clone(), - Box::new(ProjectionPlanner::new(self.plans.clone(), selected)), + Box::new(ProjectionPlanner::new(Arc::clone(&self.scan), selected)), )) } } diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 53cc685bccb..aa49bb38c3d 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -1,32 +1,49 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Scan execution over physical plans through the planner I/O and compute protocol. +//! The planners of a V2 scan, and the prepared plans they run. //! -//! A `ProjectionPlanner` emits bounded `ProjectionMorsel`s for a work chunk's selected rows. Each -//! morsel runs the physical projection plan through -//! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published -//! through a `ScanAdaptor`, so the planning driver can advance other work while reads wait. +//! 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. mod adaptor; +mod chunk; mod filter; mod projection; +mod prune; mod selection; +use std::ops::Range; use std::sync::Arc; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; use vortex_mask::Mask; +use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; use vortex_session::VortexSession; +pub(in crate::scan) 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; +#[cfg(test)] +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(super) use self::selection::plan_selection; +use self::projection::ProjectionPlanner; +use self::prune::PrunePlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; -/// Where a segment's bytes live in its source. +/// Where a segment's bytes live in the file. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct SegmentLocation { /// Byte offset of the segment from the start of the source. @@ -48,26 +65,71 @@ impl SegmentLocation { } } -/// What every split of a scan shares: the projection plan and where its segments live. -#[derive(Clone)] -pub(super) struct ScanPlans { +/// 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 { /// The session used for decoding and expression evaluation. pub(super) session: VortexSession, /// 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. pub(super) projection: PlanRef, + /// The filter's conjuncts, their plans and their proofs. pub(super) filter: Option>, /// The global row index of the root row domain's first row. pub(super) row_offset: u64, - /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans - /// always read the same complete row domain, independently of a split's selection. + /// Whole values, such as dictionary values, shared by the chunks of this scan. Their plans + /// always read the same complete row domain, independently of a chunk's selection. pub(super) values: ValueCache, + /// The rows of one chunk, the unit of work a chunk planner emits. + pub(super) chunk_rows: usize, } -/// The rows of a split selected for projection. -struct SelectedRows { - /// The split the rows belong to. +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() { + Box::new(PrunePlanner::new(Arc::clone(self), rows, selection)) + } else { + Box::new(self.chunk_planner(rows, selection, None)) + } + } + + /// 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 planner of one chunk's selected rows: its conjuncts when the scan has a filter, else + /// its projection. + pub(super) fn split_planner(self: &Arc, selected: SelectedRows) -> Box { + if self.filter.is_some() { + Box::new(FilterPlanner::new(Arc::clone(self), selected)) + } else { + Box::new(ProjectionPlanner::new(Arc::clone(self), selected)) + } + } +} + +/// The rows a planner works on: a chunk, and which of its rows are selected. +pub(super) struct SelectedRows { + /// The chunk the rows belong to. scope: WorkScope, /// Which rows of `scope.rows` are selected. mask: Mask, diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 00c0a0614f6..141b7377c26 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; use std::vec::IntoIter; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_io::request::IoConsumer; @@ -23,7 +23,7 @@ use vortex_scan::planning::planner::WorkScope; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; use crate::plan::physical::selection::selected; -use crate::scan::execution::ScanPlans; +use crate::scan::execution::PreparedScan; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; @@ -34,13 +34,13 @@ const MAX_OUTPUT_ROWS: usize = 65_536; /// Keeps sparse selections together and subdivides dense selections to avoid large copies when /// converting chunked projection results to Arrow. pub(super) struct ProjectionPlanner { - plans: ScanPlans, + scan: Arc, selected: IntoIter, } impl ProjectionPlanner { /// Creates a planner for `selected`. - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { vec![selected] } else { @@ -64,7 +64,7 @@ impl ProjectionPlanner { .collect() }; Self { - plans, + scan, selected: selected.into_iter(), } } @@ -92,7 +92,7 @@ impl Planner for ProjectionPlanner { } Ok(PlannerOutput::Morsel( selected.scope.clone(), - Box::new(ProjectionMorsel::new(self.plans.clone(), selected)), + Box::new(ProjectionMorsel::new(Arc::clone(&self.scan), selected)), )) } } @@ -100,7 +100,7 @@ impl Planner for ProjectionPlanner { /// Runs the projection plan over a split's selected rows and returns them as one batch, in row /// order. A selection that projects to no rows finishes without a batch. struct ProjectionMorsel { - plans: Option, + scan: Arc, dtype: DType, selected: SelectedRows, adaptor: Option, @@ -110,10 +110,10 @@ struct ProjectionMorsel { impl ProjectionMorsel { /// Creates a morsel projecting `selected`. - fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + fn new(scan: Arc, selected: SelectedRows) -> Self { Self { - dtype: plans.projection.dtype().clone(), - plans: Some(plans), + dtype: scan.projection.dtype().clone(), + scan, selected, adaptor: None, pieces: Vec::new(), @@ -163,14 +163,10 @@ impl Morsel for ProjectionMorsel { vortex_bail!("ProjectionMorsel: compute called after Done"); } let Some(adaptor) = self.adaptor.as_mut() else { - let plans = self - .plans - .take() - .vortex_expect("an uncompiled projection still has its plans"); let scan = Scan::try_new( - plans.session, + self.scan.session.clone(), selected( - plans.projection, + self.scan.projection.clone(), self.selected.scope.rows.clone(), self.selected.mask.clone(), ), @@ -179,9 +175,9 @@ impl Morsel for ProjectionMorsel { mask: Mask::new_true(self.selected.mask.len()), }], )? - .with_row_offset(plans.row_offset) - .with_value_cache(plans.values); - self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); + .with_row_offset(self.scan.row_offset) + .with_value_cache(self.scan.values.clone()); + self.adaptor = Some(ScanAdaptor::new(scan, Arc::clone(&self.scan.locations))); return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs new file mode 100644 index 00000000000..8545f0b8f13 --- /dev/null +++ b/vortex-layout/src/scan/execution/prune.rs @@ -0,0 +1,177 @@ +// 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. + +use std::ops::Range; +use std::sync::Arc; +use std::vec::IntoIter; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +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 super::PreparedScan; +use super::adaptor::AdaptorStep; +use super::adaptor::ScanAdaptor; +use super::filter::hint; +use super::filter::piece_mask; +use super::filter::segments_read; +use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::Split; + +/// Runs the filter's zone-map proofs over the scan's rows and emits one chunk planner for the +/// rows no proof excludes. The zone maps are the only reads, started together. +/// +/// Rows a proof excludes are never planned: a chunk none of whose rows survive is skipped by +/// the chunk planner, and the chunks that remain run their conjuncts over the kept rows only. +pub(super) struct PrunePlanner { + scan: Arc, + rows: Range, + selection: Option, + /// The rows no proof has excluded yet, over `rows`. + mask: Mask, + proofs: IntoIter, + active: Option, + pieces: Vec, + announced: bool, + /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. + next_hint: u32, + done: bool, +} + +impl PrunePlanner { + pub(super) fn new(scan: Arc, rows: Range, selection: Selection) -> Self { + let proofs: Vec<_> = scan + .filter + .iter() + .flat_map(|filter| filter.proofs().cloned()) + .collect(); + let len = usize::try_from(rows.end - rows.start).unwrap_or(usize::MAX); + Self { + scan, + rows, + selection: Some(selection), + mask: Mask::new_true(len), + proofs: proofs.into_iter(), + active: None, + pieces: Vec::new(), + announced: false, + next_hint: u32::MAX, + done: false, + } + } +} + +impl IoConsumer for PrunePlanner { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some(adaptor) = &mut self.active { + adaptor.set_io_result(request, result); + } + } +} + +impl Planner for PrunePlanner { + fn state(&self) -> State { + if self.done { + return State::Done; + } + match &self.active { + Some(adaptor) if adaptor.state() != State::Done => adaptor.state(), + _ => State::NeedsCompute, + } + } + + fn compute(&mut self) -> VortexResult { + if self.done { + return Ok(PlannerOutput::Done); + } + if !self.announced { + self.announced = true; + let segments = segments_read( + &self.scan.session, + self.proofs.as_slice().iter(), + self.rows.clone(), + &self.mask, + )?; + let mut batch = Vec::new(); + hint( + &self.scan.locations, + &mut self.next_hint, + IoIntent::Prefetch, + &segments, + &mut batch, + ); + if !batch.is_empty() { + return Ok(PlannerOutput::NeedsIO(batch)); + } + } + if let Some(adaptor) = &mut self.active { + if adaptor.state() != State::Done { + return Ok(match adaptor.compute()? { + AdaptorStep::Yield => PlannerOutput::Continue, + AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + AdaptorStep::Piece(piece) => { + self.pieces.push(piece_mask(&self.scan.session, piece)?); + PlannerOutput::Continue + } + }); + } + let input = self.mask.true_count(); + let excluded = Mask::concat(self.pieces.iter())?; + vortex_ensure_eq!(excluded.len(), input); + self.mask = self.mask.intersect_by_rank(&!excluded); + tracing::debug!( + row_start = self.rows.start, + row_end = self.rows.end, + input_rows = input, + output_rows = self.mask.true_count(), + "v2 zone proof evaluated" + ); + self.active = None; + self.pieces.clear(); + } + if self.mask.all_false() { + self.done = true; + return Ok(PlannerOutput::Done); + } + if let Some(proof) = self.proofs.next() { + let scan = Scan::try_new( + self.scan.session.clone(), + proof, + vec![Split { + rows: self.rows.clone(), + mask: self.mask.clone(), + }], + )? + .with_row_offset(self.scan.row_offset) + .with_value_cache(self.scan.values.clone()); + self.active = Some(ScanAdaptor::new(scan, Arc::clone(&self.scan.locations))); + return Ok(PlannerOutput::Continue); + } + // Every proof has run: the chunks of the kept rows are planned from here on. + self.done = true; + let selection = self.selection.take().unwrap_or_default(); + let kept = std::mem::replace(&mut self.mask, Mask::new_true(0)); + Ok(PlannerOutput::Planner( + WorkScope { + file_ordinal: 0, + rows: self.rows.clone(), + }, + Box::new( + self.scan + .chunk_planner(self.rows.clone(), selection, Some(kept)), + ), + )) + } +} diff --git a/vortex-layout/src/scan/execution/selection.rs b/vortex-layout/src/scan/execution/selection.rs index fe751b48887..d4bd8f0f41e 100644 --- a/vortex-layout/src/scan/execution/selection.rs +++ b/vortex-layout/src/scan/execution/selection.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Resolves known selections before allocating projection planners or their masks. +//! Resolves known selections before allocating a chunk's planner or its masks. + +use std::sync::Arc; use vortex_error::VortexExpect; use vortex_mask::Mask; @@ -9,25 +11,33 @@ use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; -use super::ScanPlans; +use super::PreparedScan; use super::SelectedRows; -use super::filter::FilterPlanner; -use super::projection::ProjectionPlanner; -/// Plans only the selected span, retaining one projection planner and IO session per work chunk. -pub(crate) fn plan_selection( - plans: &ScanPlans, +/// Plans only the selected span of a chunk, with one planner and one mask per chunk. +/// +/// `kept` is the mask zone pruning left over the scan's rows, with the row its first bit +/// stands for. A chunk none of whose rows were kept is not planned. +pub(super) fn plan_selection( + scan: &Arc, scope: WorkScope, selection: &Selection, + kept: Option<(&Mask, u64)>, ) -> Option<(WorkScope, Box)> { - let selected = select_rows(scope, selection)?; + let mut selected = select_rows(scope, selection)?; + if let Some((kept, first_row)) = kept { + let offset = usize::try_from(selected.scope.rows.start - first_row) + .vortex_expect("a work chunk fits usize"); + let kept = kept.slice(offset..offset + selected.mask.len()); + let mask = &selected.mask & &kept; + let first = mask.first()?; + let end = mask.last()? + 1; + selected.scope.rows = + selected.scope.rows.start + first as u64..selected.scope.rows.start + end as u64; + selected.mask = mask.slice(first..end); + } let scope = selected.scope.clone(); - let planner: Box = if plans.filter.is_some() { - Box::new(FilterPlanner::new(plans.clone(), selected)) - } else { - Box::new(ProjectionPlanner::new(plans.clone(), selected)) - }; - Some((scope, planner)) + Some((scope, scan.split_planner(selected))) } fn select_rows(mut scope: WorkScope, selection: &Selection) -> Option { diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index ce565846d7e..2779739b281 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -80,7 +80,8 @@ pub(super) struct SegmentIoSource { ranges: SegmentRanges, reads: Mutex>>, ready: Mutex>, - /// Reads started by prefetches, consumed by the next fetch of their segment. + /// Every read the session has started, by segment, shared by later fetches of the segment + /// until the session is cleared. prefetched: Mutex>, /// Announced segments, requested and never polled. announced: Mutex>>, @@ -146,19 +147,24 @@ impl IoSource for SegmentIoSource { }); continue; } - let mut bytes = match prefetched.remove(&segment) { - Some(read) => read - .map_err(move |err| { - vortex_err!("prefetched read of segment {segment} failed: {err}") - }) - .boxed(), - None => self - .announced - .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)), - }; + // A segment stays with the session once read, so a later stage or morsel of the + // split that reads it again shares the read instead of starting another once the + // 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() + }) + .clone() + .map_err(move |err| vortex_err!("read of segment {segment} failed: {err}")) + .boxed(); let id = request.request; // In-memory and cached segments need no allocated asynchronous task. A pending // read is polled again by FuturesUnordered before the driver can go to sleep. @@ -325,7 +331,7 @@ mod tests { } /// A fetch of a prefetched segment waits on the prefetch's read rather than reading again, - /// and only the fetch is delivered. + /// only the fetch is delivered, and a later fetch of the segment shares the same read. #[tokio::test] async fn fetch_reuses_a_prefetched_read() -> VortexResult<()> { let segments = Arc::new(CountingSegments::default()); @@ -350,6 +356,16 @@ mod tests { let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; assert_eq!(completion.request, IoRequestId(2)); assert!(completion.result.is_ok()); + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + + io.clear(); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 3)])?; + assert!( + std::future::poll_fn(|cx| io.poll_completion(cx)) + .await? + .result + .is_ok() + ); assert_eq!(segments.0.load(Ordering::Relaxed), 2); Ok(()) } @@ -379,10 +395,9 @@ mod tests { .result .is_ok() ); - assert_eq!( - segments.0.load(Ordering::Relaxed), - if id < 4 { 2 } else { 3 } - ); + // The first fetch takes one announcement and the session keeps its read, so the + // later fetches share it and the other announcement is never consumed. + assert_eq!(segments.0.load(Ordering::Relaxed), 2); } Ok(()) } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index c6b359fa328..d0b28a60c10 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -20,10 +20,10 @@ mod cache; mod io; -mod split; #[cfg(test)] mod tests; +use std::future; use std::ops::Range; use std::sync::Arc; @@ -31,6 +31,7 @@ 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; @@ -45,11 +46,12 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; 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; use crate::LayoutRef; @@ -57,48 +59,17 @@ 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; +use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; use crate::scan::execution::FilterPlans; -use crate::scan::execution::ScanPlans; +use crate::scan::execution::MAX_ROWS_PER_CHUNK; +use crate::scan::execution::PreparedScan; +use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; -use crate::scan::execution::plan_selection; +use crate::scan::execution::rows_per_chunk; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; -use crate::scan::v2::split::SplitPlan; use crate::segments::SegmentSource; -const ROW_BLOCK_SIZE: usize = 8192; -const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; -const FILTERED_ROWS_PER_CHUNK: usize = 16 * MAX_ROWS_PER_CHUNK; -const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; - -/// Limit estimated bytes per task without changing the writer-aligned row granularity. -fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { - if row_count <= MAX_ROWS_PER_CHUNK as u64 { - return MAX_ROWS_PER_CHUNK; - } - let bytes: u64 = locations - .iter() - .map(|location| u64::from(location.length)) - .sum(); - let bytes_per_row = bytes.div_ceil(row_count).max(1); - let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; - blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE -} - -/// Fixed boundaries in the file's row domain, clipped to the requested range. -fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { - if rows.is_empty() { - return Vec::new(); - } - let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64; - (first..rows.end) - .step_by(rows_per_chunk) - .map(|start| { - start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end) - }) - .collect() -} - /// The rows a byte range of the file stands for when an engine splits the file by bytes before /// opening it. Consecutive byte ranges give consecutive row ranges that tile the file, with /// boundaries on the grid's 8,192-row blocks so the executor's chunks stay block aligned. @@ -140,7 +111,6 @@ pub struct ScanBuilder { ordered: bool, row_range: Option>, selection: Selection, - concurrency: usize, map_fn: Arc VortexResult + Send + Sync>, row_offset: u64, } @@ -156,7 +126,6 @@ impl ScanBuilder { ordered: true, row_range: None, selection: Selection::default(), - concurrency: 4, map_fn: Arc::new(Ok), row_offset: 0, } @@ -226,10 +195,10 @@ impl ScanBuilder { self } - /// Sets the number of active splits per worker thread. Must be greater than zero. - pub fn with_concurrency(mut self, concurrency: usize) -> Self { + /// Kept for callers of the default executor: a V2 scan runs on one task, its chunks + /// overlapping their reads inside one planning run, so this sets nothing. + pub fn with_concurrency(self, concurrency: usize) -> Self { assert!(concurrency > 0); - self.concurrency = concurrency; self } @@ -247,7 +216,6 @@ impl ScanBuilder { 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)), } @@ -259,7 +227,7 @@ impl ScanBuilder { let handle = self.session.handle(); Ok(stream::once(async move { handle - .spawn_cpu(move || PreparedScan::try_new(self)?.execute_stream()) + .spawn_cpu(move || ScanRun::try_new(self)?.execute_stream()) .await }) .try_flatten() @@ -267,23 +235,22 @@ impl ScanBuilder { } } -/// A scan planned over its file, ready to run its splits. -struct PreparedScan { +/// A scan prepared over its file, ready to run. +struct ScanRun { session: VortexSession, - plans: ScanPlans, - /// Serves the splits' reads. + scan: Arc, + /// Serves the scan's reads. io: Arc, - /// The scan's splits, over the root row domain, in row order. - splits: Vec>, + /// The rows of the root row domain the scan covers. + rows: Range, selection: Selection, map_fn: Arc VortexResult + Send + Sync>, ordered: bool, - concurrency: usize, } -impl PreparedScan { - /// Lowers the file's layout, plans the projection, and partitions the selected range into - /// fixed-size row chunks. +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; let cache = file.cache.unwrap_or_default(); @@ -318,7 +285,6 @@ impl PreparedScan { } else { MAX_ROWS_PER_CHUNK }; - let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); let ranges = match cache.ranges() { Some(ranges) => ranges, None => cache.insert_ranges(segment_ranges(&file.locations)), @@ -326,87 +292,71 @@ impl PreparedScan { let io = Arc::new(SegmentScanIo::new(file.segments, ranges)); Ok(Self { session: builder.session.clone(), - plans: ScanPlans { + scan: Arc::new(PreparedScan { session: builder.session, locations: file.locations, projection, filter, row_offset: builder.row_offset, values: cache.values(), - }, + chunk_rows, + }), io, - splits, + rows: builder.row_range.unwrap_or(0..row_count), selection: builder.selection, map_fn: builder.map_fn, ordered: builder.ordered, - concurrency: builder.concurrency, }) } - /// The splits with selected rows, each ready to be admitted to a planning run. - fn split_plans(&self) -> Vec { - self.splits - .iter() - .filter_map(|rows| { - let (scope, root) = plan_selection( - &self.plans, - WorkScope { - file_ordinal: 0, - rows: rows.clone(), - }, - &self.selection, - )?; - Some(SplitPlan { scope, root }) - }) - .collect() - } - - /// Runs every split as a task on the session's runtime and yields their mapped batches. + /// Runs the scan as one task on the session's runtime and yields its mapped batches. /// - /// Like the default executor, up to `concurrency` splits per thread run at once, and their - /// batches are yielded in row order when the scan is ordered. + /// The task owns one planning run, admits the scan's first planner to it, and drives it: + /// the run advances the planners and morsels the scan emits and awaits their reads between + /// steps, on whichever thread polls the task. 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 concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); let handle = self.session.handle(); - let split_plans = self.split_plans(); - let io = self.io; - let map_fn = self.map_fn; - if !self.ordered - && split_plans.iter().any(|split| { - split.scope.rows.end - split.scope.rows.start > MAX_ROWS_PER_CHUNK as u64 - }) - { - let streams = stream::iter(split_plans).map(move |split| { - let (send, recv) = mpsc::channel(1); - let task = - handle.spawn(split.run_unordered(Arc::clone(&io), Arc::clone(&map_fn), send)); - let errors = stream::once(task) - .filter_map(|result| std::future::ready(result.err().map(Err))); - stream::select(recv, errors) - }); - return Ok(streams.flatten_unordered(concurrency).boxed()); + let (send, recv) = mpsc::channel(1); + let task = handle.spawn(self.run(send)); + let errors = stream::once(task).filter_map(|result| future::ready(result.err().map(Err))); + Ok(stream::select(recv, errors).boxed()) + } + + async fn run(self, mut send: mpsc::Sender>) -> VortexResult<()> { + let mut run = Run::new(); + run.admit( + self.scan.planner(self.rows.clone(), self.selection), + WorkScope { + file_ordinal: 0, + rows: self.rows.clone(), + }, + self.io.session(), + ); + let mut held = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) if self.ordered => held.push(batch), + Progress::Batch(batch) => { + let batch = (self.map_fn)(batch.array); + let failed = batch.is_err(); + if send.send(batch).await.is_err() || failed { + return Ok(()); + } + } + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } } - // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. - let tasks = stream::iter(split_plans).map(move |split| { - let io = Arc::clone(&io); - let map_fn = Arc::clone(&map_fn); - handle.spawn(async move { - split - .run(io) - .await? - .into_iter() - .map(|array| map_fn(array)) - .collect::>>() - }) - }); - let batches = if self.ordered { - tasks.buffered(concurrency).boxed() - } else { - tasks.buffer_unordered(concurrency).boxed() - }; - Ok(batches - .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) - .try_flatten() - .boxed()) + // Morsels finish in whatever order their reads arrive. + held.sort_by_key(|batch| batch.scope.rows.start); + for batch in held { + let batch = (self.map_fn)(batch.array); + let failed = batch.is_err(); + if send.send(batch).await.is_err() || failed { + return Ok(()); + } + } + Ok(()) } } diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs deleted file mode 100644 index 624a775ad61..00000000000 --- a/vortex-layout/src/scan/v2/split.rs +++ /dev/null @@ -1,74 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::future; -use std::sync::Arc; - -use futures::SinkExt; -use futures::channel::mpsc; -use vortex_array::ArrayRef; -use vortex_error::VortexResult; -use vortex_io::request::IoService; -use vortex_scan::planning::driver::Progress; -use vortex_scan::planning::driver::Run; -use vortex_scan::planning::planner::Planner; -use vortex_scan::planning::planner::WorkScope; - -/// One split of a scan, ready to be admitted to a planning run. -/// -/// The split's planner projects the rows of `scope`, and its morsels emit one batch per -/// projection split with selected rows, in whatever order their reads arrive. -pub(super) struct SplitPlan { - /// The rows of the scan's root row domain the split covers. - pub(super) scope: WorkScope, - /// The split's first planner. - pub(super) root: Box, -} - -impl SplitPlan { - /// Publishes mapped batches as they finish, bounding how far projection runs ahead of its - /// consumer. Unordered scans do not need to retain the whole split to sort its output. - pub(super) async fn run_unordered( - self, - io: Arc, - map: Arc VortexResult + Send + Sync>, - mut send: mpsc::Sender>, - ) -> VortexResult<()> { - let mut run = Run::new(); - run.admit(self.root, self.scope, io.session()); - loop { - match run.advance()? { - Progress::Batch(batch) => { - let batch = map(batch.array); - let failed = batch.is_err(); - if send.send(batch).await.is_err() || failed { - return Ok(()); - } - } - Progress::RootDone(_) | Progress::Idle => return Ok(()), - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - } - - /// Runs the split on a planning run of its own, reading through a session of `io`, and - /// returns its batches in row order. - /// - /// The run advances inside this future, on whichever thread polls it, and the future awaits - /// its reads between steps. Nothing is handed to another thread. - pub(super) async fn run(self, io: Arc) -> VortexResult> { - let mut run = Run::new(); - run.admit(self.root, self.scope, io.session()); - let mut batches = Vec::new(); - loop { - match run.advance()? { - Progress::Batch(batch) => batches.push(batch), - Progress::RootDone(_) | Progress::Idle => break, - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - // Morsels finish in whatever order their reads arrive. - batches.sort_by_key(|batch| batch.scope.rows.start); - Ok(batches.into_iter().map(|batch| batch.array).collect()) - } -} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index ba29376bafb..2dee182544f 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -43,18 +43,24 @@ use vortex_arrow::ArrowSessionExt; 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_scan::planning::driver::Batch; +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::WorkScope; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use super::MAX_ROWS_PER_CHUNK; -use super::PreparedScan; use super::ROW_BLOCK_SIZE; use super::ScanBuilder; use super::ScanCache; use super::ScanFile; -use super::row_chunks; +use super::ScanRun; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; @@ -66,6 +72,7 @@ use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; use crate::scan::execution::SegmentLocation; +use crate::scan::execution::row_chunks; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; use crate::segments::SegmentFuture; use crate::segments::SegmentId; @@ -188,6 +195,45 @@ fn builder( )) } +/// 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> { + let mut run = Run::new(); + run.admit( + scan.scan.planner(rows.clone(), scan.selection.clone()), + WorkScope { + file_ordinal: 0, + rows, + }, + scan.io.session(), + ); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => std::future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + batches.sort_by_key(|batch| batch.scope.rows.start); + 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> { + Ok(run_batches(scan, rows) + .await? + .into_iter() + .map(|batch| batch.array) + .collect()) +} + async fn collect( dtype: DType, stream: impl futures::Stream>, @@ -214,8 +260,7 @@ async fn assert_same_as_v1( .with_some_filter(builder.filter.clone()) .with_selection(builder.selection.clone()) .with_row_offset(builder.row_offset) - .with_ordered(builder.ordered) - .with_concurrency(builder.concurrency); + .with_ordered(builder.ordered); if let Some(rows) = builder.row_range.clone() { v1 = v1.with_row_range(rows); } @@ -469,27 +514,26 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.split_plans(); + let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + let chunks = row_chunks(scan.rows.clone(), scan.scan.chunk_rows); assert_eq!( - splits.len(), + chunks.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) ); - for (index, split) in splits.into_iter().enumerate() { - let rows = split.scope.rows.clone(); + for (index, rows) in chunks.into_iter().enumerate() { let before = counted.reads.load(Ordering::Relaxed); - let arrays = split.run(Arc::clone(&scan.io)).await?; + let arrays = run_rows(&scan, rows.clone()).await?; let reads = counted.reads.load(Ordering::Relaxed) - before; let code_chunks = (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; - // The first split reads every values chunk once for the whole scan, plus its codes. - let code_chunks = usize::try_from(code_chunks)?; + // The first chunk reads every values chunk once for the whole scan; its codes are one of + // them, and the session shares that read. Later chunks read only their codes. assert_eq!( reads, if index == 0 { - STORED_CHUNKS + code_chunks + STORED_CHUNKS } else { - code_chunks + usize::try_from(code_chunks)? } ); let expected = @@ -535,7 +579,7 @@ async fn projection_skips_unsupported_unselected_layouts() -> VortexResult<()> { async fn unordered_returns_every_row() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; + let scan = ScanRun::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; let chunks: Vec = scan.execute_stream()?.try_collect().await?; let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); assert_eq!(rows, CHUNKS * CHUNK_ROWS); @@ -603,10 +647,11 @@ async fn unordered_yields_before_the_split_finishes() -> VortexResult<()> { blocked, release: gate.shared(), }); - let mut scan = - PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; - // Put both projection morsels in one split, as a dense filtered scan would. - scan.splits = std::iter::once(0..layout.row_count()).collect(); + 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(()); @@ -629,7 +674,7 @@ 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 = PreparedScan::try_new( + let scan = ScanRun::try_new( builder(&session, &segments, &layout)? .with_projection(get_item("ints", root()).bind(&dtype)?) .with_row_range(990..1010), @@ -655,8 +700,11 @@ 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 = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - assert_eq!(scan.splits, vec![0..(CHUNKS * CHUNK_ROWS) as u64]); + let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + assert_eq!( + row_chunks(scan.rows.clone(), scan.scan.chunk_rows), + vec![0..(CHUNKS * CHUNK_ROWS) as u64] + ); let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; assert_eq!(arrays.len(), 1); assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); @@ -734,13 +782,18 @@ 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 = PreparedScan::try_new( + let scan = ScanRun::try_new( builder(&session, &segments, &layout)? .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 2001]))?), )?; - let splits = scan.split_plans(); - assert_eq!(splits.len(), 1); - assert_eq!(splits[0].scope.rows, 7..2002); + let mut chunks = scan + .scan + .chunk_planner(scan.rows.clone(), scan.selection.clone(), None); + let PlannerOutput::Planner(scope, _) = chunks.compute()? else { + vortex_bail!("the selected rows were not planned"); + }; + assert_eq!(scope.rows, 7..2002); + assert!(matches!(chunks.compute()?, PlannerOutput::Done)); Ok(()) } @@ -801,3 +854,31 @@ fn partition_rows_tile_the_file_on_row_blocks() { "a range inside one block owns no rows" ); } + +/// Zone proofs run once before any chunk is planned: rows every proof 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))] +#[case::nothing(10_000, None)] +#[tokio::test] +async fn pruned_rows_get_no_chunk( + #[case] above: i32, + #[case] expected: Option>, +) -> 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?; + assert_eq!( + batches + .iter() + .map(|batch| batch.scope.rows.clone()) + .collect::>(), + expected.into_iter().collect::>() + ); + Ok(()) +} From 8e106ac55d768a0d188cbd4adadb838c3b3c5d71 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 10:14:49 +0100 Subject: [PATCH 02/23] Optimize a zoned data plan to its data child A scan reads a zoned column through the zoned node's data child and never its zone table, yet the node and its zone table stayed in every optimized data plan. A new `ZonedDataRule` replaces a zoned data plan by its data, so an optimized plan carries a zone table only where a pruning proof reads it. Legacy stats get the same treatment through `LegacyStatsRule`; the rule that dropped legacy stats beneath a zoned plan is subsumed. A zone proof over legacy stats depended on the converted zoned node surviving child optimization. `ExpressionLegacyStatsRule` now converts the legacy child under a proof directly, through the shared `is_zone_proof` check, so the pruning rewrite still meets a zoned plan. With zone tables gone from data plans, the identity projection is optimized and cached like every other projection: chunked columns stay lazy through `Concat::map_chunks`, and a projection-only scan no longer lowers the zone tables it never reads. (cherry picked from commit 91be25d4ced87c3d7fddc45e5c6b5ed202f24ff1) Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- .../src/plan/physical/optimizer/mod.rs | 18 ++- .../src/plan/physical/plans/legacy_stats.rs | 31 +++-- vortex-layout/src/plan/physical/plans/mod.rs | 3 +- .../src/plan/physical/plans/zoned.rs | 59 ++++++--- vortex-layout/src/plan/physical/tests.rs | 119 +++++++++++------- vortex-layout/src/scan/v2/mod.rs | 19 ++- 6 files changed, 154 insertions(+), 95 deletions(-) diff --git a/vortex-layout/src/plan/physical/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs index c4841740572..5159067f57e 100644 --- a/vortex-layout/src/plan/physical/optimizer/mod.rs +++ b/vortex-layout/src/plan/physical/optimizer/mod.rs @@ -24,11 +24,12 @@ use super::Take; use super::Zoned; use super::plans::EvalIdentityRule; use super::plans::ExpressionConcatRule; +use super::plans::ExpressionLegacyStatsRule; use super::plans::ExpressionPackRule; use super::plans::ExpressionTakeRule; use super::plans::ExpressionZonedRule; use super::plans::LegacyStatsRule; -use super::plans::ZonedLegacyStatsRule; +use super::plans::ZonedDataRule; static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(EvalIdentityRule); @@ -36,7 +37,11 @@ static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter = static LEGACY_STATS_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(LegacyStatsRule); -static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[&EVAL_IDENTITY_RULE, &LEGACY_STATS_RULE]); +static ZONED_DATA_RULE: PlanReduceRuleAdapter = + PlanReduceRuleAdapter::new(ZonedDataRule); + +static PLAN_RULES: PlanRuleSet = + PlanRuleSet::new(&[&EVAL_IDENTITY_RULE, &LEGACY_STATS_RULE, &ZONED_DATA_RULE]); static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter = PlanParentReduceRuleAdapter::new(ExpressionConcatRule); @@ -46,16 +51,17 @@ static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter = PlanParentReduceRuleAdapter::new(ExpressionZonedRule); - -static ZONED_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter = - PlanParentReduceRuleAdapter::new(ZonedLegacyStatsRule); +static EXPRESSION_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter< + LegacyStats, + ExpressionLegacyStatsRule, +> = PlanParentReduceRuleAdapter::new(ExpressionLegacyStatsRule); static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[ &EXPRESSION_CONCAT_RULE, &EXPRESSION_TAKE_RULE, &EXPRESSION_PACK_RULE, &EXPRESSION_ZONED_RULE, - &ZONED_LEGACY_STATS_RULE, + &EXPRESSION_LEGACY_STATS_RULE, ]); /// Attempts a static rewrite for `plan`. diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index 7eb6bec85ff..4abd2932509 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/vortex-layout/src/plan/physical/plans/legacy_stats.rs @@ -15,6 +15,7 @@ use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::layouts::zoned::legacy_stats_to_aggregates; +use crate::plan::physical::Eval; use crate::plan::physical::EvalPlan; use crate::plan::physical::Plan; use crate::plan::physical::PlanChildren; @@ -29,6 +30,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; const DATA: usize = 0; const ZONES: usize = 1; @@ -37,8 +39,9 @@ const ZONES: usize = 1; /// than by aggregate function. /// /// The optimizer rewrites it into a [`Zoned`] plan over the zone table converted to the aggregate -/// schema, or into its data alone when a [`Zoned`] plan already summarises the same data. Compiled -/// unoptimized, it reads its data and prunes nothing. +/// schema, or into its data alone when a [`Zoned`] plan already summarises the same data; a plan +/// that reads the data is then optimized to the data alone. Compiled unoptimized, it reads its +/// data and prunes nothing. #[derive(Clone, Debug)] pub struct LegacyStats; @@ -173,24 +176,28 @@ impl PlanReduceRule for LegacyStatsRule { } } -/// Drops legacy stats beneath a [`Zoned`] plan: its aggregate zone map already summarises the -/// same rows. +/// Converts the legacy stats beneath a zone proof in place, so the proof meets a [`Zoned`] plan +/// before the data path optimizes the stats away. #[derive(Debug)] -pub(crate) struct ZonedLegacyStatsRule; +pub(crate) struct ExpressionLegacyStatsRule; -impl PlanParentReduceRule for ZonedLegacyStatsRule { - type Parent = Zoned; +impl PlanParentReduceRule for ExpressionLegacyStatsRule { + type Parent = Eval; fn reduce_parent( &self, child: &LegacyStatsPlan, - parent: &ZonedPlan, - child_idx: usize, + parent: &Plan, + _child_idx: usize, ) -> VortexResult> { - if child_idx != DATA || parent.is_pruning() { + if !is_zone_proof(parent)? { return Ok(None); } - let children = vec![child.data_plan()?, parent.zones_plan()?]; - Ok(Some(parent.clone().into_plan().with_children(children)?)) + let Some(converted) = LegacyStatsRule.reduce(child)? else { + return Ok(None); + }; + Ok(Some( + EvalPlan::try_new(parent.expression().clone(), converted)?.into_plan(), + )) } } diff --git a/vortex-layout/src/plan/physical/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs index fe5a4eae225..3a215e5b7da 100644 --- a/vortex-layout/src/plan/physical/plans/mod.rs +++ b/vortex-layout/src/plan/physical/plans/mod.rs @@ -24,11 +24,11 @@ pub use eval::EvalPlan; pub use filter::Filter; pub use filter::FilterData; pub use filter::FilterPlan; +pub(crate) use legacy_stats::ExpressionLegacyStatsRule; pub use legacy_stats::LegacyStats; pub use legacy_stats::LegacyStatsData; pub use legacy_stats::LegacyStatsPlan; pub(crate) use legacy_stats::LegacyStatsRule; -pub(crate) use legacy_stats::ZonedLegacyStatsRule; pub use list_pack::ListPack; pub use list_pack::ListPackData; pub use list_pack::ListPackPlan; @@ -53,4 +53,5 @@ pub use take::TakePlan; pub(crate) use zoned::ExpressionZonedRule; pub use zoned::Zoned; pub use zoned::ZonedData; +pub(crate) use zoned::ZonedDataRule; pub use zoned::ZonedPlan; diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index d1a17702ebf..149b9c2b679 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -46,6 +46,7 @@ use crate::plan::physical::operators::OnceSource; use crate::plan::physical::operators::UNBOUNDED; use crate::plan::physical::operators::WholeSource; use crate::plan::physical::optimizer::PlanParentReduceRule; +use crate::plan::physical::optimizer::PlanReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; @@ -151,7 +152,8 @@ pub struct ZonedData { /// Its zone table has the aggregate schema of `vortex.zoned` layouts; legacy `vortex.stats` /// layouts reach it through [`LegacyStats`](crate::plan::physical::LegacyStats), which converts /// their tables. An expression containing abstract statistic functions can rewrite it into a -/// pruning plan that retains only the zone-statistics child. +/// pruning plan that retains only the zone-statistics child; a plan that reads the data is +/// optimized to its data child alone. #[derive(Clone, Debug)] pub struct Zoned; @@ -453,29 +455,46 @@ impl PlanParentReduceRule for ExpressionZonedRule { parent: &Plan, _child_idx: usize, ) -> VortexResult> { - let mut contains_stat = false; - let mut contains_root = false; - parent.expression().clone().transform_down(|expression| { - if expression - .as_scalar() - .is_some_and(|scalar_fn| scalar_fn.is::()) - { - contains_stat = true; - return Ok(Transformed { - value: expression, - order: TraversalOrder::Skip, - changed: false, - }); - } - contains_root |= expression.is_root(); - Ok(Transformed::no(expression)) - })?; - if !parent.dtype().is_boolean() || !contains_stat || contains_root { + if !is_zone_proof(parent)? { return Ok(None); } - Ok(child .with_pruning(parent.expression().clone())? .map(Plan::into_plan)) } } + +/// 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 { + let mut contains_stat = false; + let mut contains_root = false; + eval.expression().clone().transform_down(|expression| { + if expression + .as_scalar() + .is_some_and(|scalar_fn| scalar_fn.is::()) + { + contains_stat = true; + return Ok(Transformed { + value: expression, + order: TraversalOrder::Skip, + changed: false, + }); + } + contains_root |= expression.is_root(); + Ok(Transformed::no(expression)) + })?; + Ok(eval.dtype().is_boolean() && contains_stat && !contains_root) +} + +/// 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. +#[derive(Debug)] +pub(crate) struct ZonedDataRule; + +impl PlanReduceRule for ZonedDataRule { + fn reduce(&self, plan: &Plan) -> VortexResult> { + plan.data_plan() + } +} diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 44705753b3e..7eefcfcba2d 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -58,6 +58,9 @@ use crate::layouts::zoned::MAX_IS_TRUNCATED; use crate::layouts::zoned::ZonedLayout; use crate::layouts::zoned::legacy_stats_to_aggregates; use crate::layouts::zoned::zone_map::ZoneMap; +use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; +use crate::plan::physical::pipeline::Scan; use crate::segments::SegmentId; use crate::test::SESSION; @@ -1306,14 +1309,37 @@ fn stats_expression_rewrites_to_zoned_pruning_plan() -> VortexResult<()> { let mixed = mixed .as_opt::() .ok_or_else(|| vortex_err!("expression with a data reference was pushed into zones"))?; - let mixed_child = mixed.child_plan()?; - assert!(mixed_child.is::()); - assert!( - !mixed_child - .as_opt::() - .ok_or_else(|| vortex_err!("mixed expression child is not zoned"))? - .is_pruning() - ); + assert!(mixed.child_plan()?.is::()); + Ok(()) +} + +/// A plan that reads a zoned column's data is optimized to the data alone: the zone table is +/// neither a child of the plan nor among the segments it reads. +#[test] +fn data_expression_strips_zoned_to_its_data() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let (max, zones_dtype) = max_zones_dtype(&dtype)?; + let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?; + let layout = ZonedLayout::try_new( + flat(5, dtype.clone(), 0), + flat(2, zones_dtype, 1), + zone_len, + vec![max].into(), + )? + .into_layout(); + let source = make_plan(layout)?; + + let plan = EvalPlan::try_new(gt(root(), lit(4_i32)).bind(&dtype)?, source.clone())?.into_plan(); + insta::assert_snapshot!(optimize(plan)?.display_tree(), @" + root: vortex.plan.eval(bool, rows=5) expr=($ > 4i32) + child: vortex.plan.filter(i32, rows=5) + child: vortex.plan.segment_scan(i32, rows=5) + "); + + let identity = optimize(source)?; + assert!(identity.is::()); + let segments = Scan::segments(&SESSION, &identity, 0..5, &Mask::new_true(5))?; + assert_eq!(segments, vec![SegmentId::from(0)]); Ok(()) } @@ -1483,23 +1509,8 @@ fn legacy_stats_rewrite_to_zoned_over_converted_zones() -> VortexResult<()> { let (max, zones_dtype) = max_zones_dtype(&dtype)?; let legacy = legacy_max_plan(&dtype)?; - let optimized = optimize(legacy.clone())?; - let zoned = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("legacy stats did not become zoned"))?; - assert!(!zoned.is_pruning()); - assert_eq!(zoned.zone_len(), 3); - assert_eq!( - zoned - .aggregate_fns() - .iter() - .map(ToString::to_string) - .collect::>(), - vec![max.to_string()] - ); - let zones = zoned.zones_plan()?; - assert!(zones.is::()); - assert_eq!(zones.dtype(), &zones_dtype); + // Read as data, the legacy stats are dropped along with the zoned plan they convert to. + assert!(optimize(legacy.clone())?.is::()); let falsifier = gt(root(), lit(5_i32)) .bind(&dtype)? @@ -1510,7 +1521,18 @@ fn legacy_stats_rewrite_to_zoned_over_converted_zones() -> VortexResult<()> { .as_opt::() .ok_or_else(|| vortex_err!("legacy stats predicate did not become a pruning plan"))?; assert!(pruning.is_pruning()); - assert!(pruning.zones_plan()?.is::()); + assert_eq!(pruning.zone_len(), 3); + assert_eq!( + pruning + .aggregate_fns() + .iter() + .map(ToString::to_string) + .collect::>(), + vec![max.to_string()] + ); + let zones = pruning.zones_plan()?; + assert!(zones.is::()); + assert_eq!(zones.dtype(), &zones_dtype); Ok(()) } @@ -1549,43 +1571,52 @@ fn converted_legacy_zones_prune() -> VortexResult<()> { Ok(()) } +/// A zoned plan over legacy stats read as data keeps neither zone table. #[test] -fn legacy_stats_under_zoned_are_dropped() -> VortexResult<()> { +fn zoned_over_legacy_stats_read_as_data_keeps_only_the_data() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); let (max, zones_dtype) = max_zones_dtype(&dtype)?; let zones = make_plan(flat(2, zones_dtype, 2))?; - let plan = - ZonedPlan::new(legacy_max_plan(&dtype)?, zones.clone(), 3, vec![max].into()).into_plan(); + let legacy = legacy_max_plan(&dtype)?; + let data = legacy + .as_opt::() + .ok_or_else(|| vortex_err!("legacy plan is not legacy stats"))? + .data_plan()?; + let plan = ZonedPlan::new(legacy, zones, 3, vec![max].into()).into_plan(); - let optimized = optimize(plan)?; - let zoned = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("zoned plan was rewritten away"))?; - assert!(zoned.data_plan()?.is_some_and(|data| data.is::())); - assert!(PlanRef::ptr_eq(&zoned.zones_plan()?, &zones)); + assert!(PlanRef::ptr_eq(&optimize(plan)?, &data)); Ok(()) } +/// 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. #[test] fn legacy_stats_over_zoned_are_dropped() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); let (max, zones_dtype) = max_zones_dtype(&dtype)?; - let zoned = ZonedPlan::new( - make_plan(flat(5, dtype.clone(), 0))?, - make_plan(flat(2, zones_dtype, 2))?, - 3, - vec![max].into(), - ) - .into_plan(); + let data = make_plan(flat(5, dtype.clone(), 0))?; + let zones = make_plan(flat(2, zones_dtype, 2))?; + let zoned = ZonedPlan::new(data.clone(), zones.clone(), 3, vec![max].into()).into_plan(); let legacy = LegacyStatsPlan::new( - zoned.clone(), + zoned, make_plan(flat(2, legacy_max_zones_dtype(&dtype), 1))?, 3, vec![Stat::Max].into(), ) .into_plan(); - assert!(PlanRef::ptr_eq(&optimize(legacy)?, &zoned)); + assert!(PlanRef::ptr_eq(&optimize(legacy.clone())?, &data)); + + let falsifier = gt(root(), lit(5_i32)) + .bind(&dtype)? + .falsify(&SESSION)? + .ok_or_else(|| vortex_err!("filter has no falsifier"))?; + let pruning = optimize(EvalPlan::try_new(falsifier, legacy)?.into_plan())?; + let pruning = pruning + .as_opt::() + .ok_or_else(|| vortex_err!("predicate over legacy stats did not become a pruning plan"))?; + assert!(pruning.is_pruning()); + assert!(PlanRef::ptr_eq(&pruning.zones_plan()?, &zones)); Ok(()) } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index d0b28a60c10..8e898066a3b 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -263,19 +263,14 @@ impl ScanRun { .filter .map(|filter| FilterPlans::new(filter, &root, &builder.session).map(Arc::new)) .transpose()?; - // An identity projection needs no rewrites. Keep the lowered children lazy, including - // statistics that a projection-only scan never reads. + // 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 = if identity { - root - } else { - match cache.projection(&builder.projection) { - Some(plan) => plan, - None => { - let plan = - optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; - cache.insert_projection(builder.projection, plan) - } + let projection = match cache.projection(&builder.projection) { + Some(plan) => plan, + None => { + let plan = optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; + cache.insert_projection(builder.projection, plan) } }; let chunk_rows = if filter.is_some() { From 38f6a81747b60ee87e019ac45eb7b0fa6a441df9 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 12:32:13 +0100 Subject: [PATCH 03/23] Share one V2 scan cache per file across DataFusion partitions DataFusion opens a file once per partition, so each partition lowered the layout, planned the filter and proved the zone maps on its own. A scan cache now keeps each distinct filter's plans along with the zones their proofs excluded, and the opener keeps one cache per file that concurrent partitions find through a weak handle and that the last scan holding it frees. A prune planner skips the zone maps of a proof an earlier scan of the file already proved: they are not read again. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-datafusion/src/persistent/opener.rs | 35 +++++++++++++ vortex-datafusion/src/persistent/source.rs | 4 ++ .../src/plan/physical/plans/zoned.rs | 9 ++++ vortex-layout/src/scan/execution/prune.rs | 27 +++++++++- vortex-layout/src/scan/v2/cache.rs | 52 ++++++++++++++++++- vortex-layout/src/scan/v2/mod.rs | 21 ++++++-- vortex-layout/src/scan/v2/tests.rs | 37 +++++++++++++ 7 files changed, 180 insertions(+), 5 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 71dfb22ffc1..cd17179b410 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -56,6 +56,8 @@ use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; 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::metrics::Label; @@ -112,6 +114,9 @@ pub(crate) struct VortexOpener { /// boundary and the file's metadata segments are read once, as they are for V1 through the /// 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. + pub scan_caches: Arc>, /// Shared V1 full-file natural splits keyed by file path. pub natural_splits: Arc>>, /// Shared physical Arrow schemas of each file, keyed by file path. @@ -158,6 +163,7 @@ impl FileOpener for VortexOpener { let limit = self.limit; let layout_readers = Arc::clone(&self.layout_readers); let segment_sources = Arc::clone(&self.segment_sources); + let scan_caches = Arc::clone(&self.scan_caches); let natural_splits = Arc::clone(&self.natural_splits); let file_schemas = Arc::clone(&self.file_schemas); let has_output_ordering = self.has_output_ordering; @@ -423,6 +429,7 @@ impl FileOpener for VortexOpener { &file.object_meta.location, scan_file.segments, ); + 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) @@ -643,6 +650,28 @@ fn contains_lambda(projection: &ProjectionExprs) -> DFResult { Ok(false) } +/// Returns the scan cache other partitions of `path` are planning with, or registers a new one. +/// A scan holds its cache until its stream ends, so the cache lives as long as one scan. +fn shared_scan_cache(scan_caches: &DashMap, path: &Path) -> ScanCache { + match scan_caches.entry(path.clone()) { + Entry::Occupied(mut entry) => { + if let Some(shared) = entry.get().upgrade() { + tracing::trace!("reusing scan cache for {path}"); + shared + } else { + let fresh = ScanCache::default(); + entry.insert(fresh.downgrade()); + fresh + } + } + Entry::Vacant(entry) => { + let fresh = ScanCache::default(); + entry.insert(fresh.downgrade()); + fresh + } + } +} + /// Returns the segment source other partitions of `path` are reading through, or registers /// `fresh` as that source. Like the layout readers, a source lives only as long as one scan. fn shared_segment_source( @@ -1074,6 +1103,7 @@ mod tests { df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), segment_sources: Default::default(), + scan_caches: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1482,6 +1512,7 @@ mod tests { df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), segment_sources: Default::default(), + scan_caches: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1572,6 +1603,7 @@ mod tests { df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), segment_sources: Default::default(), + scan_caches: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1753,6 +1785,7 @@ mod tests { df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), segment_sources: Default::default(), + scan_caches: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1816,6 +1849,7 @@ mod tests { df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), segment_sources: Default::default(), + scan_caches: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -2030,6 +2064,7 @@ mod tests { df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), segment_sources: Default::default(), + scan_caches: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 5a86abee50c..8ff7b735cee 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -34,6 +34,7 @@ use object_store::ObjectStore; use object_store::path::Path; use vortex::file::VORTEX_FILE_EXTENSION; use vortex::layout::LayoutReader; +use vortex::layout::scan::v2::WeakScanCache; use vortex::layout::segments::SegmentSource; use vortex::metrics::DefaultMetricsRegistry; use vortex::metrics::MetricsRegistry; @@ -203,6 +204,7 @@ pub struct VortexSource { layout_readers: Arc>>, /// Shared V2 segment sources, keyed by path and alive only while a scan reads the file. segment_sources: Arc>>, + scan_caches: Arc>, /// Shared full-file natural splits keyed by path. natural_splits: Arc>>, /// Shared physical Arrow schemas of each file, keyed by path. @@ -240,6 +242,7 @@ impl VortexSource { df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), segment_sources: Arc::new(DashMap::default()), + scan_caches: Arc::new(DashMap::default()), natural_splits: Arc::new(DashMap::default()), file_schemas: Arc::new(DashMap::default()), expression_converter, @@ -366,6 +369,7 @@ impl VortexSource { df_metrics: self.df_metrics.clone(), layout_readers: Arc::clone(&self.layout_readers), segment_sources: Arc::clone(&self.segment_sources), + scan_caches: Arc::clone(&self.scan_caches), natural_splits: Arc::clone(&self.natural_splits), file_schemas: Arc::clone(&self.file_schemas), has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered, diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index 149b9c2b679..8b1729e953e 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -241,6 +241,15 @@ 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 8545f0b8f13..9bda28e035d 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -19,6 +19,7 @@ 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_session::VortexSession; use super::PreparedScan; use super::adaptor::AdaptorStep; @@ -27,9 +28,25 @@ 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) +} + /// Runs the filter's zone-map proofs over the scan's rows and emits one chunk planner for the /// rows no proof excludes. The zone maps are the only reads, started together. /// @@ -98,9 +115,17 @@ 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); + } + } let segments = segments_read( &self.scan.session, - self.proofs.as_slice().iter(), + unproved.into_iter(), self.rows.clone(), &self.mask, )?; diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 3daa0975a0e..5a26e28ec97 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -4,6 +4,7 @@ //! State derived from a [`ScanFile`](super::ScanFile) that repeated V2 scans can share. use std::sync::Arc; +use std::sync::Weak; use parking_lot::Mutex; use rustc_hash::FxHashMap; @@ -13,6 +14,7 @@ use vortex_array::expr::ExactBoundExpr; use super::io::SegmentRanges; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; +use crate::scan::execution::FilterPlans; /// Reusable state for repeated V2 scans of one file, the V2 counterpart of a cached /// [`LayoutReader`](crate::LayoutReader) tree. @@ -21,7 +23,9 @@ use crate::plan::physical::pipeline::ValueCache; /// 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, and dictionary values stay keyed by the cached plan that produced them. +/// 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 +/// that produced them. /// /// A cache describes one file's layout, segment locations, and segment source. Use a new cache /// when any of them changes. Cached plans and values live as long as the cache, one set per @@ -36,6 +40,9 @@ 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. + filters: Mutex>>, values: ValueCache, } @@ -80,7 +87,50 @@ impl ScanCache { Arc::clone(self.0.ranges.lock().get_or_insert(ranges)) } + /// The plans of `filter` an earlier scan kept. + pub(super) fn filter_plans(&self, filter: &BoundExpression) -> Option> { + self.0 + .filters + .lock() + .get(&ExactBoundExpr(filter.clone())) + .cloned() + } + + /// Keeps the plans of `filter`, returning the plans already kept if another scan got there + /// first. + pub(super) fn insert_filter_plans( + &self, + filter: BoundExpression, + plans: Arc, + ) -> Arc { + Arc::clone( + self.0 + .filters + .lock() + .entry(ExactBoundExpr(filter)) + .or_insert(plans), + ) + } + pub(super) fn values(&self) -> ValueCache { self.0.values.clone() } + + /// A handle that finds this cache while a scan holds it, without keeping it alive itself. + pub fn downgrade(&self) -> WeakScanCache { + WeakScanCache(Arc::downgrade(&self.0)) + } +} + +/// A [`ScanCache`] handle that does not keep the cache alive. Engines that open a file once per +/// partition keep one per file, so concurrent partitions share a cache that is freed with the +/// last scan holding it. +#[derive(Clone)] +pub struct WeakScanCache(Weak); + +impl WeakScanCache { + /// The cache, while a scan still holds it. + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(ScanCache) + } } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 8e898066a3b..194c88e74f7 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -15,8 +15,9 @@ //! 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, segment index, and dictionary -//! values with every other scan of that file, like a cached layout reader does for V1 scans. +//! 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; @@ -54,6 +55,7 @@ use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; pub use self::cache::ScanCache; +pub use self::cache::WeakScanCache; use crate::LayoutRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; @@ -246,6 +248,8 @@ struct ScanRun { selection: Selection, map_fn: Arc VortexResult + Send + Sync>, ordered: bool, + /// Held until the scan ends, so scans of the file that start meanwhile share it. + cache: ScanCache, } impl ScanRun { @@ -261,7 +265,16 @@ impl ScanRun { let row_count = root.row_count(); let filter = builder .filter - .map(|filter| FilterPlans::new(filter, &root, &builder.session).map(Arc::new)) + .map(|filter| -> VortexResult<_> { + match cache.filter_plans(&filter) { + Some(plans) => Ok(plans), + None => { + let plans = + Arc::new(FilterPlans::new(filter.clone(), &root, &builder.session)?); + Ok(cache.insert_filter_plans(filter, plans)) + } + } + }) .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. @@ -301,6 +314,7 @@ impl ScanRun { selection: builder.selection, map_fn: builder.map_fn, ordered: builder.ordered, + cache, }) } @@ -319,6 +333,7 @@ impl ScanRun { } async fn run(self, mut send: mpsc::Sender>) -> VortexResult<()> { + let _cache = &self.cache; let mut run = Run::new(); run.admit( self.scan.planner(self.rows.clone(), self.selection), diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 2dee182544f..21980158b72 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -882,3 +882,40 @@ 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: +/// 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<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, zoned_table().as_ref()).await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments = Arc::clone(&counted) as Arc; + let mut file = scan_file(&segments, &layout)?; + let cache = ScanCache::default(); + file.cache = Some(cache.clone()); + let filter = gt(get_item("ints", root()), lit(2_999i32)).bind(layout.dtype())?; + let projection = get_item("ints", root()).bind(layout.dtype())?; + assert!(cache.filter_plans(&filter).is_none()); + + let mut reads = Vec::new(); + for _ in 0..2 { + let before = counted.reads.load(Ordering::Relaxed); + 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?; + assert_eq!(actual.len(), 1_000); + reads.push(counted.reads.load(Ordering::Relaxed) - before); + } + 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!(reads, vec![2, 1]); + Ok(()) +} From 6894be051278f2425a0e7b36ad0892c8c6c61409 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 13:35:22 +0100 Subject: [PATCH 04/23] Drive a V2 scan's run from its stream, with compute on spawn_cpu The stream a scan returns now owns the planning run and drives it as it is polled. Every CPU step of the run, the planners' and morsels' compute, is a `spawn_cpu` task; every read is already a `spawn_io` task of the segment source; the stream's poll awaits one or the other and does no scan work on the engine's thread. The scan advances only as fast as its consumer pulls, so the task and channel that ran ahead of it go. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-layout/src/scan/v2/mod.rs | 105 ++++++++++++++++++++----------- 1 file changed, 70 insertions(+), 35 deletions(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 194c88e74f7..f9b57684d33 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -24,6 +24,7 @@ mod io; #[cfg(test)] mod tests; +use std::collections::VecDeque; use std::future; use std::ops::Range; use std::sync::Arc; @@ -32,11 +33,9 @@ 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; @@ -45,8 +44,11 @@ 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::runtime::Handle; use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::driver::Batch; use vortex_scan::planning::driver::Progress; use vortex_scan::planning::driver::Run; use vortex_scan::planning::planner::WorkScope; @@ -318,22 +320,14 @@ impl ScanRun { }) } - /// Runs the scan as one task on the session's runtime and yields its mapped batches. + /// Runs the scan on a planning run the returned stream owns and drives as it is polled. /// - /// The task owns one planning run, admits the scan's first planner to it, and drives it: - /// the run advances the planners and morsels the scan emits and awaits their reads between - /// steps, on whichever thread polls the task. Batches of an unordered scan are yielded as - /// morsels finish; an ordered scan yields them in row order once every morsel has. + /// Every CPU step of the run, the planners' and morsels' compute, is a [`Handle::spawn_cpu`] + /// task; every read is a `spawn_io` task of the segment source; the stream's own poll only + /// awaits one or the other, so the engine thread polling it does no scan work itself and the + /// scan advances only as fast as its consumer pulls. 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 (send, recv) = mpsc::channel(1); - let task = handle.spawn(self.run(send)); - let errors = stream::once(task).filter_map(|result| future::ready(result.err().map(Err))); - Ok(stream::select(recv, errors).boxed()) - } - - async fn run(self, mut send: mpsc::Sender>) -> VortexResult<()> { - let _cache = &self.cache; let mut run = Run::new(); run.admit( self.scan.planner(self.rows.clone(), self.selection), @@ -343,30 +337,71 @@ impl ScanRun { }, self.io.session(), ); - let mut held = Vec::new(); + let driver = Driver { + handle: self.session.handle(), + run: Some(run), + map_fn: self.map_fn, + ordered: self.ordered, + held: VecDeque::new(), + done: false, + cache: self.cache, + }; + Ok(stream::try_unfold(driver, Driver::next).boxed()) + } +} + +/// The state of a scan's stream: its planning run and what it holds back. +struct Driver { + handle: Handle, + /// Taken while a CPU step runs on its task and put back with its result. + run: Option, + map_fn: Arc VortexResult + Send + Sync>, + ordered: bool, + /// Batches of an ordered scan, in row order once the run is done. + held: VecDeque, + done: bool, + /// Held until the scan ends, so scans of the file that start meanwhile share it. + #[expect(dead_code, reason = "kept alive, never read")] + cache: ScanCache, +} + +impl Driver { + /// The next mapped batch, or `None` once the run has retired everything. + async fn next(mut self) -> VortexResult> { loop { - match run.advance()? { - Progress::Batch(batch) if self.ordered => held.push(batch), + if self.done { + return match self.held.pop_front() { + Some(batch) => Ok(Some(((self.map_fn)(batch.array)?, self))), + None => Ok(None), + }; + } + let run = self + .run + .take() + .ok_or_else(|| vortex_err!("the scan's run is missing"))?; + let (run, progress) = self + .handle + .spawn_cpu(move || { + let mut run = run; + let progress = run.advance(); + (run, progress) + }) + .await; + let run = self.run.insert(run); + match progress? { + Progress::Batch(batch) if self.ordered => self.held.push_back(batch), Progress::Batch(batch) => { - let batch = (self.map_fn)(batch.array); - let failed = batch.is_err(); - if send.send(batch).await.is_err() || failed { - return Ok(()); - } + return Ok(Some(((self.map_fn)(batch.array)?, self))); + } + Progress::RootDone(_) | Progress::Idle => { + self.done = true; + // Morsels finish in whatever order their reads arrive. + self.held + .make_contiguous() + .sort_by_key(|batch| batch.scope.rows.start); } - Progress::RootDone(_) | Progress::Idle => break, Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, } } - // Morsels finish in whatever order their reads arrive. - held.sort_by_key(|batch| batch.scope.rows.start); - for batch in held { - let batch = (self.map_fn)(batch.array); - let failed = batch.is_err(); - if send.send(batch).await.is_err() || failed { - return Ok(()); - } - } - Ok(()) } } From 5a66bc1221c4970e21379de8f6d79882813206bd Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 13:35:22 +0100 Subject: [PATCH 05/23] Build a prune planner with the planner it continues into A prune planner no longer knows about chunks or selections. It is built with the planner it continues into, as a function of the rows its proofs keep, and emits that planner with the kept mask once every proof has run. `PreparedScan::planner` supplies the chunk planner. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-layout/src/scan/execution/mod.rs | 8 +++++- vortex-layout/src/scan/execution/prune.rs | 32 ++++++++++++----------- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index aa49bb38c3d..e371fdb539c 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -99,7 +99,13 @@ impl PreparedScan { .as_ref() .is_some_and(|filter| filter.proofs().next().is_some()); if proofs && !rows.is_empty() { - Box::new(PrunePlanner::new(Arc::clone(self), rows, selection)) + 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)) } diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 9bda28e035d..56e5808238e 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -9,6 +9,7 @@ use std::vec::IntoIter; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; use vortex_io::request::IoConsumer; use vortex_io::request::IoIntent; use vortex_io::request::IoRequestId; @@ -18,7 +19,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_scan::selection::Selection; use vortex_session::VortexSession; use super::PreparedScan; @@ -47,15 +47,17 @@ fn is_proved(plan: &PlanRef, session: &VortexSession) -> VortexResult { Ok(true) } -/// Runs the filter's zone-map proofs over the scan's rows and emits one chunk planner for the -/// rows no proof excludes. The zone maps are the only reads, started together. -/// -/// Rows a proof excludes are never planned: a chunk none of whose rows survive is skipped by -/// the chunk planner, and the chunks that remain run their conjuncts over the kept rows only. +/// The planner a prune planner continues into, given the rows the proofs kept. +pub(super) type Next = 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 { scan: Arc, rows: Range, - selection: Option, + /// What to emit once every proof has run, taken when it is. + next: Option, /// The rows no proof has excluded yet, over `rows`. mask: Mask, proofs: IntoIter, @@ -68,7 +70,7 @@ pub(super) struct PrunePlanner { } impl PrunePlanner { - pub(super) fn new(scan: Arc, rows: Range, selection: Selection) -> Self { + pub(super) fn new(scan: Arc, rows: Range, next: Next) -> Self { let proofs: Vec<_> = scan .filter .iter() @@ -78,7 +80,7 @@ impl PrunePlanner { Self { scan, rows, - selection: Some(selection), + next: Some(next), mask: Mask::new_true(len), proofs: proofs.into_iter(), active: None, @@ -184,19 +186,19 @@ impl Planner for PrunePlanner { self.active = Some(ScanAdaptor::new(scan, Arc::clone(&self.scan.locations))); return Ok(PlannerOutput::Continue); } - // Every proof has run: the chunks of the kept rows are planned from here on. + // Every proof has run: the planner this one continues into gets the kept rows. self.done = true; - let selection = self.selection.take().unwrap_or_default(); + let next = self + .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)); Ok(PlannerOutput::Planner( WorkScope { file_ordinal: 0, rows: self.rows.clone(), }, - Box::new( - self.scan - .chunk_planner(self.rows.clone(), selection, Some(kept)), - ), + next(kept), )) } } From a74502fa8ff247154f0814c56c636647d748d389 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 13:58:57 +0100 Subject: [PATCH 06/23] Spawn a scan's chunks as independent planners and run them beside each other A planner may now emit a child with `PlannerOutput::Spawn`: the run does not queue it but hands it to its owner as `Progress::Spawn`, and the planner continues at once. The owner runs it anywhere; the blocking driver admits it as another root of the same run. The chunk planner spawns every chunk. The scan's stream runs each on a run and a session of its own, stepped on `spawn_cpu`, up to `concurrency` chunks per worker thread at once, merging their batches as they finish. A scan with fewer engine partitions than cores had lost the parallelism of one task per split, 2 to 4 times slower on selective queries with one partition; it is back, without a planner knowing about threads. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-layout/src/scan/execution/chunk.rs | 10 +- vortex-layout/src/scan/v2/mod.rs | 204 +++++++++++++++------- vortex-layout/src/scan/v2/tests.rs | 6 +- vortex-scan/src/planning/driver/mod.rs | 14 +- vortex-scan/src/planning/driver/tests.rs | 1 + vortex-scan/src/planning/planner.rs | 6 + 6 files changed, 169 insertions(+), 72 deletions(-) diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 268841b0d2f..6808c707c48 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -58,11 +58,9 @@ pub(crate) fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec, rows: Range, @@ -112,7 +110,7 @@ impl Planner for ChunkPlanner { }; if let Some((scope, planner)) = plan_selection(&self.scan, scope, &self.selection, kept) { - return Ok(PlannerOutput::Planner(scope, planner)); + return Ok(PlannerOutput::Spawn(scope, planner)); } } Ok(PlannerOutput::Done) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index f9b57684d33..8b34548ffe8 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -24,10 +24,10 @@ mod io; #[cfg(test)] mod tests; -use std::collections::VecDeque; use std::future; use std::ops::Range; use std::sync::Arc; +use std::task::Poll; use arrow_array::RecordBatch; use arrow_schema::ArrowError; @@ -46,15 +46,17 @@ 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::Batch; 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; @@ -115,6 +117,7 @@ pub struct ScanBuilder { ordered: bool, row_range: Option>, selection: Selection, + concurrency: usize, map_fn: Arc VortexResult + Send + Sync>, row_offset: u64, } @@ -130,6 +133,7 @@ impl ScanBuilder { ordered: true, row_range: None, selection: Selection::default(), + concurrency: 4, map_fn: Arc::new(Ok), row_offset: 0, } @@ -199,10 +203,10 @@ impl ScanBuilder { self } - /// Kept for callers of the default executor: a V2 scan runs on one task, its chunks - /// overlapping their reads inside one planning run, so this sets nothing. - pub fn with_concurrency(self, concurrency: usize) -> 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 } @@ -220,6 +224,7 @@ impl ScanBuilder { 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)), } @@ -250,6 +255,7 @@ struct ScanRun { 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, } @@ -316,92 +322,162 @@ impl ScanRun { selection: builder.selection, map_fn: builder.map_fn, ordered: builder.ordered, + concurrency: builder.concurrency, cache, }) } - /// Runs the scan on a planning run the returned stream owns and drives as it is polled. + /// Runs the scan on planning runs the returned stream owns and drives as it is polled. /// - /// Every CPU step of the run, the planners' and morsels' compute, is a [`Handle::spawn_cpu`] - /// task; every read is a `spawn_io` task of the segment source; the stream's own poll only - /// awaits one or the other, so the engine thread polling it does no scan work itself and the - /// scan advances only as fast as its consumer pulls. Batches of an unordered scan are yielded - /// as morsels finish; an ordered scan yields them in row order once every morsel has. + /// The scan's first planner runs on one run. Every chunk it spawns runs on a run of its own, + /// with a session of its own, up to `concurrency` chunks per worker thread at once. Every CPU + /// step of a run, the planners' and morsels' compute and the batch mapping, is a + /// [`Handle::spawn_cpu`] task; every read is a `spawn_io` task of the segment source; the + /// stream's own poll awaits one or the other, so the engine thread polling it does no scan + /// work itself, and a chunk is started only when the consumer has room for its batches. + /// 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 mut run = Run::new(); - run.admit( - self.scan.planner(self.rows.clone(), self.selection), - WorkScope { - file_ordinal: 0, - rows: self.rows.clone(), - }, - self.io.session(), - ); - let driver = Driver { - handle: self.session.handle(), - run: Some(run), - map_fn: self.map_fn, - ordered: self.ordered, - held: VecDeque::new(), - done: false, - cache: self.cache, + 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(), }; - Ok(stream::try_unfold(driver, Driver::next).boxed()) + 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 batches = run_stream(handle, root, scope, io.session(), map_fn, true) + .map_ok(move |yielded| match yielded { + Yielded::Batch(start, batch) => { + stream::once(future::ready(Ok((start, batch)))).boxed() + } + Yielded::Spawn(scope, chunk) => run_stream( + chunk_handle.clone(), + chunk, + scope, + chunk_io.session(), + Arc::clone(&chunk_map), + false, + ) + .try_filter_map(|yielded| { + future::ready(Ok(match yielded { + Yielded::Batch(start, batch) => Some((start, batch)), + Yielded::Spawn(..) => None, + })) + }) + .boxed(), + }) + .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()) } } -/// The state of a scan's stream: its planning run and what it holds back. -struct Driver { +/// What a planning run's stream yields: a mapped batch with the row it starts at, or a planner +/// a planner of the run spawned as independent. +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 a 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>, - ordered: bool, - /// Batches of an ordered scan, in row order once the run is done. - held: VecDeque, - done: bool, - /// Held until the scan ends, so scans of the file that start meanwhile share it. - #[expect(dead_code, reason = "kept alive, never read")] - cache: ScanCache, + /// Whether a spawned planner is yielded to the owner, or admitted to this run as another + /// root and run here. + yield_spawns: bool, } -impl Driver { - /// The next mapped batch, or `None` once the run has retired everything. - async fn next(mut self) -> VortexResult> { +/// Runs `root` over `scope` on a planning run of its own, reading through `io`, and yields its +/// mapped batches and, when `yield_spawns`, the planners it spawns. +fn run_stream( + handle: Handle, + root: Box, + scope: WorkScope, + io: Arc, + map_fn: Arc VortexResult + Send + Sync>, + yield_spawns: 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, + yield_spawns, + }; + stream::try_unfold(state, |mut state| async move { loop { - if self.done { - return match self.held.pop_front() { - Some(batch) => Ok(Some(((self.map_fn)(batch.array)?, self))), - None => Ok(None), - }; - } - let run = self + let run = state .run .take() .ok_or_else(|| vortex_err!("the scan's run is missing"))?; - let (run, progress) = self + let map_fn = Arc::clone(&state.map_fn); + let (run, step) = state .handle .spawn_cpu(move || { let mut run = run; - let progress = run.advance(); - (run, progress) + 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 = self.run.insert(run); - match progress? { - Progress::Batch(batch) if self.ordered => self.held.push_back(batch), - Progress::Batch(batch) => { - return Ok(Some(((self.map_fn)(batch.array)?, self))); + let run = state.run.insert(run); + match step? { + Step::Batch(start, batch) => { + return Ok(Some((Yielded::Batch(start, batch), state))); } - Progress::RootDone(_) | Progress::Idle => { - self.done = true; - // Morsels finish in whatever order their reads arrive. - self.held - .make_contiguous() - .sort_by_key(|batch| batch.scope.rows.start); + Step::Spawn(scope, child) if state.yield_spawns => { + return Ok(Some((Yielded::Spawn(scope, child), state))); } - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + Step::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&state.io)); + } + 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() } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 21980158b72..eabf863eb47 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -214,6 +214,10 @@ async fn run_batches( loop { match run.advance()? { Progress::Batch(batch) => batches.push(batch), + Progress::Spawn(scope, chunk) => { + run.admit(chunk, scope, scan.io.session()); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => break, Progress::Waiting => std::future::poll_fn(|cx| run.poll_completion(cx)).await?, } @@ -789,7 +793,7 @@ async fn selected_split_scope_is_bounded_before_projection() -> VortexResult<()> let mut chunks = scan .scan .chunk_planner(scan.rows.clone(), scan.selection.clone(), None); - let PlannerOutput::Planner(scope, _) = chunks.compute()? else { + let PlannerOutput::Spawn(scope, _) = chunks.compute()? else { vortex_bail!("the selected rows were not planned"); }; assert_eq!(scope.rows, 7..2002); diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs index 8cb78099272..e9cf27da3c2 100644 --- a/vortex-scan/src/planning/driver/mod.rs +++ b/vortex-scan/src/planning/driver/mod.rs @@ -65,6 +65,9 @@ pub enum Progress { Batch(Batch), /// Every item descended from the root has retired, and its IO session was cleared. RootDone(RootId), + /// A planner declared this child independent. The run has not queued it: the owner runs it, + /// on this run as another root through [`Run::admit`] or on a run of its own elsewhere. + Spawn(WorkScope, Box), /// Every live item waits for IO. Await [`Run::poll_completion`], or pass the next completion /// to [`Run::complete`]. Waiting, @@ -462,6 +465,10 @@ impl Run { self.spawn(work.root, scope, Item::Morsel(morsel), path); self.queue.insert(work); } + Output::Planner(PlannerOutput::Spawn(scope, child)) => { + self.events.push_back(Progress::Spawn(scope, child)); + self.queue.insert(work); + } Output::Morsel(MorselOutput::Batch(array)) => { if array.is_empty() { vortex_bail!("morsel for scope {:?} returned an empty batch", work.scope); @@ -527,7 +534,8 @@ impl Driver { } } - /// Runs `root` and everything it spawns, returning batches in emission order. + /// Runs `root` and everything it spawns, returning batches in emission order. A planner + /// spawned as independent runs as another root of the same run. pub fn run(&self, root: Box) -> VortexResult> { let mut run = Run::new(); run.step_limit = self.step_limit; @@ -540,6 +548,10 @@ impl Driver { loop { match run.advance()? { Progress::Batch(batch) => batches.push(batch), + Progress::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&self.io)); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => return Ok(batches), Progress::Waiting => run.complete(self.io.wait()?)?, } diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs index 389ad24372d..9461cd4cd36 100644 --- a/vortex-scan/src/planning/driver/tests.rs +++ b/vortex-scan/src/planning/driver/tests.rs @@ -503,6 +503,7 @@ fn drain(run: &mut Run) -> VortexResult> { events.push(match run.advance()? { Progress::Batch(batch) => format!("batch {:?} {:?}", batch.root, batch.scope.rows), Progress::RootDone(root) => format!("done {root:?}"), + Progress::Spawn(scope, _) => format!("spawn {:?}", scope.rows), Progress::Waiting => { events.push("waiting".to_string()); return Ok(events); diff --git a/vortex-scan/src/planning/planner.rs b/vortex-scan/src/planning/planner.rs index 91d875ff317..cf795c99d7b 100644 --- a/vortex-scan/src/planning/planner.rs +++ b/vortex-scan/src/planning/planner.rs @@ -51,6 +51,11 @@ pub enum PlannerOutput { NeedsIO(IoBatch), /// A child planner, with the scope it is responsible for. Planner(WorkScope, Box), + /// A child planner independent of this one, with the scope it is responsible for. The run + /// hands it to its owner as [`Progress::Spawn`](super::driver::Progress::Spawn) instead of + /// running it: the owner may run it anywhere, including on another thread, and this planner + /// continues at once. + Spawn(WorkScope, Box), /// A morsel authorised to produce arrays for the scope. Morsel(WorkScope, Box), } @@ -75,6 +80,7 @@ impl fmt::Debug for PlannerOutput { Self::Continue => f.write_str("Continue"), Self::NeedsIO(batch) => f.debug_tuple("NeedsIO").field(batch).finish(), Self::Planner(scope, _) => f.debug_tuple("Planner").field(scope).finish(), + Self::Spawn(scope, _) => f.debug_tuple("Spawn").field(scope).finish(), Self::Morsel(scope, _) => f.debug_tuple("Morsel").field(scope).finish(), } } From 52daa27395ba87609c7a33acc73a237117fafbed Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 14:16:03 +0100 Subject: [PATCH 07/23] Dedup a planner's IO hints with a set Hints were deduplicated by scanning the batch for each segment, which is quadratic in a chunk's segments: a wide projection over a filtered chunk announces hundreds. A set of the segments seen replaces the scan. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-layout/src/scan/execution/filter.rs | 14 +++++++++----- vortex-layout/src/scan/execution/prune.rs | 9 +++++++++ 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 3eb9c291b56..48718307577 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -5,6 +5,7 @@ use std::ops::Range; use std::sync::Arc; use bit_vec::BitVec; +use rustc_hash::FxHashSet; use vortex_array::VortexSessionExecute; use vortex_array::builtins::ArrayBuiltins; use vortex_array::expr::BoundExpression; @@ -163,23 +164,23 @@ impl FilterPlanner { } } -/// Publishes a hint for each of `segments` not already in `batch`. +/// Publishes a hint for each of `segments` not in `seen`, and adds it there. pub(super) fn hint( locations: &[SegmentLocation], next_hint: &mut u32, intent: IoIntent, segments: &[SegmentId], + seen: &mut FxHashSet, batch: &mut IoBatch, ) { - for segment in segments { - let target = locations[**segment as usize].target(); - if batch.iter().any(|request| request.target == target) { + for &segment in segments { + if !seen.insert(segment) { continue; } batch.push(IoRequest { intent, request: IoRequestId(*next_hint), - target, + target: locations[*segment as usize].target(), }); *next_hint -= 1; } @@ -231,12 +232,14 @@ impl Planner for FilterPlanner { &selected.mask, )?; let mut batch = Vec::new(); + let mut seen = FxHashSet::default(); let locations = &self.scan.locations; hint( locations, &mut self.next_hint, IoIntent::Prefetch, &predicates, + &mut seen, &mut batch, ); hint( @@ -244,6 +247,7 @@ impl Planner for FilterPlanner { &mut self.next_hint, IoIntent::Announce, &projection, + &mut seen, &mut batch, ); if !batch.is_empty() { diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 56e5808238e..1fd153f989c 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -7,6 +7,7 @@ use std::ops::Range; use std::sync::Arc; use std::vec::IntoIter; +use rustc_hash::FxHashSet; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -132,11 +133,13 @@ impl Planner for PrunePlanner { &self.mask, )?; let mut batch = Vec::new(); + let mut seen = FxHashSet::default(); hint( &self.scan.locations, &mut self.next_hint, IoIntent::Prefetch, &segments, + &mut seen, &mut batch, ); if !batch.is_empty() { @@ -188,6 +191,12 @@ impl Planner for PrunePlanner { } // Every proof 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" + ); let next = self .next .take() From 26c4265840f6e9236f1b111ab0231a44deb4ea7b Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 14:38:15 +0100 Subject: [PATCH 08/23] Run a V2 scan's chunks beside each other only as far as the engine does not A spawned chunk runs on one task of the runtime that drives its run to the end and awaits only its reads, instead of a `spawn_cpu` task per step: a chunk's run waits for reads many times, and a hop at each wait was latency on its critical path that no other chunk's work filled. At a concurrency of one, chunks run on the scan's own run as further roots: the run's steps already interleave their work, and a run, a session and a task per chunk only added a hand-off per chunk, measured at 2 to 3 percent on TPC-H, for parallelism the engine provides. DataFusion sets the concurrency, unless configured, to the cores per partition, so a scan over one partition per core runs its chunks inline and a single-partition scan runs them beside each other. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 --- vortex-datafusion/src/persistent/opener.rs | 23 ++++- vortex-datafusion/src/persistent/source.rs | 1 + vortex-layout/src/scan/v2/mod.rs | 106 ++++++++++++++------- 3 files changed, 93 insertions(+), 37 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index cd17179b410..26620ea2522 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -66,6 +66,7 @@ use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; +use vortex_utils::parallelism::get_available_parallelism; use crate::VortexAccessPlan; use crate::convert::exprs::ExpressionConverter; @@ -130,6 +131,9 @@ pub(crate) struct VortexOpener { pub projection_pushdown: bool, pub scan_concurrency: Option, pub scan_v2: bool, + /// How many partitions the engine runs the scan on, so a V2 scan parallelises inside a + /// partition only as far as the engine has not. + pub partitions: usize, } impl FileOpener for VortexOpener { @@ -168,6 +172,7 @@ impl FileOpener for VortexOpener { let file_schemas = Arc::clone(&self.file_schemas); let has_output_ordering = self.has_output_ordering; let scan_concurrency = self.scan_concurrency; + let partitions = self.partitions.max(1); let scan_v2 = self.scan_v2; let expr_converter = Arc::clone(&self.expression_converter); @@ -449,9 +454,15 @@ impl FileOpener for VortexOpener { { builder = builder.with_selection(selection.clone()); } - if let Some(concurrency) = scan_concurrency { - builder = builder.with_concurrency(concurrency); - } + // Unless set, chunks run beside each other only as far as the engine's partitions + // leave cores idle: at one partition per core they run in the scan's own run. + let concurrency = scan_concurrency.unwrap_or_else(|| { + get_available_parallelism() + .unwrap_or(1) + .div_ceil(partitions) + .max(1) + }); + builder = builder.with_concurrency(concurrency); builder.map(map_batch).into_stream() } else { // Share V1 layout readers across partitions of the file. @@ -1104,6 +1115,7 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), + partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1513,6 +1525,7 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), + partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1604,6 +1617,7 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), + partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1786,6 +1800,7 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), + partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1850,6 +1865,7 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), + partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -2065,6 +2081,7 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), + partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 8ff7b735cee..644fea4e71c 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -378,6 +378,7 @@ impl VortexSource { projection_pushdown: self.options.projection_pushdown, scan_concurrency: self.options.scan_concurrency, scan_v2: self.options.scan_v2, + partitions: base_config.file_groups.len(), }; Ok(opener) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 8b34548ffe8..87e830c1a13 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -33,9 +33,11 @@ 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; @@ -327,16 +329,16 @@ impl ScanRun { }) } - /// Runs the scan on planning runs the returned stream owns and drives as it is polled. + /// Runs the scan on planning runs the returned stream owns. /// - /// The scan's first planner runs on one run. Every chunk it spawns runs on a run of its own, - /// with a session of its own, up to `concurrency` chunks per worker thread at once. Every CPU - /// step of a run, the planners' and morsels' compute and the batch mapping, is a - /// [`Handle::spawn_cpu`] task; every read is a `spawn_io` task of the segment source; the - /// stream's own poll awaits one or the other, so the engine thread polling it does no scan - /// work itself, and a chunk is started only when the consumer has room for its batches. - /// Batches of an unordered scan are yielded as morsels finish; an ordered scan yields them in - /// row order once every morsel has. + /// 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); @@ -349,26 +351,19 @@ impl ScanRun { let map_fn = self.map_fn; let (chunk_handle, chunk_io, chunk_map) = (handle.clone(), Arc::clone(&io), Arc::clone(&map_fn)); - let batches = run_stream(handle, root, scope, io.session(), map_fn, true) + 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) => run_stream( - chunk_handle.clone(), + Yielded::Spawn(scope, chunk) => chunk_stream( + &chunk_handle, chunk, scope, chunk_io.session(), Arc::clone(&chunk_map), - false, - ) - .try_filter_map(|yielded| { - future::ready(Ok(match yielded { - Yielded::Batch(start, batch) => Some((start, batch)), - Yielded::Spawn(..) => None, - })) - }) - .boxed(), + ), }) .try_flatten_unordered(in_flight); let cache = self.cache; @@ -390,8 +385,8 @@ impl ScanRun { } } -/// What a planning run's stream yields: a mapped batch with the row it starts at, or a planner -/// a planner of the run spawned as independent. +/// 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), @@ -406,27 +401,29 @@ enum Step { Waiting, } -/// The state of a run's stream. +/// 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 planner is yielded to the owner, or admitted to this run as another - /// root and run here. - yield_spawns: bool, + /// 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`, and yields its -/// mapped batches and, when `yield_spawns`, the planners it spawns. +/// 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>, - yield_spawns: bool, + inline: bool, ) -> BoxStream<'static, VortexResult>> { let mut run = Run::new(); run.admit(root, scope, Arc::clone(&io)); @@ -435,7 +432,7 @@ fn run_stream( run: Some(run), io, map_fn, - yield_spawns, + inline, }; stream::try_unfold(state, |mut state| async move { loop { @@ -467,11 +464,11 @@ fn run_stream( Step::Batch(start, batch) => { return Ok(Some((Yielded::Batch(start, batch), state))); } - Step::Spawn(scope, child) if state.yield_spawns => { - return Ok(Some((Yielded::Spawn(scope, child), state))); + Step::Spawn(scope, child) if state.inline => { + run.admit(child, scope, Arc::clone(&state.io)); } Step::Spawn(scope, child) => { - run.admit(child, scope, Arc::clone(&state.io)); + return Ok(Some((Yielded::Spawn(scope, child), state))); } Step::RootDone if run.live_roots() > 0 => {} Step::RootDone | Step::Idle => return Ok(None), @@ -481,3 +478,44 @@ fn run_stream( }) .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() +} From c441bc4795d041cb8b82f18f59a326f1b8773745 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 23 Oct 2026 15:07:22 +0100 Subject: [PATCH 09/23] Link the run driver by its imported name in the V2 module doc rustdoc rejects the explicit target as redundant now that `Run` is imported here. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Vr65ZwPZz4LPjkvDA6ZUuE Signed-off-by: Joe Isaacs --- vortex-layout/src/scan/v2/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 87e830c1a13..d60f6b76289 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -4,7 +4,7 @@ //! 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`](vortex_scan::planning::driver::Run) driver advances +//! 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. //! From 33a5b2481340659ae6304867ae78d0b9ac122272 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 23 Oct 2026 14:23:54 +0000 Subject: [PATCH 10/23] Pass an execution context to ListArray::try_new in the nested empty list test #10533 added `nested_empty_list_roundtrip` calling `ListArray::try_new` without the `&mut ExecutionCtx` argument it now takes, which broke the `vortex-file` test build on develop. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01BL868qtvZQwUKzuhNuWAS2 --- vortex-file/src/tests.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 368e56d970f..40e03b3ab94 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2047,17 +2047,20 @@ async fn nested_list_of_list_roundtrip() -> VortexResult<()> { #[tokio::test] #[cfg_attr(miri, ignore)] async fn nested_empty_list_roundtrip(#[case] validity: Validity) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let inner = ListArray::try_new( PrimitiveArray::new::(buffer![], Validity::NonNullable).into_array(), buffer![0u32].into_array(), Validity::NonNullable, + &mut ctx, )? .into_array(); - let outer = ListArray::try_new(inner, buffer![0u32, 0, 0].into_array(), validity)?.into_array(); + let outer = ListArray::try_new(inner, buffer![0u32, 0, 0].into_array(), validity, &mut ctx)? + .into_array(); let array = StructArray::from_fields(&[("nested", outer)])?.into_array(); let result = write_read_roundtrip(array.clone()).await?; - assert_arrays_eq!(result, array, &mut SESSION.create_execution_ctx()); + assert_arrays_eq!(result, array, &mut ctx); Ok(()) } From a07e7f13fade1e0495b536376740c1d6733ad6e2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 18:01:49 +0000 Subject: [PATCH 11/23] Run a V2 scan's chunks inline at concurrency one, each root on its own IO session At one partition per core the opener passes a concurrency of one, but inline was gated on concurrency times the core count, so every partition ran a core's worth of chunks at once. Chunks admitted inline or spawned inside a chunk now each get their own IO session, since retiring a root clears its session's reads. Also logs how long preparing a scan takes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-layout/src/scan/v2/mod.rs | 70 +++++++++++++++----------------- 1 file changed, 33 insertions(+), 37 deletions(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index d60f6b76289..02b5507f213 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -28,6 +28,7 @@ use std::future; use std::ops::Range; use std::sync::Arc; use std::task::Poll; +use std::time::Instant; use arrow_array::RecordBatch; use arrow_schema::ArrowError; @@ -48,7 +49,6 @@ 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; @@ -205,7 +205,8 @@ impl ScanBuilder { self } - /// Sets the number of chunks run at once per worker thread. Must be greater than zero. + /// Sets the number of chunks run at once per worker thread. Must be greater than zero. At one, + /// the chunks run on the scan's own run instead, for engines that parallelise across scans. pub fn with_concurrency(mut self, concurrency: usize) -> Self { assert!(concurrency > 0); self.concurrency = concurrency; @@ -238,7 +239,12 @@ impl ScanBuilder { let handle = self.session.handle(); Ok(stream::once(async move { handle - .spawn_cpu(move || ScanRun::try_new(self)?.execute_stream()) + .spawn_cpu(move || { + let start = Instant::now(); + let run = ScanRun::try_new(self)?; + tracing::debug!(elapsed = ?start.elapsed(), "prepared V2 scan"); + run.execute_stream() + }) .await }) .try_flatten() @@ -268,34 +274,26 @@ impl ScanRun { fn try_new(builder: ScanBuilder) -> VortexResult { let file = builder.file; let cache = file.cache.unwrap_or_default(); - let root = match cache.root() { - Some(root) => root, - None => cache.insert_root(lower(&file.layout)?), - }; + let root = cache.root_or_try_init(|| lower(&file.layout))?; let row_count = root.row_count(); + let session = &builder.session; let filter = builder .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)?); - Ok(cache.insert_filter_plans(filter, plans)) - } - } + .map(|filter| { + cache.filter_plans_or_try_init(&filter, || { + FilterPlans::new(filter.clone(), &root, session) + }) }) .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) { - Some(plan) => plan, - None => { - let plan = optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; - cache.insert_projection(builder.projection, plan) - } - }; + let projection = cache.projection_or_try_init(&builder.projection, || { + optimize(plan_row_idx_expression( + builder.projection.clone(), + root.clone(), + )?) + })?; let chunk_rows = if filter.is_some() { FILTERED_ROWS_PER_CHUNK } else if identity { @@ -303,10 +301,7 @@ impl ScanRun { } else { MAX_ROWS_PER_CHUNK }; - let ranges = match cache.ranges() { - Some(ranges) => ranges, - None => cache.insert_ranges(segment_ranges(&file.locations)), - }; + let ranges = cache.ranges_or_init(|| segment_ranges(&file.locations)); let io = Arc::new(SegmentScanIo::new(file.segments, ranges)); Ok(Self { session: builder.session.clone(), @@ -351,8 +346,8 @@ impl ScanRun { 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) + let inline = self.concurrency <= 1; + let batches = run_stream(handle, root, scope, Arc::clone(&io), map_fn, inline) .map_ok(move |yielded| match yielded { Yielded::Batch(start, batch) => { stream::once(future::ready(Ok((start, batch)))).boxed() @@ -361,7 +356,7 @@ impl ScanRun { &chunk_handle, chunk, scope, - chunk_io.session(), + Arc::clone(&chunk_io), Arc::clone(&chunk_map), ), }) @@ -406,7 +401,8 @@ struct RunState { handle: Handle, /// Taken while a CPU step runs on its task and put back with its result. run: Option, - io: Arc, + /// Gives every root of the run a session of its own: retiring a root clears its session. + 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, @@ -421,12 +417,12 @@ fn run_stream( handle: Handle, root: Box, scope: WorkScope, - io: Arc, + 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)); + run.admit(root, scope, io.session()); let state = RunState { handle, run: Some(run), @@ -465,7 +461,7 @@ fn run_stream( return Ok(Some((Yielded::Batch(start, batch), state))); } Step::Spawn(scope, child) if state.inline => { - run.admit(child, scope, Arc::clone(&state.io)); + run.admit(child, scope, state.io.session()); } Step::Spawn(scope, child) => { return Ok(Some((Yielded::Spawn(scope, child), state))); @@ -489,13 +485,13 @@ fn chunk_stream( handle: &Handle, root: Box, scope: WorkScope, - io: Arc, + 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)); + run.admit(root, scope, io.session()); loop { match run.advance()? { Progress::Batch(batch) => { @@ -507,7 +503,7 @@ fn chunk_stream( } } Progress::Spawn(scope, child) => { - run.admit(child, scope, Arc::clone(&io)); + run.admit(child, scope, io.session()); } Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => return Ok(()), From 338e5503d825182357191024ae37cca1c9757f4c Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 18:01:49 +0000 Subject: [PATCH 12/23] Build each V2 scan cache entry once when partitions miss together Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-layout/src/scan/v2/cache.rs | 97 ++++++++++++++---------------- 1 file changed, 44 insertions(+), 53 deletions(-) diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 5a26e28ec97..b7620556181 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -6,10 +6,12 @@ use std::sync::Arc; use std::sync::Weak; +use once_cell::sync::OnceCell; use parking_lot::Mutex; use rustc_hash::FxHashMap; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_error::VortexResult; use super::io::SegmentRanges; use crate::plan::physical::PlanRef; @@ -35,81 +37,62 @@ pub struct ScanCache(Arc); #[derive(Default)] struct Inner { - root: Mutex>, - /// Optimized plans of non-identity projections. Values are cached by the plan that produced - /// them, so a projection must reuse its plan for its values to be found again. - projections: Mutex>, - ranges: Mutex>, + root: OnceCell, + /// Optimized plans of projections. Values are cached by the plan that produced them, so a + /// projection must reuse its plan for its values to be found again. + projections: Mutex>>>, + ranges: OnceCell, /// 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. - filters: Mutex>>, + filters: Mutex>>>>, values: ValueCache, } impl ScanCache { - /// The lowered root plan an earlier scan kept. - pub(super) fn root(&self) -> Option { - self.0.root.lock().clone() - } - - /// Keeps `root`, returning the plan already kept if another scan got there first. - pub(super) fn insert_root(&self, root: PlanRef) -> PlanRef { - self.0.root.lock().get_or_insert(root).clone() + /// The lowered root plan, lowered by `lower` unless an earlier scan kept it. Scans that miss + /// together wait for one of them to lower it. + pub(super) fn root_or_try_init( + &self, + lower: impl FnOnce() -> VortexResult, + ) -> VortexResult { + self.0.root.get_or_try_init(lower).cloned() } - /// The optimized plan of `projection` an earlier scan kept. - pub(super) fn projection(&self, projection: &BoundExpression) -> Option { - self.0 - .projections - .lock() - .get(&ExactBoundExpr(projection.clone())) + /// The optimized plan of `projection`, planned by `plan` unless an earlier scan kept it. + pub(super) fn projection_or_try_init( + &self, + projection: &BoundExpression, + plan: impl FnOnce() -> VortexResult, + ) -> VortexResult { + cell(&self.0.projections, projection) + .get_or_try_init(plan) .cloned() } - /// Keeps the optimized plan of `projection`, returning the plan already kept if another scan - /// got there first. - pub(super) fn insert_projection(&self, projection: BoundExpression, plan: PlanRef) -> PlanRef { - self.0 - .projections - .lock() - .entry(ExactBoundExpr(projection)) - .or_insert(plan) - .clone() - } - - /// The segment byte-range index an earlier scan kept. - pub(super) fn ranges(&self) -> Option { - self.0.ranges.lock().clone() - } - - /// Keeps `ranges`, returning the index already kept if another scan got there first. - pub(super) fn insert_ranges(&self, ranges: SegmentRanges) -> SegmentRanges { - Arc::clone(self.0.ranges.lock().get_or_insert(ranges)) + /// The segment byte-range index, built by `index` unless an earlier scan kept it. + pub(super) fn ranges_or_init(&self, index: impl FnOnce() -> SegmentRanges) -> SegmentRanges { + Arc::clone(self.0.ranges.get_or_init(index)) } /// The plans of `filter` an earlier scan kept. + #[cfg(test)] pub(super) fn filter_plans(&self, filter: &BoundExpression) -> Option> { self.0 .filters .lock() .get(&ExactBoundExpr(filter.clone())) - .cloned() + .and_then(|plans| plans.get().cloned()) } - /// Keeps the plans of `filter`, returning the plans already kept if another scan got there - /// first. - pub(super) fn insert_filter_plans( + /// The plans of `filter`, planned by `plan` unless an earlier scan kept them. + pub(super) fn filter_plans_or_try_init( &self, - filter: BoundExpression, - plans: Arc, - ) -> Arc { - Arc::clone( - self.0 - .filters - .lock() - .entry(ExactBoundExpr(filter)) - .or_insert(plans), - ) + filter: &BoundExpression, + plan: impl FnOnce() -> VortexResult, + ) -> VortexResult> { + cell(&self.0.filters, filter) + .get_or_try_init(|| plan().map(Arc::new)) + .cloned() } pub(super) fn values(&self) -> ValueCache { @@ -122,6 +105,14 @@ impl ScanCache { } } +/// The cell of `key` in `map`, so scans that miss together build its entry once. +fn cell( + map: &Mutex>>>, + key: &BoundExpression, +) -> Arc> { + Arc::clone(map.lock().entry(ExactBoundExpr(key.clone())).or_default()) +} + /// A [`ScanCache`] handle that does not keep the cache alive. Engines that open a file once per /// partition keep one per file, so concurrent partitions share a cache that is freed with the /// last scan holding it. From efa7556148de0d13b54b89530cd39d9744eec524 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 18:01:49 +0000 Subject: [PATCH 13/23] Register a whole IO batch before polling any of its reads Polling each read as it was registered let the segment source issue a prefetch before the announcements it could coalesce with existed. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-layout/src/scan/v2/io.rs | 165 +++++++++++++++++++++----------- 1 file changed, 108 insertions(+), 57 deletions(-) diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 2779739b281..5dfd21cf6af 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -87,7 +87,7 @@ pub(super) struct SegmentIoSource { announced: Mutex>>, } -/// A read started ahead of the fetch that needs it. +/// A segment read the session shares among every request of the segment. type Prefetched = Shared>>>; impl SegmentIoSource { @@ -101,71 +101,69 @@ impl SegmentIoSource { announced: Mutex::default(), } } + + /// Registers a read of `segment`, taking over an announcement of it if there is one. + fn read(&self, segment: SegmentId) -> Prefetched { + self.announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared() + } } impl IoSource for SegmentIoSource { fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { - let reads = self.reads.lock(); - let mut ready = self.ready.lock(); - let mut prefetched = self.prefetched.lock(); - for request in batch { - let IoTarget::Range { offset, len, .. } = request.target else { - vortex_bail!( - "SegmentIoSource only serves byte ranges, not {:?}", - request.target - ); - }; - let segment = *self - .ranges - .get(&(offset, len)) - .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; - if request.intent == IoIntent::Announce { - self.announced - .lock() - .entry(segment) - .or_default() - .push_back(self.segments.request(segment)); - continue; - } - 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(); - // Polling registers the read as wanted, so the source starts it now. - drop( - read.clone() - .poll_unpin(&mut Context::from_waker(Waker::noop())), + let mut started = Vec::with_capacity(batch.len()); + { + let mut prefetched = self.prefetched.lock(); + for request in batch { + let IoTarget::Range { offset, len, .. } = request.target else { + vortex_bail!( + "SegmentIoSource only serves byte ranges, not {:?}", + request.target ); - read - }); - continue; - } - // A segment stays with the session once read, so a later stage or morsel of the - // split that reads it again shares the read instead of starting another once the - // first has completed. - let mut bytes = prefetched - .entry(segment) - .or_insert_with(|| { + }; + let segment = *self + .ranges + .get(&(offset, len)) + .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; + if request.intent == IoIntent::Announce { self.announced .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) - .map_err(Arc::new) - .boxed() - .shared() - }) - .clone() + .entry(segment) + .or_default() + .push_back(self.segments.request(segment)); + continue; + } + // A segment stays with the session once read, so a later stage or morsel of the + // split that reads it again shares the read instead of starting another once the + // first has completed. + let read = prefetched + .entry(segment) + .or_insert_with(|| self.read(segment)) + .clone(); + let fetch = (request.intent == IoIntent::Fetch).then_some(request.request); + started.push((segment, read, fetch)); + } + } + + // Every read of the batch is registered before any is polled, so a coalescing source + // sees the whole batch, announcements included, before it picks a read to issue. + let reads = self.reads.lock(); + let mut ready = self.ready.lock(); + for (segment, mut read, fetch) in started { + let Some(id) = fetch else { + // Polling registers the read as wanted, so the source starts it now. + drop(read.poll_unpin(&mut Context::from_waker(Waker::noop()))); + continue; + }; + let mut bytes = read .map_err(move |err| vortex_err!("read of segment {segment} failed: {err}")) .boxed(); - let id = request.request; // In-memory and cached segments need no allocated asynchronous task. A pending // read is polled again by FuturesUnordered before the driver can go to sleep. if let Poll::Ready(result) = bytes.poll_unpin(&mut Context::from_waker(Waker::noop())) { @@ -235,6 +233,7 @@ mod tests { use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use std::task::Poll; use futures::FutureExt; use futures::channel::oneshot; @@ -298,6 +297,58 @@ mod tests { } } + /// Logs each segment's registration and first poll; its reads never complete. + #[derive(Default)] + struct LoggingSegments(Arc>>); + + impl SegmentSource for LoggingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.0.lock().push(format!("register {id}")); + let log = Arc::clone(&self.0); + let mut polled = false; + future::poll_fn(move |_| { + if !std::mem::replace(&mut polled, true) { + log.lock().push(format!("poll {id}")); + } + Poll::Pending + }) + .boxed() + } + } + + /// A coalescing source sees every read of a batch, announcements included, before any of + /// them is polled and so made eligible to issue. + #[test] + fn batch_registers_every_read_before_polling() -> VortexResult<()> { + let segments = LoggingSegments::default(); + let log = Arc::clone(&segments.0); + let locations = [0, 4, 8].map(|offset| SegmentLocation { + offset, + length: 4, + alignment: Alignment::none(), + }); + let io = SegmentIoSource::new(Arc::new(segments), segment_ranges(&locations)); + let batch = [ + (IoIntent::Prefetch, 0), + (IoIntent::Fetch, 4), + (IoIntent::Announce, 8), + ] + .into_iter() + .enumerate() + .map(|(id, (intent, offset))| IoRequest { + intent, + request: IoRequestId(u32::try_from(id).unwrap_or(u32::MAX)), + target: IoTarget::range(offset, 4), + }) + .collect(); + io.submit(IoOwnerId(0), batch)?; + assert_eq!( + *log.lock(), + ["register 0", "register 1", "register 2", "poll 0", "poll 1"] + ); + Ok(()) + } + #[tokio::test] async fn pending_read_wakes_after_an_immediate_completion() -> VortexResult<()> { let (sender, receiver) = oneshot::channel(); From 64e12b68dbf2de83d7159d3b3ad0f227c888a781 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 18:01:49 +0000 Subject: [PATCH 14/23] Share the mask-driving loop of the filter and prune planners, and align projection windows MaskScan drives a boolean plan and combines its pieces, taking a single piece without copying. Dense projections are cut on the file's row grid instead of from the trimmed chunk start, so writer chunks are not split between two morsels. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-layout/src/scan/execution/filter.rs | 106 ++++++++++++------ .../src/scan/execution/projection.rs | 21 ++-- vortex-layout/src/scan/execution/prune.rs | 49 +++----- 3 files changed, 104 insertions(+), 72 deletions(-) diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 48718307577..4ce6a65482d 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -124,8 +124,7 @@ pub(super) struct FilterPlanner { selected: Option, filter: Arc, remaining: BitVec, - active: Option<(usize, ScanAdaptor)>, - pieces: Vec, + active: Option<(usize, MaskScan)>, announced: bool, /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. next_hint: u32, @@ -157,7 +156,6 @@ impl FilterPlanner { selected: Some(selected), filter, active: None, - pieces: Vec::new(), announced: false, next_hint: u32::MAX, } @@ -187,15 +185,72 @@ pub(super) fn hint( } /// Executes a scan's boolean piece to a mask, nulls as false. -pub(super) fn piece_mask(session: &VortexSession, piece: Piece) -> VortexResult { +fn piece_mask(session: &VortexSession, piece: Piece) -> VortexResult { let mut ctx = session.create_execution_ctx(); piece.array.fill_null(false)?.execute::(&mut ctx) } +/// A boolean plan running over one split, its pieces gathered into one mask. +pub(super) struct MaskScan { + adaptor: ScanAdaptor, + pieces: Vec, +} + +/// What one [`MaskScan::compute`] produced. +pub(super) enum MaskStep { + /// The scan is still running: its owner's compute returns this. + Running(PlannerOutput), + /// The mask over the split's selected rows, nulls as false. + Finished(Mask), +} + +impl MaskScan { + pub(super) fn new(scan: &PreparedScan, plan: PlanRef, split: Split) -> VortexResult { + let pipeline = Scan::try_new(scan.session.clone(), plan, vec![split])? + .with_row_offset(scan.row_offset) + .with_value_cache(scan.values.clone()); + Ok(Self { + adaptor: ScanAdaptor::new(pipeline, Arc::clone(&scan.locations)), + pieces: Vec::new(), + }) + } + + /// The owner's state while the scan runs: once it has finished, its mask is ready to take. + pub(super) fn state(&self) -> State { + match self.adaptor.state() { + State::Done => State::NeedsCompute, + state => state, + } + } + + pub(super) fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + self.adaptor.set_io_result(request, result); + } + + pub(super) fn compute(&mut self, session: &VortexSession) -> VortexResult { + if self.adaptor.state() == State::Done { + // A single piece is already the mask; concatenating would copy its bits. + let mask = match self.pieces.as_slice() { + [mask] => mask.clone(), + pieces => Mask::concat(pieces.iter())?, + }; + return Ok(MaskStep::Finished(mask)); + } + Ok(MaskStep::Running(match self.adaptor.compute()? { + AdaptorStep::Yield => PlannerOutput::Continue, + AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + AdaptorStep::Piece(piece) => { + self.pieces.push(piece_mask(session, piece)?); + PlannerOutput::Continue + } + })) + } +} + impl IoConsumer for FilterPlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { - if let Some((_, adaptor)) = &mut self.active { - adaptor.set_io_result(request, result); + if let Some((_, conjunct)) = &mut self.active { + conjunct.set_io_result(request, result); } } } @@ -206,8 +261,8 @@ impl Planner for FilterPlanner { return State::Done; } match &self.active { - Some((_, adaptor)) if adaptor.state() != State::Done => adaptor.state(), - _ => State::NeedsCompute, + Some((_, conjunct)) => conjunct.state(), + None => State::NeedsCompute, } } @@ -254,19 +309,12 @@ impl Planner for FilterPlanner { return Ok(PlannerOutput::NeedsIO(batch)); } } - if let Some((index, adaptor)) = &mut self.active { - if adaptor.state() != State::Done { - return Ok(match adaptor.compute()? { - AdaptorStep::Yield => PlannerOutput::Continue, - AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), - AdaptorStep::Piece(piece) => { - self.pieces.push(piece_mask(&self.scan.session, piece)?); - PlannerOutput::Continue - } - }); - } + if let Some((index, conjunct)) = &mut self.active { + let result = match conjunct.compute(&self.scan.session)? { + MaskStep::Running(output) => return Ok(output), + MaskStep::Finished(result) => result, + }; let input = selected.mask.true_count(); - let result = Mask::concat(self.pieces.iter())?; vortex_ensure_eq!(result.len(), input); selected.mask = selected.mask.intersect_by_rank(&result); tracing::debug!( @@ -281,7 +329,6 @@ impl Planner for FilterPlanner { .expression .report_selectivity(*index, selected.mask.true_count() as f64 / input as f64); self.active = None; - self.pieces.clear(); } if selected.mask.all_false() { self.selected = None; @@ -289,20 +336,15 @@ impl Planner for FilterPlanner { } if let Some(index) = self.filter.expression.next_conjunct(&self.remaining) { self.remaining.set(index, false); - let scan = Scan::try_new( - self.scan.session.clone(), + let conjunct = MaskScan::new( + &self.scan, self.filter.predicates[index].clone(), - vec![Split { + Split { rows: selected.scope.rows.clone(), mask: selected.mask.clone(), - }], - )? - .with_row_offset(self.scan.row_offset) - .with_value_cache(self.scan.values.clone()); - self.active = Some(( - index, - ScanAdaptor::new(scan, Arc::clone(&self.scan.locations)), - )); + }, + )?; + self.active = Some((index, conjunct)); return Ok(PlannerOutput::Continue); } let selected = self.selected.take().vortex_expect("selected rows exist"); diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 141b7377c26..8b238510f14 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -7,6 +7,7 @@ use std::vec::IntoIter; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_io::request::IoConsumer; @@ -28,11 +29,13 @@ use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; +use crate::scan::execution::chunk::row_chunks; const MAX_OUTPUT_ROWS: usize = 65_536; /// Keeps sparse selections together and subdivides dense selections to avoid large copies when -/// converting chunked projection results to Arrow. +/// converting chunked projection results to Arrow. Subdivisions lie on a fixed grid of the file's +/// rows, so a writer chunk aligned to it is never split between two morsels. pub(super) struct ProjectionPlanner { scan: Arc, selected: IntoIter, @@ -44,19 +47,21 @@ impl ProjectionPlanner { let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { vec![selected] } else { - (0..selected.mask.len()) - .step_by(MAX_OUTPUT_ROWS) - .filter_map(|start| { - let end = (start + MAX_OUTPUT_ROWS).min(selected.mask.len()); - let mask = selected.mask.slice(start..end); + let offset = |row: u64| { + usize::try_from(row - selected.scope.rows.start) + .vortex_expect("a work chunk fits usize") + }; + row_chunks(selected.scope.rows.clone(), MAX_OUTPUT_ROWS) + .into_iter() + .filter_map(|rows| { + let mask = selected.mask.slice(offset(rows.start)..offset(rows.end)); if mask.all_false() { return None; } Some(SelectedRows { scope: WorkScope { file_ordinal: selected.scope.file_ordinal, - rows: selected.scope.rows.start + start as u64 - ..selected.scope.rows.start + end as u64, + rows, }, mask, }) diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 1fd153f989c..d82dd7690e0 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -23,14 +23,12 @@ use vortex_scan::planning::planner::WorkScope; use vortex_session::VortexSession; use super::PreparedScan; -use super::adaptor::AdaptorStep; -use super::adaptor::ScanAdaptor; +use super::filter::MaskScan; +use super::filter::MaskStep; 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. @@ -62,8 +60,7 @@ pub(super) struct PrunePlanner { /// The rows no proof has excluded yet, over `rows`. mask: Mask, proofs: IntoIter, - active: Option, - pieces: Vec, + active: Option, announced: bool, /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. next_hint: u32, @@ -85,7 +82,6 @@ impl PrunePlanner { mask: Mask::new_true(len), proofs: proofs.into_iter(), active: None, - pieces: Vec::new(), announced: false, next_hint: u32::MAX, done: false, @@ -95,8 +91,8 @@ impl PrunePlanner { impl IoConsumer for PrunePlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { - if let Some(adaptor) = &mut self.active { - adaptor.set_io_result(request, result); + if let Some(proof) = &mut self.active { + proof.set_io_result(request, result); } } } @@ -107,8 +103,8 @@ impl Planner for PrunePlanner { return State::Done; } match &self.active { - Some(adaptor) if adaptor.state() != State::Done => adaptor.state(), - _ => State::NeedsCompute, + Some(proof) => proof.state(), + None => State::NeedsCompute, } } @@ -146,19 +142,12 @@ impl Planner for PrunePlanner { return Ok(PlannerOutput::NeedsIO(batch)); } } - if let Some(adaptor) = &mut self.active { - if adaptor.state() != State::Done { - return Ok(match adaptor.compute()? { - AdaptorStep::Yield => PlannerOutput::Continue, - AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), - AdaptorStep::Piece(piece) => { - self.pieces.push(piece_mask(&self.scan.session, piece)?); - PlannerOutput::Continue - } - }); - } + if let Some(proof) = &mut self.active { + let excluded = match proof.compute(&self.scan.session)? { + MaskStep::Running(output) => return Ok(output), + MaskStep::Finished(excluded) => excluded, + }; let input = self.mask.true_count(); - let excluded = Mask::concat(self.pieces.iter())?; vortex_ensure_eq!(excluded.len(), input); self.mask = self.mask.intersect_by_rank(&!excluded); tracing::debug!( @@ -169,24 +158,20 @@ impl Planner for PrunePlanner { "v2 zone proof evaluated" ); self.active = None; - self.pieces.clear(); } if self.mask.all_false() { self.done = true; return Ok(PlannerOutput::Done); } if let Some(proof) = self.proofs.next() { - let scan = Scan::try_new( - self.scan.session.clone(), + self.active = Some(MaskScan::new( + &self.scan, proof, - vec![Split { + Split { rows: self.rows.clone(), mask: self.mask.clone(), - }], - )? - .with_row_offset(self.scan.row_offset) - .with_value_cache(self.scan.values.clone()); - 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. From 6168ba0ae79b660931ae7112abf0303add8572c4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 18:01:49 +0000 Subject: [PATCH 15/23] Drop identity struct evals and skip lowering legacy stats for data reads An Eval rebuilding every root field of a Pack in order is the identity. Legacy stats only need converting to zones under a proof, which the expression rule already handles; a data read reduces straight to data. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- .../src/plan/physical/optimizer/mod.rs | 8 +-- vortex-layout/src/plan/physical/plans/eval.rs | 14 ++++-- .../src/plan/physical/plans/legacy_stats.rs | 50 ++++++++++--------- vortex-layout/src/plan/physical/plans/mod.rs | 2 +- vortex-layout/src/plan/physical/plans/pack.rs | 5 +- .../src/plan/physical/plans/zoned.rs | 5 +- vortex-layout/src/plan/physical/tests.rs | 7 +-- 7 files changed, 51 insertions(+), 40 deletions(-) diff --git a/vortex-layout/src/plan/physical/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs index 073bc712d36..b9fff16070a 100644 --- a/vortex-layout/src/plan/physical/optimizer/mod.rs +++ b/vortex-layout/src/plan/physical/optimizer/mod.rs @@ -28,7 +28,7 @@ use super::plans::ExpressionLegacyStatsRule; use super::plans::ExpressionPackRule; use super::plans::ExpressionTakeRule; use super::plans::ExpressionZonedRule; -use super::plans::LegacyStatsRule; +use super::plans::LegacyStatsDataRule; use super::plans::ZonedDataRule; use crate::plan::physical::plans::EvalEvalRule; use crate::plan::physical::plans::EvalPackRule; @@ -36,8 +36,8 @@ use crate::plan::physical::plans::EvalPackRule; static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(EvalIdentityRule); -static LEGACY_STATS_RULE: PlanReduceRuleAdapter = - PlanReduceRuleAdapter::new(LegacyStatsRule); +static LEGACY_STATS_DATA_RULE: PlanReduceRuleAdapter = + PlanReduceRuleAdapter::new(LegacyStatsDataRule); static ZONED_DATA_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(ZonedDataRule); @@ -45,7 +45,7 @@ static ZONED_DATA_RULE: PlanReduceRuleAdapter = static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[ &EVAL_IDENTITY_RULE, &PlanReduceRuleAdapter::new(EvalPackRule), - &LEGACY_STATS_RULE, + &LEGACY_STATS_DATA_RULE, &ZONED_DATA_RULE, ]); diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index b6a2f1ff620..1b6117d5f17 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -32,6 +32,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::pack::is_identity_expression; /// Applies an expression to the output of its child. #[derive(Clone, Debug)] @@ -149,17 +150,20 @@ fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> V Ok(()) } -/// Removes an [`Eval`] whose expression is the identity expression. +/// Removes an [`Eval`] whose expression is the identity: the root, or a non-nullable struct +/// rebuilt from every field of the root in order, as pushing a projection into a +/// [`PackPlan`] leaves. #[derive(Debug)] pub(crate) struct EvalIdentityRule; impl PlanReduceRule for EvalIdentityRule { fn reduce(&self, plan: &Plan) -> VortexResult> { - if plan.expression().is_root() { - Ok(Some(plan.child_plan()?)) - } else { - Ok(None) + let expression = plan.expression(); + if !expression.is_root() && !expression.dtype().is_struct() { + return Ok(None); } + let child = plan.child_plan()?; + Ok(is_identity_expression(expression, child.dtype())?.then_some(child)) } } diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index 4abd2932509..8092cbc308b 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/vortex-layout/src/plan/physical/plans/legacy_stats.rs @@ -38,10 +38,10 @@ const ZONES: usize = 1; /// Data alongside a legacy `vortex.stats` zone table, whose fields are named by [`Stat`] rather /// than by aggregate function. /// -/// The optimizer rewrites it into a [`Zoned`] plan over the zone table converted to the aggregate -/// schema, or into its data alone when a [`Zoned`] plan already summarises the same data; a plan -/// that reads the data is then optimized to the data alone. Compiled unoptimized, it reads its -/// data and prunes nothing. +/// Beneath a zone proof the optimizer rewrites it into a [`Zoned`] plan over the zone table +/// converted to the aggregate schema, or into its data alone when a [`Zoned`] plan already +/// summarises the same data; a plan that reads the data is optimized to the data alone. Compiled +/// unoptimized, it reads its data and prunes nothing. #[derive(Clone, Debug)] pub struct LegacyStats; @@ -155,24 +155,31 @@ impl PlanVTable for LegacyStats { /// A plan whose data is already [`Zoned`] keeps only that data: legacy stats are used only when no /// aggregate zone map summarises the same rows. A zero zone length, as older files may record, or /// a table with no stat an aggregate can stand for, leaves only the data too. +fn to_zoned(plan: &LegacyStatsPlan) -> VortexResult { + let data = plan.data_plan()?; + if plan.zone_len() == 0 || data.is::() { + return Ok(data); + } + let zones = plan.zones_plan()?; + let (aggregate_fns, conversion) = + legacy_stats_to_aggregates(plan.dtype(), zones.dtype(), plan.stats()); + if aggregate_fns.is_empty() { + return Ok(data); + } + let zones = EvalPlan::try_new(conversion.bind(zones.dtype())?, zones)?.into_plan(); + Ok(ZonedPlan::new(data, zones, plan.zone_len(), aggregate_fns).into_plan()) +} + +/// Replaces a legacy stats plan by its data, as [`ZonedDataRule`](super::ZonedDataRule) does for +/// the [`Zoned`] plan it would convert to. A zone proof over it is converted by +/// [`ExpressionLegacyStatsRule`] before this runs, so the zone table is never lowered or +/// converted only to be dropped. #[derive(Debug)] -pub(crate) struct LegacyStatsRule; +pub(crate) struct LegacyStatsDataRule; -impl PlanReduceRule for LegacyStatsRule { +impl PlanReduceRule for LegacyStatsDataRule { fn reduce(&self, plan: &Plan) -> VortexResult> { - let data = plan.data_plan()?; - if plan.zone_len() == 0 || data.is::() { - return Ok(Some(data)); - } - let zones = plan.zones_plan()?; - let (aggregate_fns, conversion) = - legacy_stats_to_aggregates(plan.dtype(), zones.dtype(), plan.stats()); - if aggregate_fns.is_empty() { - return Ok(Some(data)); - } - let zones = EvalPlan::try_new(conversion.bind(zones.dtype())?, zones)?.into_plan(); - let zoned = ZonedPlan::new(data, zones, plan.zone_len(), aggregate_fns); - Ok(Some(zoned.into_plan())) + plan.data_plan().map(Some) } } @@ -193,11 +200,8 @@ impl PlanParentReduceRule for ExpressionLegacyStatsRule { if !is_zone_proof(parent)? { return Ok(None); } - let Some(converted) = LegacyStatsRule.reduce(child)? else { - return Ok(None); - }; Ok(Some( - EvalPlan::try_new(parent.expression().clone(), converted)?.into_plan(), + EvalPlan::try_new(parent.expression().clone(), to_zoned(child)?)?.into_plan(), )) } } diff --git a/vortex-layout/src/plan/physical/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs index 3d37390cfcb..4dd6e655999 100644 --- a/vortex-layout/src/plan/physical/plans/mod.rs +++ b/vortex-layout/src/plan/physical/plans/mod.rs @@ -29,8 +29,8 @@ pub use filter::FilterPlan; pub(crate) use legacy_stats::ExpressionLegacyStatsRule; pub use legacy_stats::LegacyStats; pub use legacy_stats::LegacyStatsData; +pub(crate) use legacy_stats::LegacyStatsDataRule; pub use legacy_stats::LegacyStatsPlan; -pub(crate) use legacy_stats::LegacyStatsRule; pub use list_pack::ListPack; pub use list_pack::ListPackData; pub use list_pack::ListPackPlan; diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index 3ead88f8892..40e82b4c855 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -516,7 +516,10 @@ fn expanded_struct_root( bound_pack(fields.names().clone(), children) } -fn is_identity_expression(expression: &BoundExpression, input_dtype: &DType) -> VortexResult { +pub(super) fn is_identity_expression( + expression: &BoundExpression, + input_dtype: &DType, +) -> VortexResult { if expression.is_root() { return Ok(expression.dtype() == input_dtype); } diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index 8b1729e953e..6ae9bec28d4 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -476,6 +476,9 @@ 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 { + if !eval.dtype().is_boolean() { + return Ok(false); + } let mut contains_stat = false; let mut contains_root = false; eval.expression().clone().transform_down(|expression| { @@ -493,7 +496,7 @@ pub(crate) fn is_zone_proof(eval: &Plan) -> VortexResult { contains_root |= expression.is_root(); Ok(Transformed::no(expression)) })?; - Ok(eval.dtype().is_boolean() && contains_stat && !contains_root) + Ok(contains_stat && !contains_root) } /// Replaces a zoned data plan by its data: a scan reads the data through it and never its zone diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 6698b6a0fe3..9a6c5852f65 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -646,12 +646,9 @@ fn empty_projection_prunes_row_idx_child_fields() -> VortexResult<()> { ); let plan = make_row_idx_plan(projection, make_plan(layout)?)?; + // The empty projection over the empty row domain is the identity, so only the domain remains. let optimized = optimize(plan)?; - let projection = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("optimized plan has no projection expression"))?; - let child = projection.child_plan()?; - let empty_struct = child + let empty_struct = optimized .as_opt::() .ok_or_else(|| vortex_err!("empty projection did not prune the RowIdx child"))?; From c1ce3af22166801eafeb7ed1fa2424c549a829ab Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 18:20:19 +0000 Subject: [PATCH 16/23] Expect segment ids as they display in the batch registration test Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-layout/src/scan/v2/io.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 5dfd21cf6af..943bfe9683b 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -344,7 +344,13 @@ mod tests { io.submit(IoOwnerId(0), batch)?; assert_eq!( *log.lock(), - ["register 0", "register 1", "register 2", "poll 0", "poll 1"] + [ + "register SegmentId(0)", + "register SegmentId(1)", + "register SegmentId(2)", + "poll SegmentId(0)", + "poll SegmentId(1)", + ] ); Ok(()) } From aae7554f8a18604495b9a406a80cad224caac6af Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 23:44:32 +0000 Subject: [PATCH 17/23] Project a chunk on one scan, a split and a batch per output window Each 64Ki-row window of a chunk was a morsel with a scan of its own, so a segment spanning several windows was decoded once per window. One morsel now runs one scan per chunk with a split per window: segments are decoded once, kernels run a window at a time, and a window's batch is yielded as soon as it has all its rows. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-layout/src/scan/execution/adaptor.rs | 21 +-- .../src/scan/execution/projection.rs | 135 +++++++++++------- vortex-layout/src/scan/v2/tests.rs | 47 ++++++ 3 files changed, 132 insertions(+), 71 deletions(-) diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index fc6053d55ca..85b834bb74d 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -3,7 +3,6 @@ //! Adapts pipeline execution to the planner I/O and compute protocol. -use std::ops::Range; use std::sync::Arc; use rustc_hash::FxHashMap; @@ -25,11 +24,11 @@ use crate::plan::physical::pipeline::Turn; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentId; -/// Rows a running plan produced, in row order. +/// Rows a running plan produced, in row order within its split. #[derive(Clone)] pub(super) struct Piece { - /// Orders the piece among the outputs of one scan. - pub(super) rows: Range, + /// The index of the split the rows belong to. + pub(super) split: usize, /// The selected rows the piece holds. pub(super) array: ArrayRef, } @@ -54,8 +53,6 @@ pub(super) struct ScanAdaptor { /// A turn already taken while batching reads, retained for the next compute. pending: Option, done: bool, - /// Orders the pieces handed out. - next_piece: u64, locations: Arc<[SegmentLocation]>, /// Published and undelivered requests, with the scan read and segment each one answers. outstanding: FxHashMap, @@ -72,7 +69,6 @@ impl ScanAdaptor { scan, pending: None, done: false, - next_piece: 0, locations, outstanding: FxHashMap::default(), failed: None, @@ -96,7 +92,7 @@ impl ScanAdaptor { } Ok(match self.step()? { Turn::Waiting | Turn::Done => AdaptorStep::Yield, - Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)), + Turn::Output(split, array) => AdaptorStep::Piece(Piece { split, array }), Turn::Read(read) => { // Publish every read the scan has ready together, so the owner can coalesce them. let mut batch = vec![self.request(read.id, read.segment_id)?]; @@ -130,15 +126,6 @@ impl ScanAdaptor { Ok(turn) } - fn piece(&mut self, array: ArrayRef) -> Piece { - let start = self.next_piece; - self.next_piece += 1; - Piece { - rows: start..start + 1, - array, - } - } - fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { let location = self.location(segment_id)?; let id = self.next_id; diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 8b238510f14..f12beb1d1fd 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; -use std::vec::IntoIter; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; @@ -19,7 +19,6 @@ use vortex_scan::planning::morsel::MorselOutput; 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 crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; @@ -27,50 +26,23 @@ use crate::plan::physical::selection::selected; use crate::scan::execution::PreparedScan; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; -use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; use crate::scan::execution::chunk::row_chunks; const MAX_OUTPUT_ROWS: usize = 65_536; -/// Keeps sparse selections together and subdivides dense selections to avoid large copies when -/// converting chunked projection results to Arrow. Subdivisions lie on a fixed grid of the file's -/// rows, so a writer chunk aligned to it is never split between two morsels. +/// Emits one morsel projecting a chunk's selected rows. pub(super) struct ProjectionPlanner { scan: Arc, - selected: IntoIter, + selected: Option, } impl ProjectionPlanner { /// Creates a planner for `selected`. pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { - let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { - vec![selected] - } else { - let offset = |row: u64| { - usize::try_from(row - selected.scope.rows.start) - .vortex_expect("a work chunk fits usize") - }; - row_chunks(selected.scope.rows.clone(), MAX_OUTPUT_ROWS) - .into_iter() - .filter_map(|rows| { - let mask = selected.mask.slice(offset(rows.start)..offset(rows.end)); - if mask.all_false() { - return None; - } - Some(SelectedRows { - scope: WorkScope { - file_ordinal: selected.scope.file_ordinal, - rows, - }, - mask, - }) - }) - .collect() - }; Self { scan, - selected: selected.into_iter(), + selected: Some(selected), } } } @@ -81,7 +53,7 @@ impl IoConsumer for ProjectionPlanner { impl Planner for ProjectionPlanner { fn state(&self) -> State { - if !self.selected.as_slice().is_empty() { + if self.selected.is_some() { State::NeedsCompute } else { State::Done @@ -89,7 +61,7 @@ impl Planner for ProjectionPlanner { } fn compute(&mut self) -> VortexResult { - let Some(selected) = self.selected.next() else { + let Some(selected) = self.selected.take() else { return Ok(PlannerOutput::Done); }; if selected.mask.all_false() { @@ -102,14 +74,24 @@ impl Planner for ProjectionPlanner { } } -/// Runs the projection plan over a split's selected rows and returns them as one batch, in row -/// order. A selection that projects to no rows finishes without a batch. +/// Runs the projection plan over a chunk's selected rows on one scan, so each segment is read +/// and decoded once, and yields a batch per window of the rows, in row order. +/// +/// Sparse selections stay one window. Dense selections are cut into windows of +/// `MAX_OUTPUT_ROWS` on a fixed grid of the file's rows, each a split of the scan, so kernels run +/// over a window at a time and converting a batch to Arrow copies no more than a window. struct ProjectionMorsel { scan: Arc, dtype: DType, selected: SelectedRows, adaptor: Option, - pieces: Vec, + /// The rows each split projects to, in split order. + split_rows: Vec, + /// The split whose pieces `pieces` holds. + split: usize, + pieces: Vec, + /// The rows `pieces` holds. + held: usize, done: bool, } @@ -121,25 +103,54 @@ impl ProjectionMorsel { scan, selected, adaptor: None, + split_rows: Vec::new(), + split: 0, pieces: Vec::new(), + held: 0, done: false, } } - fn finish(&mut self) -> VortexResult { - self.done = true; - self.pieces.sort_by_key(|piece| piece.rows.start); + /// The windows of the selected rows that select anything, each a split of the scan, with + /// the rows each projects to. + fn splits(&self) -> (Vec, Vec) { + let rows = &self.selected.scope.rows; + let mask = &self.selected.mask; + if mask.true_count() <= MAX_OUTPUT_ROWS { + let split = Split { + rows: rows.clone(), + mask: Mask::new_true(mask.len()), + }; + return (vec![split], vec![mask.true_count()]); + } + let offset = + |row: u64| usize::try_from(row - rows.start).vortex_expect("a work chunk fits usize"); + row_chunks(rows.clone(), MAX_OUTPUT_ROWS) + .into_iter() + .filter_map(|window| { + let (start, end) = (offset(window.start), offset(window.end)); + let selected = mask.slice(start..end).true_count(); + let split = Split { + rows: window, + mask: Mask::new_true(end - start), + }; + (selected > 0).then_some((split, selected)) + }) + .unzip() + } + + /// The pieces of the split held so far as one batch, or `None` when they hold no rows. + fn take_batch(&mut self) -> VortexResult> { + self.held = 0; let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) .into_iter() - .map(|piece| piece.array) .filter(|array| !array.is_empty()) .collect(); - let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), - 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), - }; - Ok(MorselOutput::Batch(array)) + Ok(match arrays.len() { + 0 => None, + 1 => Some(arrays.remove(0)), + _ => Some(ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array()), + }) } } @@ -168,6 +179,8 @@ impl Morsel for ProjectionMorsel { vortex_bail!("ProjectionMorsel: compute called after Done"); } let Some(adaptor) = self.adaptor.as_mut() else { + let (splits, split_rows) = self.splits(); + self.split_rows = split_rows; let scan = Scan::try_new( self.scan.session.clone(), selected( @@ -175,10 +188,7 @@ impl Morsel for ProjectionMorsel { self.selected.scope.rows.clone(), self.selected.mask.clone(), ), - vec![Split { - rows: self.selected.scope.rows.clone(), - mask: Mask::new_true(self.selected.mask.len()), - }], + splits, )? .with_row_offset(self.scan.row_offset) .with_value_cache(self.scan.values.clone()); @@ -186,14 +196,31 @@ impl Morsel for ProjectionMorsel { return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { - return self.finish(); + return Ok(match self.take_batch()? { + Some(batch) => MorselOutput::Batch(batch), + None => { + self.done = true; + MorselOutput::Done + } + }); } Ok(match adaptor.compute()? { AdaptorStep::Yield => MorselOutput::Continue, AdaptorStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), AdaptorStep::Piece(piece) => { - self.pieces.push(piece); - MorselOutput::Continue + // One split runs at a time, so a piece of a later split ends the one before it. + let mut batch = None; + if piece.split != self.split { + batch = self.take_batch()?; + self.split = piece.split; + } + self.held += piece.array.len(); + self.pieces.push(piece.array); + // A split that has projected all its rows yields without waiting for the next. + if batch.is_none() && self.split_rows.get(self.split) == Some(&self.held) { + batch = self.take_batch()?; + } + batch.map_or(MorselOutput::Continue, MorselOutput::Batch) } }) } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index eabf863eb47..46e0b81000b 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -673,6 +673,53 @@ async fn unordered_yields_before_the_split_finishes() -> VortexResult<()> { Ok(()) } +/// A dense chunk over one stored segment is projected by one morsel, a batch per output window. +#[tokio::test] +async fn dense_chunk_yields_a_batch_per_window() -> VortexResult<()> { + const ROWS: i32 = 150_000; + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0..ROWS).into_array(); + let dtype = array.dtype().clone(); + let layout = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter([Ok((sequence.advance(), array))])) + .sendable(), + eof, + &session, + ) + .await?; + let segments: Arc = segments; + let mut scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + Arc::get_mut(&mut scan.scan) + .ok_or_else(|| vortex_err!("the prepared scan is not shared yet"))? + .chunk_rows = usize::try_from(ROWS)?; + let batches = run_batches(&scan, 0..layout.row_count()).await?; + assert_eq!( + batches + .iter() + .map(|batch| batch.array.len()) + .collect::>(), + [65_536, 65_536, 18_928] + ); + let actual = ChunkedArray::try_new( + batches + .into_iter() + .map(|batch| batch.array) + .collect::>(), + layout.dtype().clone(), + )?; + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter(0..ROWS), + &mut session.create_execution_ctx() + ); + Ok(()) +} + #[tokio::test] async fn primitive_column() -> VortexResult<()> { let session = new_session().with_tokio(); From cc817bf7709623d8377fa4ea0b4b17f9eb6d76ce Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 10:28:06 +0000 Subject: [PATCH 18/23] Size a V2 scan's chunk concurrency by the partitions reading its file Concurrency 1 made a scan run its chunks inline, and the opener derived it from the engine's partition count. A file the engine does not split, such as a small dimension table, was then read by one partition with its chunks run one after another while other cores idled. The opener now divides the cores by the partitions that split this file, so a whole-file scan runs its chunks beside each other and a file split a partition per core still runs inline. A V2 scan's concurrency now counts its chunks in flight, one per worker thread by default, instead of per worker thread. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-datafusion/src/persistent/opener.rs | 24 +++++++++------------- vortex-datafusion/src/persistent/source.rs | 1 - vortex-layout/src/scan/v2/mod.rs | 9 ++++---- 3 files changed, 15 insertions(+), 19 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 26620ea2522..51b36789560 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -131,9 +131,6 @@ pub(crate) struct VortexOpener { pub projection_pushdown: bool, pub scan_concurrency: Option, pub scan_v2: bool, - /// How many partitions the engine runs the scan on, so a V2 scan parallelises inside a - /// partition only as far as the engine has not. - pub partitions: usize, } impl FileOpener for VortexOpener { @@ -172,7 +169,6 @@ impl FileOpener for VortexOpener { let file_schemas = Arc::clone(&self.file_schemas); let has_output_ordering = self.has_output_ordering; let scan_concurrency = self.scan_concurrency; - let partitions = self.partitions.max(1); let scan_v2 = self.scan_v2; let expr_converter = Arc::clone(&self.expression_converter); @@ -439,6 +435,11 @@ impl FileOpener for VortexOpener { .with_projection(scan_projection) .with_some_filter(native_filter) .with_ordered(has_output_ordering); + // How many partitions split the file, assuming they cut it evenly. + let readers = byte_range.as_ref().map_or(1, |bytes| { + let len = (bytes.end - bytes.start).max(1); + usize::try_from(file.object_meta.size.div_ceil(len)).unwrap_or(usize::MAX) + }); if let Some(bytes) = byte_range { // Every partition maps the same byte boundary to the same row boundary. let rows = partition_row_range(bytes, file.object_meta.size, vxf.row_count()); @@ -454,12 +455,13 @@ impl FileOpener for VortexOpener { { builder = builder.with_selection(selection.clone()); } - // Unless set, chunks run beside each other only as far as the engine's partitions - // leave cores idle: at one partition per core they run in the scan's own run. - let concurrency = scan_concurrency.unwrap_or_else(|| { + // Unless set, chunks run beside each other only as far as the cores outnumber the + // partitions reading this file: a file split a partition per core is read in each + // scan's own run, a file read whole by one partition on every core. + let concurrency = scan_concurrency.filter(|&c| c > 0).unwrap_or_else(|| { get_available_parallelism() .unwrap_or(1) - .div_ceil(partitions) + .div_ceil(readers) .max(1) }); builder = builder.with_concurrency(concurrency); @@ -1115,7 +1117,6 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), - partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1525,7 +1526,6 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), - partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1617,7 +1617,6 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), - partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1800,7 +1799,6 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), - partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1865,7 +1863,6 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), - partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -2081,7 +2078,6 @@ mod tests { layout_readers: Default::default(), segment_sources: Default::default(), scan_caches: Default::default(), - partitions: 1, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 644fea4e71c..8ff7b735cee 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -378,7 +378,6 @@ impl VortexSource { projection_pushdown: self.options.projection_pushdown, scan_concurrency: self.options.scan_concurrency, scan_v2: self.options.scan_v2, - partitions: base_config.file_groups.len(), }; Ok(opener) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 02b5507f213..ecce8cf7db3 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -135,7 +135,7 @@ impl ScanBuilder { ordered: true, row_range: None, selection: Selection::default(), - concurrency: 4, + concurrency: get_available_parallelism().unwrap_or(1), map_fn: Arc::new(Ok), row_offset: 0, } @@ -205,8 +205,9 @@ impl ScanBuilder { self } - /// Sets the number of chunks run at once per worker thread. Must be greater than zero. At one, - /// the chunks run on the scan's own run instead, for engines that parallelise across scans. + /// Sets the number of the scan's chunks run at once, one per worker thread by default. Must be + /// greater than zero. At one, the chunks run on the scan's own run instead, for engines that + /// parallelise across scans. pub fn with_concurrency(mut self, concurrency: usize) -> Self { assert!(concurrency > 0); self.concurrency = concurrency; @@ -336,7 +337,7 @@ impl ScanRun { /// 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 in_flight = self.concurrency; let root = self.scan.planner(self.rows.clone(), self.selection); let scope = WorkScope { file_ordinal: 0, From 9b71061665b8d4e9e453df448fb1d3af895128fe Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 12:04:37 +0000 Subject: [PATCH 19/23] Keep at least 256Ki rows in each V2 scan partition On a many-core machine DataFusion split SF1 lineitem into a partition per core, about 62K rows each, so every partition paid the V2 scan's fixed work (planning, zone proofs, compiling each conjunct) for one output window of rows. A V2 scan now caps the engine's target partitions by the files' row counts, so each partition reads at least 256Ki rows; the engine repartitions after the scan as before. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-datafusion/src/persistent/source.rs | 28 ++++++++++++++++++++++ vortex-datafusion/src/persistent/tests.rs | 7 +++--- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 8ff7b735cee..11ba970345c 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -49,6 +49,29 @@ use crate::convert::exprs::ExpressionConverter; use crate::persistent::reader::DefaultVortexReaderFactory; use crate::persistent::reader::VortexReaderFactory; +/// The fewest rows a V2 scan partition reads when the engine splits files. Each partition plans, +/// proves zone maps and compiles its plans once, so a partition much smaller than a few output +/// windows spends more on that than on reading rows. +const V2_MIN_ROWS_PER_PARTITION: usize = 256 * 1024; + +/// Caps `target_partitions` so no V2 partition reads fewer than [`V2_MIN_ROWS_PER_PARTITION`] +/// rows, when the files' row counts are known. +fn v2_target_partitions(target_partitions: usize, config: &FileScanConfig) -> usize { + let rows = config + .file_groups + .iter() + .flat_map(|group| group.iter()) + .map(|file| { + file.statistics + .as_ref() + .and_then(|statistics| statistics.num_rows.get_value().copied()) + }) + .sum::>(); + rows.map_or(target_partitions, |rows| { + target_partitions.min(rows.div_ceil(V2_MIN_ROWS_PER_PARTITION).max(1)) + }) +} + /// File scan implementation for reading one or more `.vortex` files. /// /// `VortexSource` is the lower-level read component underneath @@ -467,6 +490,11 @@ impl FileSource for VortexSource { if config.file_compression_type.is_compressed() { return Ok(None); } + let target_partitions = if self.options.scan_v2 { + v2_target_partitions(target_partitions, config) + } else { + target_partitions + }; let Some(file_groups) = FileGroupPartitioner::new() .with_target_partitions(target_partitions) .with_repartition_file_min_size(repartition_file_min_size) diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 5d91242c91d..b4d1475ed8f 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -486,9 +486,10 @@ async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits( let session = VortexSession::default(); let path = object_store::path::Path::parse("/split-aligned-repartition.vortex")?; - let chunk_1_len = 2_000; - let chunk_2_len = 5_000; - let chunk_3_len = 6_000; + // Enough rows that a V2 scan, which keeps a minimum of rows per partition, still splits. + let chunk_1_len = 200_000; + let chunk_2_len = 250_000; + let chunk_3_len = 300_000; let row_count = chunk_1_len + chunk_2_len + chunk_3_len; let chunk_1 = StructArray::try_new( From 5f20be003f6cb4c010165407112b298b7322984f Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 12:20:18 +0000 Subject: [PATCH 20/23] Size a V2 scan's filtered chunks to keep its concurrency busy A filtered chunk was always 1Mi rows, so a partition of a few hundred thousand rows ran one chunk at a time however many chunks its concurrency allowed, leaving cores idle on CPU-heavy predicates. A filtered chunk now takes the scan's rows divided by its concurrency, rounded to the row grid and kept between 64Ki and 1Mi rows. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX Signed-off-by: Claude --- vortex-layout/src/scan/execution/chunk.rs | 12 ++++++++++++ vortex-layout/src/scan/execution/mod.rs | 2 +- vortex-layout/src/scan/v2/mod.rs | 7 ++++--- vortex-layout/src/scan/v2/tests.rs | 11 +++++++++++ 4 files changed, 28 insertions(+), 4 deletions(-) diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 6808c707c48..22c55e984c1 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -44,6 +44,18 @@ pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> u blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE } +/// The rows of a filtered chunk for a scan of `rows` running `concurrency` chunks at once: as +/// large as [`FILTERED_ROWS_PER_CHUNK`], but small enough that every chunk the scan can run at +/// once has rows to filter, and never below [`MAX_ROWS_PER_CHUNK`]. +pub(crate) fn filtered_rows_per_chunk(rows: u64, concurrency: usize) -> usize { + let per_chunk = rows.div_ceil(concurrency.max(1) as u64); + let blocks = per_chunk.div_ceil(ROW_BLOCK_SIZE as u64); + let rows = usize::try_from(blocks) + .unwrap_or(usize::MAX) + .saturating_mul(ROW_BLOCK_SIZE); + rows.clamp(MAX_ROWS_PER_CHUNK, FILTERED_ROWS_PER_CHUNK) +} + /// Fixed boundaries in the file's row domain, clipped to the requested range. pub(crate) fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { if rows.is_empty() { diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index e371fdb539c..5a28ad788a7 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -30,9 +30,9 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; pub(in crate::scan) 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; +pub(crate) use self::chunk::filtered_rows_per_chunk; #[cfg(test)] pub(crate) use self::chunk::row_chunks; pub(crate) use self::chunk::rows_per_chunk; diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index ecce8cf7db3..3c86a6106b2 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -67,12 +67,12 @@ 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; -use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; use crate::scan::execution::FilterPlans; use crate::scan::execution::MAX_ROWS_PER_CHUNK; use crate::scan::execution::PreparedScan; use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; +use crate::scan::execution::filtered_rows_per_chunk; use crate::scan::execution::rows_per_chunk; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; @@ -295,8 +295,9 @@ impl ScanRun { root.clone(), )?) })?; + let rows = builder.row_range.clone().unwrap_or(0..row_count); let chunk_rows = if filter.is_some() { - FILTERED_ROWS_PER_CHUNK + filtered_rows_per_chunk(rows.end - rows.start, builder.concurrency) } else if identity { rows_per_chunk(row_count, &file.locations) } else { @@ -316,7 +317,7 @@ impl ScanRun { chunk_rows, }), io, - rows: builder.row_range.unwrap_or(0..row_count), + rows, selection: builder.selection, map_fn: builder.map_fn, ordered: builder.ordered, diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index b01eb48f8a4..ded60947e6a 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -72,6 +72,7 @@ use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; use crate::scan::execution::SegmentLocation; +use crate::scan::execution::filtered_rows_per_chunk; use crate::scan::execution::row_chunks; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; use crate::segments::SegmentFuture; @@ -465,6 +466,16 @@ fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: us assert_eq!(rows_per_chunk(rows, &locations), expected); } +/// A filtered chunk shrinks so every chunk a scan runs at once has rows, within its bounds. +#[rstest] +#[case::one_at_a_time(6_000_000, 1, 1_048_576)] +#[case::split_across_concurrency(262_144, 4, 65_536)] +#[case::rounded_to_blocks(300_000, 4, 81_920)] +#[case::at_least_a_window(100_000, 96, 65_536)] +fn filtered_chunk_rows(#[case] rows: u64, #[case] concurrency: usize, #[case] expected: usize) { + assert_eq!(filtered_rows_per_chunk(rows, concurrency), expected); +} + #[test] fn smaller_chunks_keep_file_alignment() { assert_eq!( From 09d0360984a488d9b6a5ef4bf0f6228a0d0793fe Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 22:38:48 +0000 Subject: [PATCH 21/23] Size a V2 scan's projection chunks to keep its concurrency busy A partition keeps at least 256Ki rows, so a small, wide table such as TPC-H SF10's customer (1.5M rows) is read by 6 partitions instead of one per core. Each ran 4 projection chunks of 64Ki rows, so only 24 chunks decoded the table on 96 cores, and Q10 ran 19-24% slower than develop. A projection chunk is now cut like a filtered one: small enough that every chunk the scan runs at once has rows, down to one row block, and never larger than before. Also correct ProjectionMorsel's doc: each split decodes the segments it reads; a chunk's splits share the reads and dictionary values, not the decoded segments. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX --- vortex-layout/src/scan/execution/chunk.rs | 17 +++++++++++--- vortex-layout/src/scan/execution/mod.rs | 1 + .../src/scan/execution/projection.rs | 5 ++-- vortex-layout/src/scan/v2/mod.rs | 10 +++++--- vortex-layout/src/scan/v2/tests.rs | 23 ++++++++++++++++++- 5 files changed, 47 insertions(+), 9 deletions(-) diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 3575ed1f8fe..3128b2b93cc 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -48,12 +48,23 @@ pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> u /// large as [`FILTERED_ROWS_PER_CHUNK`], but small enough that every chunk the scan can run at /// once has rows to filter, and never below [`MAX_ROWS_PER_CHUNK`]. pub(crate) fn filtered_rows_per_chunk(rows: u64, concurrency: usize) -> usize { + rows_per_concurrent_chunk(rows, concurrency).clamp(MAX_ROWS_PER_CHUNK, FILTERED_ROWS_PER_CHUNK) +} + +/// The rows of a projection chunk for a scan of `rows` running `concurrency` chunks at once: as +/// large as `max_rows`, but small enough that every chunk the scan can run at once has rows to +/// project, and never below [`ROW_BLOCK_SIZE`]. +pub(crate) fn projected_rows_per_chunk(rows: u64, concurrency: usize, max_rows: usize) -> usize { + rows_per_concurrent_chunk(rows, concurrency).clamp(ROW_BLOCK_SIZE, max_rows) +} + +/// `rows` split evenly over `concurrency` chunks, rounded up to the row grid. +fn rows_per_concurrent_chunk(rows: u64, concurrency: usize) -> usize { let per_chunk = rows.div_ceil(concurrency.max(1) as u64); let blocks = per_chunk.div_ceil(ROW_BLOCK_SIZE as u64); - let rows = usize::try_from(blocks) + usize::try_from(blocks) .unwrap_or(usize::MAX) - .saturating_mul(ROW_BLOCK_SIZE); - rows.clamp(MAX_ROWS_PER_CHUNK, FILTERED_ROWS_PER_CHUNK) + .saturating_mul(ROW_BLOCK_SIZE) } /// Fixed boundaries in the file's row domain, clipped to the requested range. diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 9ab1baa91cf..0a9ff0b7679 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -34,6 +34,7 @@ pub use self::chunk::ChunkPlanner; pub(crate) use self::chunk::MAX_ROWS_PER_CHUNK; pub(crate) use self::chunk::ROW_BLOCK_SIZE; pub(crate) use self::chunk::filtered_rows_per_chunk; +pub(crate) use self::chunk::projected_rows_per_chunk; #[cfg(test)] pub(crate) use self::chunk::row_chunks; pub(crate) use self::chunk::rows_per_chunk; diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index f12beb1d1fd..f3da4b4218c 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -74,8 +74,9 @@ impl Planner for ProjectionPlanner { } } -/// Runs the projection plan over a chunk's selected rows on one scan, so each segment is read -/// and decoded once, and yields a batch per window of the rows, in row order. +/// Runs the projection plan over a chunk's selected rows on one scan, so the chunk compiles one +/// plan and shares its reads and dictionary values, and yields a batch per window of the rows, +/// in row order. /// /// Sparse selections stay one window. Dense selections are cut into windows of /// `MAX_OUTPUT_ROWS` on a fixed grid of the file's rows, each a split of the scan, so kernels run diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 1d7465f05d4..d933ede8ac2 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -53,6 +53,7 @@ pub use crate::scan::execution::PrunePlanner; use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; use crate::scan::execution::filtered_rows_per_chunk; +use crate::scan::execution::projected_rows_per_chunk; use crate::scan::execution::rows_per_chunk; use crate::scan::v2::io::segment_ranges; use crate::segments::SegmentSource; @@ -194,10 +195,13 @@ impl ScanBuilder { let rows = self.row_range.unwrap_or(0..row_count); let chunk_rows = if filter.is_some() { filtered_rows_per_chunk(rows.end - rows.start, self.concurrency) - } else if identity { - rows_per_chunk(row_count, &file.locations) } else { - MAX_ROWS_PER_CHUNK + let max_rows = if identity { + rows_per_chunk(row_count, &file.locations) + } else { + MAX_ROWS_PER_CHUNK + }; + projected_rows_per_chunk(rows.end - rows.start, self.concurrency, max_rows) }; let ranges = cache.ranges_or_init(|| segment_ranges(&file.locations)); let values = cache.values(); diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 2e6a5fbb419..9015e37b9a9 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -73,6 +73,7 @@ use crate::layouts::zoned::writer::ZonedStrategy; use crate::scan::arrow::to_record_batch; use crate::scan::execution::SegmentLocation; use crate::scan::execution::filtered_rows_per_chunk; +use crate::scan::execution::projected_rows_per_chunk; use crate::scan::execution::row_chunks; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; use crate::segments::SegmentFuture; @@ -506,6 +507,21 @@ fn filtered_chunk_rows(#[case] rows: u64, #[case] concurrency: usize, #[case] ex assert_eq!(filtered_rows_per_chunk(rows, concurrency), expected); } +/// A projection chunk shrinks so every chunk a scan runs at once has rows, within its bounds: +/// the 1.5M rows of TPC-H SF10's `customer` cut into 6 partitions of 16 chunks each must not +/// run 4 chunks of 64Ki rows each. +#[rstest] +#[case::one_at_a_time(6_000_000, 1, 65_536)] +#[case::split_across_concurrency(250_000, 16, 16_384)] +#[case::rounded_to_blocks(100_000, 4, 32_768)] +#[case::at_least_a_block(10_000, 96, 8_192)] +fn projected_chunk_rows(#[case] rows: u64, #[case] concurrency: usize, #[case] expected: usize) { + assert_eq!( + projected_rows_per_chunk(rows, concurrency, MAX_ROWS_PER_CHUNK), + expected + ); +} + #[test] fn smaller_chunks_keep_file_alignment() { assert_eq!( @@ -561,7 +577,12 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); + // One chunk at a time, so chunks are as large as a projection chunk on any machine. + let scan = Arc::new( + builder(&session, &segments, &layout)? + .with_concurrency(1) + .prepare()?, + ); let chunks = row_chunks(scan.rows().clone(), scan.chunk_rows); assert_eq!( chunks.len(), From 6a4c019c522b03b6d3448b692712bddc69972745 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 23:33:01 +0000 Subject: [PATCH 22/23] Keep the engine's partitions for an unfiltered V2 scan The minimum rows per V2 partition cut TPC-H SF10's customer from one partition per core to 6 and part to 8. Every row of an unfiltered scan flows into the operators above it, which run a task per partition, so Q10 and Q14 ran 13-17% slower than develop on 96 cores however many chunks each partition ran at once. Forcing customer and part to one partition reproduces it locally: Q10 and Q14 at SF10 are 13% slower on 4 cores. The minimum now applies only to a filtered scan, whose partitions each prove and plan the filter and emit a fraction of the rows they read. Locally the SF10 slowdown is gone, and TPC-H SF1 split 96 ways keeps most of what the minimum saves: a geomean of 87 ms, against 83 ms with the minimum on every scan and 104 ms without it. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GwhHqFVpQqPgQQkkK2y9kX --- vortex-datafusion/src/persistent/source.rs | 10 ++++--- vortex-datafusion/src/persistent/tests.rs | 32 ++++++++++++++++++++++ 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 11ba970345c..c738df95550 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -49,9 +49,11 @@ use crate::convert::exprs::ExpressionConverter; use crate::persistent::reader::DefaultVortexReaderFactory; use crate::persistent::reader::VortexReaderFactory; -/// The fewest rows a V2 scan partition reads when the engine splits files. Each partition plans, -/// proves zone maps and compiles its plans once, so a partition much smaller than a few output -/// windows spends more on that than on reading rows. +/// The fewest rows a filtered V2 scan partition reads when the engine splits files. Each +/// partition proves zone maps, plans its filter and compiles its plans once, so a partition much +/// smaller than a few output windows spends more on that than on reading rows. An unfiltered scan +/// keeps the engine's partitions: every row it reads flows into the operators above it, which +/// run a task per partition. const V2_MIN_ROWS_PER_PARTITION: usize = 256 * 1024; /// Caps `target_partitions` so no V2 partition reads fewer than [`V2_MIN_ROWS_PER_PARTITION`] @@ -490,7 +492,7 @@ impl FileSource for VortexSource { if config.file_compression_type.is_compressed() { return Ok(None); } - let target_partitions = if self.options.scan_v2 { + let target_partitions = if self.options.scan_v2 && self.vortex_predicate.is_some() { v2_target_partitions(target_partitions, config) } else { target_partitions diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index b4d1475ed8f..15509f5f941 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -583,6 +583,38 @@ async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits( Ok(()) } +/// A filtered V2 scan of a file smaller than the minimum rows per partition reads it in one +/// partition. An unfiltered scan keeps the engine's partitions: every row it reads flows into the +/// operators above it. +#[tokio::test] +async fn v2_scan_keeps_partitions_unless_filtered() -> anyhow::Result<()> { + let store = Arc::new(InMemory::new()) as _; + let session = VortexSession::default(); + let path = object_store::path::Path::parse("/small.vortex")?; + let table = StructArray::try_new( + ["value"].into(), + vec![Buffer::from_iter(0_i32..100_000).into_array()], + 100_000, + Validity::NonNullable, + )? + .into_array(); + let mut writer = ObjectStoreWrite::new(Arc::clone(&store), &path).await?; + session + .write_options() + .write(&mut writer, table.to_array_stream()) + .await?; + writer.shutdown().await?; + + let ctx = make_session(store, true, true); + let unfiltered = count_query_partitions(&ctx, "SELECT value FROM '/small.vortex'").await?; + let filtered = + count_query_partitions(&ctx, "SELECT value FROM '/small.vortex' WHERE value > 10").await?; + + assert_eq!(unfiltered, 4); + assert_eq!(filtered, 1); + Ok(()) +} + /// Roundtrip an `arrow.uuid` extension column through a Vortex file: write the column directly /// via the session-aware Arrow→Vortex conversion, then `SELECT *` and assert both the field /// metadata and the underlying values survive the trip. From 143c86efa39628ce462dc50f946b728f736f4193 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 17:18:48 +0000 Subject: [PATCH 23/23] DCO Remediation Commit for Claude I, Claude , hereby add my Signed-off-by to this commit: 09d0360984a488d9b6a5ef4bf0f6228a0d0793fe I, Claude , hereby add my Signed-off-by to this commit: 6a4c019c522b03b6d3448b692712bddc69972745 Signed-off-by: Claude