From c56ec58c71a5f99983134960d6306f2396af9e03 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 22 Oct 2026 12:28:04 +0100 Subject: [PATCH 01/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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(()) }