diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index fbdeb027554..cb22673cb6a 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -3,22 +3,25 @@ use std::io::Cursor; use std::path::Path; -use std::sync::Arc; use std::time::Duration; use std::time::Instant; use anyhow::Result; use anyhow::bail; +use arrow_schema::Field; use async_trait::async_trait; use bytes::Bytes; use futures::StreamExt; use futures::pin_mut; use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; use vortex::dtype::FieldNames; use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; +use vortex::file::v2::ScanDriver; +use vortex::io::session::RuntimeSessionExt; use vortex::layout::scan::v2::ScanBuilder; use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; @@ -96,13 +99,17 @@ async fn vortex_decompress_read(data: Bytes) -> Result<()> { .optimize_recursive()?; scan = scan.with_projection(projection); } - let schema = Arc::new(SESSION.arrow().to_arrow_schema(scan.dtype())?); + let schema = SESSION.arrow().to_arrow_schema(scan.dtype())?; + let field = Field::new_struct("", schema.fields().clone(), false); - let stream = scan.into_record_batch_stream(schema)?; + let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan); pin_mut!(stream); - while let Some(batch) = stream.next().await { - let _batch = batch?; + while let Some(chunk) = stream.next().await { + let mut ctx = SESSION.create_execution_ctx(); + let _batch = SESSION + .arrow() + .execute_arrow(chunk?, Some(&field), &mut ctx)?; } Ok(()) } diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 164c402ecc9..f57ee614c08 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -35,6 +35,8 @@ use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; +use vortex::file::v2::ScanDriver; +use vortex::io::session::RuntimeSessionExt; use vortex::layout::scan::v2::ScanBuilder; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; @@ -190,9 +192,9 @@ impl RandomAccessor for VortexRandomAccessor { async fn take(&self, indices: &[u64]) -> anyhow::Result { let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?; - let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file()) - .with_row_indices(indices) - .into_stream()?; + let scan = + ScanBuilder::new(SESSION.clone(), self.file.scan_file()).with_row_indices(indices); + let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan); let array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream) .read_all() .await?; diff --git a/vortex-datafusion/src/persistent/mod.rs b/vortex-datafusion/src/persistent/mod.rs index 53fefacafa4..1df70f35edc 100644 --- a/vortex-datafusion/src/persistent/mod.rs +++ b/vortex-datafusion/src/persistent/mod.rs @@ -28,6 +28,7 @@ mod format; pub mod metrics; mod opener; pub mod reader; +mod scan_v2; mod sink; mod source; mod stream; diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 26620ea2522..208691f0b0e 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -53,15 +53,20 @@ use vortex::error::VortexExpect; use vortex::expr::and_collect; use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; +use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::layout::scan::v2::ChunkPlanner; +use vortex::layout::scan::v2::PrunePlanner; use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; use vortex::layout::scan::v2::ScanCache; use vortex::layout::scan::v2::WeakScanCache; use vortex::layout::scan::v2::partition_row_range; use vortex::layout::segments::SegmentSource; +use vortex::mask::Mask; use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; +use vortex::scan::planning::planner::Planner; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; use vortex_utils::aliases::dash_map::DashMap; @@ -77,6 +82,7 @@ use crate::metrics::PARTITION_LABEL; use crate::metrics::PATH_LABEL; use crate::persistent::cache::CachedVortexMetadata; use crate::persistent::reader::VortexReaderFactory; +use crate::persistent::scan_v2::ScanDriver; use crate::persistent::stream::PrunableStream; #[derive(Clone)] @@ -116,7 +122,7 @@ pub(crate) struct VortexOpener { /// shared layout reader. pub segment_sources: Arc>>, /// The V2 scan cache each open file's partitions share: its lowered layout, filter and - /// projection plans, proved zones and dictionary values, built once per file. + /// projection plans, pruned zones and dictionary values, built once per file. pub scan_caches: Arc>, /// Shared V1 full-file natural splits keyed by file path. pub natural_splits: Arc>>, @@ -437,8 +443,7 @@ impl FileOpener for VortexOpener { scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta.location)); let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file) .with_projection(scan_projection) - .with_some_filter(native_filter) - .with_ordered(has_output_ordering); + .with_some_filter(native_filter); if let Some(bytes) = byte_range { // Every partition maps the same byte boundary to the same row boundary. let rows = partition_row_range(bytes, file.object_meta.size, vxf.row_count()); @@ -462,8 +467,35 @@ impl FileOpener for VortexOpener { .div_ceil(partitions) .max(1) }); - builder = builder.with_concurrency(concurrency); - builder.map(map_batch).into_stream() + let handle = vxf.session().handle(); + // Lowering and planning a wide file is real CPU: it runs where compute runs. + let scan = handle + .spawn_cpu(move || builder.prepare()) + .await + .map_err(|e| exec_datafusion_err!("Failed to plan Vortex scan: {e}"))?; + let scan = Arc::new(scan); + + // The planners, composed: with a filter, its pruning plans first over the whole + // partition, continuing into the chunks of the rows they keep; without a + // filter, the chunks directly. A chunk plans its conjuncts, then its + // projection. + let chunks = { + let scan = Arc::clone(&scan); + move |kept: Option| -> Box { + let rows = scan.rows().clone(); + let selection = scan.selection().clone(); + Box::new(ChunkPlanner::new(scan, rows, selection, kept)) + } + }; + let root: Box = if scan.filter().is_some() { + let rows = scan.rows().clone(); + Box::new(PrunePlanner::new(Arc::clone(&scan), rows, Box::new(chunks))) + } else { + chunks(None) + }; + + let driver = ScanDriver::new(handle, concurrency, has_output_ordering); + Ok(driver.run(scan, root, map_batch)) } else { // Share V1 layout readers across partitions of the file. let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { diff --git a/vortex-datafusion/src/persistent/scan_v2.rs b/vortex-datafusion/src/persistent/scan_v2.rs new file mode 100644 index 00000000000..10ea7b969ea --- /dev/null +++ b/vortex-datafusion/src/persistent/scan_v2.rs @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Drives a V2 scan for DataFusion. +//! +//! A planner says what it needs: compute, or reads. This driver decides where each runs. The +//! opener composes the planners over a [`PreparedScan`] and hands the driver the first; the +//! driver runs it and everything it emits on planning runs it owns: the root run stepped on +//! `spawn_cpu` tasks as the stream is polled, each spawned chunk on a task of its own, or on the +//! root run when the engine's partitions already fill the cores. Every read is a `spawn_io` task +//! of the file's segment source. The DataFusion thread polling the stream does no scan work +//! itself. + +use std::future; +use std::sync::Arc; +use std::task::Poll; + +use arrow_array::RecordBatch; +use futures::SinkExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::channel::mpsc; +use futures::stream; +use futures::stream::BoxStream; +use vortex::array::ArrayRef; +use vortex::error::VortexResult; +use vortex::error::vortex_err; +use vortex::io::request::IoSource; +use vortex::io::runtime::Handle; +use vortex::layout::scan::v2::PreparedScan; +use vortex::scan::planning::driver::Progress; +use vortex::scan::planning::driver::Run; +use vortex::scan::planning::planner::Planner; +use vortex::scan::planning::planner::WorkScope; +use vortex_utils::parallelism::get_available_parallelism; + +/// Maps a scan's batches to Arrow on the task that produced them. +type MapBatch = Arc VortexResult + Send + Sync>; + +/// Runs one partition's V2 scan on DataFusion's runtime. +pub(crate) struct ScanDriver { + handle: Handle, + /// Chunks run beside each other per worker thread; at 1 they run on the scan's own run. + concurrency: usize, + /// Whether batches are yielded in row order, which holds them until the scan is done. + ordered: bool, +} + +impl ScanDriver { + pub(crate) fn new(handle: Handle, concurrency: usize, ordered: bool) -> Self { + Self { + handle, + concurrency, + ordered, + } + } + + /// Runs `root`, the first planner composed over `scan`, and everything it emits, and + /// streams their batches mapped by `map`. + pub(crate) fn run( + self, + scan: Arc, + root: Box, + map: impl Fn(ArrayRef) -> VortexResult + Send + Sync + 'static, + ) -> BoxStream<'static, VortexResult> { + let map: MapBatch = Arc::new(map); + let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); + let inline = in_flight <= 1; + let scope = WorkScope { + file_ordinal: 0, + rows: scan.rows().clone(), + }; + let io = scan.io(); + let (chunk_handle, chunk_io, chunk_map) = + (self.handle.clone(), Arc::clone(&io), Arc::clone(&map)); + let batches = run_stream(self.handle, root, scope, io.session(), map, inline) + .map_ok(move |yielded| match yielded { + Yielded::Batch(start, batch) => { + stream::once(future::ready(Ok((start, batch)))).boxed() + } + Yielded::Spawn(scope, chunk) => chunk_stream( + &chunk_handle, + chunk, + scope, + chunk_io.session(), + Arc::clone(&chunk_map), + ), + }) + .try_flatten_unordered(in_flight); + // The prepared scan holds the file's shared cache: keep it until the stream ends, so + // scans of the file that start meanwhile find it. + let batches = batches.chain(stream::poll_fn(move |_| { + let _ = &scan; + Poll::Ready(None) + })); + if !self.ordered { + return batches.map_ok(|(_, batch)| batch).boxed(); + } + stream::once(async move { + let mut held: Vec<(u64, RecordBatch)> = batches.try_collect().await?; + // Morsels finish in whatever order their reads arrive. + held.sort_by_key(|(start, _)| *start); + VortexResult::Ok(stream::iter(held.into_iter().map(|(_, batch)| Ok(batch)))) + }) + .try_flatten() + .boxed() + } +} + +/// What the root run's stream yields: a mapped batch with the row it starts at, or a chunk a +/// planner of the run spawned. +enum Yielded { + Batch(u64, RecordBatch), + Spawn(WorkScope, Box), +} + +/// One CPU step of a run, taken on a `spawn_cpu` task with the run. +enum Step { + Batch(u64, RecordBatch), + Spawn(WorkScope, Box), + RootDone, + Idle, + Waiting, +} + +/// The state of the root run's stream. +struct RunState { + handle: Handle, + /// Taken while a CPU step runs on its task and put back with its result. + run: Option, + io: Arc, + map: MapBatch, + /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. + inline: bool, +} + +/// Runs `root` over `scope` on a planning run of its own, reading through `io`, each CPU step +/// on a `spawn_cpu` task, and yields its mapped batches and, unless `inline`, the chunks it +/// spawns. Inline, a spawned chunk runs on this run beside the others: the run's steps already +/// interleave every chunk's work, and a run per chunk would only add a hand-off per chunk for +/// parallelism the engine provides. +fn run_stream( + handle: Handle, + root: Box, + scope: WorkScope, + io: Arc, + map: MapBatch, + inline: bool, +) -> BoxStream<'static, VortexResult> { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + let state = RunState { + handle, + run: Some(run), + io, + map, + inline, + }; + stream::try_unfold(state, |mut state| async move { + loop { + let run = state + .run + .take() + .ok_or_else(|| vortex_err!("the scan's run is missing"))?; + let map = Arc::clone(&state.map); + let (run, step) = state + .handle + .spawn_cpu(move || { + let mut run = run; + let step = run.advance().and_then(|progress| { + Ok(match progress { + Progress::Batch(batch) => { + Step::Batch(batch.scope.rows.start, map(batch.array)?) + } + Progress::Spawn(scope, child) => Step::Spawn(scope, child), + Progress::RootDone(_) => Step::RootDone, + Progress::Idle => Step::Idle, + Progress::Waiting => Step::Waiting, + }) + }); + (run, step) + }) + .await; + let run = state.run.insert(run); + match step? { + Step::Batch(start, batch) => { + return Ok(Some((Yielded::Batch(start, batch), state))); + } + Step::Spawn(scope, child) if state.inline => { + run.admit(child, scope, Arc::clone(&state.io)); + } + Step::Spawn(scope, child) => { + return Ok(Some((Yielded::Spawn(scope, child), state))); + } + Step::RootDone if run.live_roots() > 0 => {} + Step::RootDone | Step::Idle => return Ok(None), + Step::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }) + .boxed() +} + +/// Runs a spawned chunk's `root` over `scope` on a planning run of its own, on one task that +/// drives the run to its end and awaits only its reads, and yields its mapped batches with the +/// row each starts at. A planner the chunk spawns in turn runs as another root of the same run. +/// +/// One task per chunk, not one per step: a chunk's run waits for reads many times, and a task +/// hop at each wait is latency on the chunk's critical path that no other chunk's work fills. +fn chunk_stream( + handle: &Handle, + root: Box, + scope: WorkScope, + io: Arc, + map: MapBatch, +) -> BoxStream<'static, VortexResult<(u64, RecordBatch)>> { + let (mut send, recv) = mpsc::channel(1); + let task = handle.spawn(async move { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + loop { + match run.advance()? { + Progress::Batch(batch) => { + let start = batch.scope.rows.start; + let batch = map(batch.array); + let failed = batch.is_err(); + if send.send(batch.map(|batch| (start, batch))).await.is_err() || failed { + return Ok(()); + } + } + Progress::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&io)); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => return Ok(()), + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }); + let errors = stream::once(task) + .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); + stream::select(recv, errors).boxed() +} diff --git a/vortex-file/src/v2/mod.rs b/vortex-file/src/v2/mod.rs index 988916b65d1..f9ffd62549d 100644 --- a/vortex-file/src/v2/mod.rs +++ b/vortex-file/src/v2/mod.rs @@ -2,5 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod file_stats_reader; +mod scan; pub use file_stats_reader::*; +pub use scan::ScanDriver; diff --git a/vortex-file/src/v2/scan.rs b/vortex-file/src/v2/scan.rs new file mode 100644 index 00000000000..14753e58777 --- /dev/null +++ b/vortex-file/src/v2/scan.rs @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Drives a V2 scan of a file for callers without an engine of their own. +//! +//! A planner says what it needs: compute, or reads. This driver decides where each runs. It +//! composes the planners over a [`PreparedScan`], the pruning plans first when the scan has a +//! filter and the chunks of the rows they keep after, and runs them on planning runs it owns: +//! the root run stepped on `spawn_cpu` tasks as the stream is polled, each spawned chunk on a +//! task of its own, up to the concurrency, or on the root run at a concurrency of one. Every +//! read is a `spawn_io` task of the file's segment source. + +use std::future; +use std::sync::Arc; +use std::task::Poll; + +use futures::SinkExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::channel::mpsc; +use futures::stream; +use futures::stream::BoxStream; +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::request::IoSource; +use vortex_io::runtime::Handle; +use vortex_layout::scan::v2::ChunkPlanner; +use vortex_layout::scan::v2::PreparedScan; +use vortex_layout::scan::v2::PrunePlanner; +use vortex_layout::scan::v2::ScanBuilder; +use vortex_mask::Mask; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::WorkScope; +use vortex_utils::parallelism::get_available_parallelism; + +/// Runs a V2 scan on the session's runtime. +pub struct ScanDriver { + handle: Handle, + concurrency: usize, + ordered: bool, +} + +impl ScanDriver { + /// A driver spawning on `handle`, running `concurrency` chunks beside each other per worker + /// thread, or every chunk on the scan's own run at one, and yielding batches in row order + /// when `ordered`, which holds them until the scan is done. + pub fn new(handle: Handle, concurrency: usize, ordered: bool) -> Self { + assert!(concurrency > 0); + Self { + handle, + concurrency, + ordered, + } + } + + /// Prepares `builder` where compute runs, composes its planners, and streams their batches. + /// Preparation errors are the stream's first item. + pub fn run(self, builder: ScanBuilder) -> BoxStream<'static, VortexResult> { + let handle = self.handle.clone(); + stream::once(async move { + handle + .spawn_cpu(move || { + let scan = Arc::new(builder.prepare()?); + VortexResult::Ok(self.execute(scan)) + }) + .await + }) + .try_flatten() + .boxed() + } + + fn execute(self, scan: Arc) -> BoxStream<'static, VortexResult> { + let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); + let inline = in_flight <= 1; + let scope = WorkScope { + file_ordinal: 0, + rows: scan.rows().clone(), + }; + let io = scan.io(); + let root = compose(&scan); + let (chunk_handle, chunk_io) = (self.handle.clone(), Arc::clone(&io)); + let batches = run_stream(self.handle, root, scope, io.session(), inline) + .map_ok(move |yielded| match yielded { + Yielded::Batch(start, batch) => { + stream::once(future::ready(Ok((start, batch)))).boxed() + } + Yielded::Spawn(scope, chunk) => { + chunk_stream(&chunk_handle, chunk, scope, chunk_io.session()) + } + }) + .try_flatten_unordered(in_flight); + // The prepared scan holds the file's shared cache: keep it until the stream ends, so + // scans of the file that start meanwhile find it. + let batches = batches.chain(stream::poll_fn(move |_| { + let _ = &scan; + Poll::Ready(None) + })); + if !self.ordered { + return batches.map_ok(|(_, batch)| batch).boxed(); + } + stream::once(async move { + let mut held: Vec<(u64, ArrayRef)> = batches.try_collect().await?; + // Morsels finish in whatever order their reads arrive. + held.sort_by_key(|(start, _)| *start); + VortexResult::Ok(stream::iter(held.into_iter().map(|(_, batch)| Ok(batch)))) + }) + .try_flatten() + .boxed() + } +} + +/// The planners of a scan, composed: with a filter, its pruning plans first over the whole scan, +/// continuing into the chunks of the rows they keep; without, the chunks directly. +fn compose(scan: &Arc) -> Box { + let rows = scan.rows().clone(); + let chunks = { + let scan = Arc::clone(scan); + let rows = rows.clone(); + move |kept: Option| -> Box { + let selection = scan.selection().clone(); + Box::new(ChunkPlanner::new(scan, rows, selection, kept)) + } + }; + if scan.filter().is_some() { + Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(chunks))) + } else { + chunks(None) + } +} + +/// What the root run's stream yields: a batch with the row it starts at, or a chunk a planner of +/// the run spawned. +enum Yielded { + Batch(u64, ArrayRef), + Spawn(WorkScope, Box), +} + +/// One CPU step of a run, taken on a `spawn_cpu` task with the run. +enum Step { + Batch(u64, ArrayRef), + Spawn(WorkScope, Box), + RootDone, + Idle, + Waiting, +} + +/// The state of the root run's stream. +struct RunState { + handle: Handle, + /// Taken while a CPU step runs on its task and put back with its result. + run: Option, + io: Arc, + /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. + inline: bool, +} + +/// Runs `root` over `scope` on a planning run of its own, reading through `io`, each CPU step +/// on a `spawn_cpu` task, and yields its batches and, unless `inline`, the chunks it spawns. +fn run_stream( + handle: Handle, + root: Box, + scope: WorkScope, + io: Arc, + inline: bool, +) -> BoxStream<'static, VortexResult> { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + let state = RunState { + handle, + run: Some(run), + io, + inline, + }; + stream::try_unfold(state, |mut state| async move { + loop { + let run = state + .run + .take() + .ok_or_else(|| vortex_err!("the scan's run is missing"))?; + let (run, step) = state + .handle + .spawn_cpu(move || { + let mut run = run; + let step = run.advance().map(|progress| match progress { + Progress::Batch(batch) => Step::Batch(batch.scope.rows.start, batch.array), + Progress::Spawn(scope, child) => Step::Spawn(scope, child), + Progress::RootDone(_) => Step::RootDone, + Progress::Idle => Step::Idle, + Progress::Waiting => Step::Waiting, + }); + (run, step) + }) + .await; + let run = state.run.insert(run); + match step? { + Step::Batch(start, batch) => { + return Ok(Some((Yielded::Batch(start, batch), state))); + } + Step::Spawn(scope, child) if state.inline => { + run.admit(child, scope, Arc::clone(&state.io)); + } + Step::Spawn(scope, child) => { + return Ok(Some((Yielded::Spawn(scope, child), state))); + } + Step::RootDone if run.live_roots() > 0 => {} + Step::RootDone | Step::Idle => return Ok(None), + Step::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }) + .boxed() +} + +/// Runs a spawned chunk's `root` over `scope` on a planning run of its own, on one task that +/// drives the run to its end and awaits only its reads, and yields its batches with the row each +/// starts at. A planner the chunk spawns in turn runs as another root of the same run. +fn chunk_stream( + handle: &Handle, + root: Box, + scope: WorkScope, + io: Arc, +) -> BoxStream<'static, VortexResult<(u64, ArrayRef)>> { + let (mut send, recv) = mpsc::channel(1); + let task = handle.spawn(async move { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + loop { + match run.advance()? { + Progress::Batch(batch) => { + if send + .send(Ok((batch.scope.rows.start, batch.array))) + .await + .is_err() + { + return Ok(()); + } + } + Progress::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&io)); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => return Ok(()), + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }); + let errors = stream::once(task) + .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); + stream::select(recv, errors).boxed() +} diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index 4abd2932509..67164f308b1 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/vortex-layout/src/plan/physical/plans/legacy_stats.rs @@ -30,7 +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; +use crate::plan::physical::plans::zoned::is_zone_pruning; const DATA: usize = 0; const ZONES: usize = 1; @@ -176,7 +176,7 @@ impl PlanReduceRule for LegacyStatsRule { } } -/// Converts the legacy stats beneath a zone proof in place, so the proof meets a [`Zoned`] plan +/// Converts the legacy stats beneath a zone pruning plan in place, so the plan meets a [`Zoned`] plan /// before the data path optimizes the stats away. #[derive(Debug)] pub(crate) struct ExpressionLegacyStatsRule; @@ -190,7 +190,7 @@ impl PlanParentReduceRule for ExpressionLegacyStatsRule { parent: &Plan, _child_idx: usize, ) -> VortexResult> { - if !is_zone_proof(parent)? { + if !is_zone_pruning(parent)? { return Ok(None); } let Some(converted) = LegacyStatsRule.reduce(child)? else { diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index 8b1729e953e..d4944588136 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -241,15 +241,6 @@ impl ZonedPlan { self.child_required(index) } - /// Whether an earlier execution already proved this pruning plan's expression over its zones, - /// so compiling it emits the kept result without reading the zone table. - pub(crate) fn is_proved(&self, session: &VortexSession) -> VortexResult { - match (self.pruned_zones(), self.pruning_expression()) { - (Some(pruned), Some(expression)) => Ok(pruned.get(expression, session)?.is_some()), - _ => Ok(false), - } - } - /// Returns whether this plan represents a zone-backed pruning proof. pub fn is_pruning(&self) -> bool { self.data().pruning.is_some() @@ -464,7 +455,7 @@ impl PlanParentReduceRule for ExpressionZonedRule { parent: &Plan, _child_idx: usize, ) -> VortexResult> { - if !is_zone_proof(parent)? { + if !is_zone_pruning(parent)? { return Ok(None); } Ok(child @@ -474,8 +465,8 @@ impl PlanParentReduceRule for ExpressionZonedRule { } /// Whether `eval` is a boolean expression over abstract statistics alone, which a zone table can -/// prove without reading the data it summarises. -pub(crate) fn is_zone_proof(eval: &Plan) -> VortexResult { +/// prune without reading the data it summarises. +pub(crate) fn is_zone_pruning(eval: &Plan) -> VortexResult { let mut contains_stat = false; let mut contains_root = false; eval.expression().clone().transform_down(|expression| { @@ -497,8 +488,8 @@ pub(crate) fn is_zone_proof(eval: &Plan) -> VortexResult { } /// Replaces a zoned data plan by its data: a scan reads the data through it and never its zone -/// table. Pruning proofs are derived from the filter over the lowered plan before this runs, so -/// an optimized plan carries a zone table only where a proof reads it. +/// table. Pruning plans are derived from the filter over the lowered plan before this runs, so +/// an optimized plan carries a zone table only where a pruning plan reads it. #[derive(Debug)] pub(crate) struct ZonedDataRule; diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 7eefcfcba2d..876bc55efb1 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -1589,7 +1589,7 @@ fn zoned_over_legacy_stats_read_as_data_keeps_only_the_data() -> VortexResult<() } /// Legacy stats over a zoned plan are dropped: read as data only the data remains, and a zone -/// proof is proved over the aggregate zones, never the legacy table. +/// pruning plan is evaluated over the aggregate zones, never the legacy table. #[test] fn legacy_stats_over_zoned_are_dropped() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 6808c707c48..382425d9bb0 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -59,19 +59,21 @@ pub(crate) fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec, rows: Range, chunks: IntoIter>, selection: Selection, - /// The rows zone pruning kept, over `rows`, or `None` when nothing was proved. + /// The rows zone pruning kept, over `rows`, or `None` when nothing was pruned. kept: Option, } impl ChunkPlanner { - pub(super) fn new( + /// The planner of the chunks of `rows`, planning only what `selection` and `kept`, the rows + /// zone pruning left over `rows`, leave something of. + pub fn new( scan: Arc, rows: Range, selection: Selection, diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 48718307577..87a24df4664 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -44,15 +44,15 @@ use crate::scan::filter::FilterExpr; use crate::segments::SegmentId; /// Plans and adaptive conjunct ordering shared by every chunk of a scan. -pub(in crate::scan) struct FilterPlans { +pub struct FilterPlans { expression: FilterExpr, predicates: Vec, - /// A metadata-only proof per conjunct, where the conjunct has one. + /// A metadata-only pruning plan per conjunct, where the conjunct has one. pruning: Vec>, } impl FilterPlans { - pub(in crate::scan) fn new( + pub(crate) fn new( expression: BoundExpression, root: &PlanRef, session: &VortexSession, @@ -71,7 +71,7 @@ impl FilterPlans { .iter() .map(|expr| { expr.falsify(session)? - .map(|proof| optimize(plan_row_idx_expression(proof, root.clone())?)) + .map(|pruning| optimize(plan_row_idx_expression(pruning, root.clone())?)) .transpose() }) .map(|plan| { @@ -89,13 +89,13 @@ impl FilterPlans { }) } - /// The pruning proofs, run once over a scan's rows before its chunks are planned. - pub(in crate::scan) fn proofs(&self) -> impl Iterator { + /// The pruning plans, run once over a scan's rows before its chunks are planned. + pub(crate) fn pruning(&self) -> impl Iterator { self.pruning.iter().flatten() } } -/// A pruning proof must not fetch data when its requested statistics are unavailable. +/// A pruning plan must not fetch data when its requested statistics are unavailable. fn metadata_only(plan: &PlanRef) -> VortexResult { if plan .as_opt::() diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index e371fdb539c..0185f61d6a9 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -3,13 +3,13 @@ //! The planners of a V2 scan, and the prepared plans they run. //! -//! A [`PreparedScan`] holds a scan's plans over one file and makes its planners. The planners -//! form a chain, each emitting the next through the planning protocol: a [`PrunePlanner`] -//! proves the filter's zone maps over the scan's rows and emits one [`ChunkPlanner`] for the rows -//! they keep; the chunk planner cuts those rows into chunks and emits a [`FilterPlanner`] or a -//! [`ProjectionPlanner`] per chunk; a filter planner runs its chunk's conjuncts and emits a -//! projection planner for the rows that survive; a projection planner emits one morsel per piece -//! of its rows. A stage with nothing to do is never made. +//! A [`PreparedScan`] holds a scan's plans over one file. An engine composes the planners over it +//! and drives them: a [`PrunePlanner`] evaluates the filter's pruning plans over the scan's rows and +//! emits the planner it was built to continue into, with the rows they keep; a [`ChunkPlanner`] +//! cuts those rows into chunks and spawns a [`FilterPlanner`] or a [`ProjectionPlanner`] per +//! chunk; a filter planner runs its chunk's conjuncts and emits a projection planner for the rows +//! that survive; a projection planner emits one morsel per piece of its rows. A planner only +//! says what it needs, compute or reads; where either runs is the engine's driver's decision. mod adaptor; mod chunk; @@ -22,6 +22,7 @@ use std::ops::Range; use std::sync::Arc; use vortex_buffer::Alignment; +use vortex_io::request::IoService; use vortex_io::request::IoTarget; use vortex_mask::Mask; use vortex_scan::planning::planner::Planner; @@ -29,7 +30,7 @@ use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_session::VortexSession; -pub(in crate::scan) use self::chunk::ChunkPlanner; +pub use self::chunk::ChunkPlanner; pub(crate) use self::chunk::FILTERED_ROWS_PER_CHUNK; pub(crate) use self::chunk::MAX_ROWS_PER_CHUNK; pub(crate) use self::chunk::ROW_BLOCK_SIZE; @@ -37,11 +38,16 @@ pub(crate) use self::chunk::ROW_BLOCK_SIZE; pub(crate) use self::chunk::row_chunks; pub(crate) use self::chunk::rows_per_chunk; use self::filter::FilterPlanner; -pub(super) use self::filter::FilterPlans; +pub use self::filter::FilterPlans; use self::projection::ProjectionPlanner; -use self::prune::PrunePlanner; +pub use self::prune::PruneNext; +pub use self::prune::PrunePlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; +use crate::scan::v2::ScanCache; +use crate::scan::v2::io::SegmentRanges; +use crate::scan::v2::io::SegmentScanIo; +use crate::segments::SegmentSource; /// Where a segment's bytes live in the file. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -67,14 +73,25 @@ impl SegmentLocation { /// A scan's plans over one file, built once from the file's footer and read by every planner of /// the scan. Nothing here touches a segment, and no planner plans again. -pub(crate) struct PreparedScan { +pub struct PreparedScan { /// The session used for decoding and expression evaluation. pub(super) session: VortexSession, + /// The rows of the root row domain the scan covers. + pub(super) rows: Range, + /// The rows of `rows` the scan selects. + pub(super) selection: Selection, + /// Fetches segment bytes by id. + pub(super) segments: Arc, + /// The segments' byte ranges, back to their ids. + pub(super) ranges: SegmentRanges, + /// Held for the scan's life, so scans of the file that start meanwhile share it. + #[expect(dead_code, reason = "kept alive, never read")] + pub(super) cache: ScanCache, /// The location of every segment, indexed by segment id. pub(super) locations: Arc<[SegmentLocation]>, /// The optimized projection plan, over the scan's root row domain. pub(super) projection: PlanRef, - /// The filter's conjuncts, their plans and their proofs. + /// The filter's conjuncts, their plans and their pruning plans. pub(super) filter: Option>, /// The global row index of the root row domain's first row. pub(super) row_offset: u64, @@ -86,40 +103,28 @@ pub(crate) struct PreparedScan { } impl PreparedScan { - /// The first planner a scan of `rows` needs. Every later planner is emitted by the one before - /// it: the zone proofs run first when the filter has any, then one planner per chunk of the - /// rows they keep. - pub(crate) fn planner( - self: &Arc, - rows: Range, - selection: Selection, - ) -> Box { - let proofs = self - .filter - .as_ref() - .is_some_and(|filter| filter.proofs().next().is_some()); - if proofs && !rows.is_empty() { - let scan = Arc::clone(self); - let chunks = rows.clone(); - Box::new(PrunePlanner::new( - Arc::clone(self), - rows, - Box::new(move |kept| Box::new(scan.chunk_planner(chunks, selection, Some(kept)))), - )) - } else { - Box::new(self.chunk_planner(rows, selection, None)) - } + /// The rows of the file's row domain the scan covers. + pub fn rows(&self) -> &Range { + &self.rows + } + + /// The rows of [`rows`](Self::rows) the scan selects. + pub fn selection(&self) -> &Selection { + &self.selection + } + + /// The filter's plans, when the scan has a filter. + pub fn filter(&self) -> Option<&FilterPlans> { + self.filter.as_deref() } - /// The planner cutting `rows` into chunks, planning only the chunks `selection` and `kept` - /// leave something of. - pub(super) fn chunk_planner( - self: &Arc, - rows: Range, - selection: Selection, - kept: Option, - ) -> ChunkPlanner { - ChunkPlanner::new(Arc::clone(self), rows, selection, kept) + /// The service a driver takes IO sessions from: every read a planner of the scan requests is + /// served from it by segment. + pub fn io(&self) -> Arc { + Arc::new(SegmentScanIo::new( + Arc::clone(&self.segments), + Arc::clone(&self.ranges), + )) } /// The planner of one chunk's selected rows: its conjuncts when the scan has a filter, else diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 1fd153f989c..325f5f897d7 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Proves a filter's zone maps once over a scan's rows, before any chunk is planned. +//! Evaluates a filter's pruning plans once over a scan's rows, before any chunk is planned. use std::ops::Range; use std::sync::Arc; @@ -20,7 +20,6 @@ use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; use vortex_scan::planning::planner::WorkScope; -use vortex_session::VortexSession; use super::PreparedScan; use super::adaptor::AdaptorStep; @@ -29,39 +28,27 @@ use super::filter::hint; use super::filter::piece_mask; use super::filter::segments_read; use crate::plan::physical::PlanRef; -use crate::plan::physical::Zoned; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; -/// Whether every zone proof in `plan` has already been proved, so compiling it reads no zones. -fn is_proved(plan: &PlanRef, session: &VortexSession) -> VortexResult { - if let Some(zoned) = plan.as_opt::() - && zoned.is_pruning() - { - return zoned.is_proved(session); - } - for child in plan.children().iter() { - if !is_proved(&child?, session)? { - return Ok(false); - } - } - Ok(true) -} - -/// The planner a prune planner continues into, given the rows the proofs kept. -pub(super) type Next = Box Box + Send>; +/// The planner a prune planner continues into, given the rows pruning kept, or `None` when +/// no pruning plan ran. +pub type PruneNext = Box) -> Box + Send>; -/// Runs the filter's zone-map proofs over the scan's rows and emits the planner it was built -/// to continue into, with the rows no proof excludes. The zone maps are the only reads, -/// started together. Rows every proof excludes end the scan here: nothing is emitted. -pub(super) struct PrunePlanner { +/// Evaluates the filter's pruning plans over the scan's rows and emits the planner it was +/// built to continue into, with the rows no plan pruned. The zone maps are the only reads, +/// started together. Rows every plan prunes end the scan here: nothing is emitted. A filter +/// without pruning plans emits the next planner at once. +pub struct PrunePlanner { scan: Arc, rows: Range, - /// What to emit once every proof has run, taken when it is. - next: Option, - /// The rows no proof has excluded yet, over `rows`. + /// What to emit once every pruning plan has run, taken when it is. + next: Option, + /// Whether any pruning plan ran, so the kept rows mean something. + pruned: bool, + /// The rows no pruning plan has excluded yet, over `rows`. mask: Mask, - proofs: IntoIter, + pruning: IntoIter, active: Option, pieces: Vec, announced: bool, @@ -71,19 +58,21 @@ pub(super) struct PrunePlanner { } impl PrunePlanner { - pub(super) fn new(scan: Arc, rows: Range, next: Next) -> Self { - let proofs: Vec<_> = scan + /// The planner of `scan`'s filter's pruning plans over `rows`, continuing into `next`. + pub fn new(scan: Arc, rows: Range, next: PruneNext) -> Self { + let pruning: Vec<_> = scan .filter .iter() - .flat_map(|filter| filter.proofs().cloned()) + .flat_map(|filter| filter.pruning().cloned()) .collect(); let len = usize::try_from(rows.end - rows.start).unwrap_or(usize::MAX); Self { scan, rows, next: Some(next), + pruned: false, mask: Mask::new_true(len), - proofs: proofs.into_iter(), + pruning: pruning.into_iter(), active: None, pieces: Vec::new(), announced: false, @@ -118,17 +107,11 @@ impl Planner for PrunePlanner { } if !self.announced { self.announced = true; - // A proof whose zones an earlier scan of the file proved reads nothing: a shared - // plan keeps its proved zones, so only unproved proofs need their zone maps. - let mut unproved = Vec::new(); - for proof in self.proofs.as_slice() { - if !is_proved(proof, &self.scan.session)? { - unproved.push(proof); - } - } + // A pruning plan whose zones an earlier scan of the file pruned compiles to its kept + // result and lists no segment, so only unpruned plans prefetch their zone maps. let segments = segments_read( &self.scan.session, - unproved.into_iter(), + self.pruning.as_slice().iter(), self.rows.clone(), &self.mask, )?; @@ -166,19 +149,20 @@ impl Planner for PrunePlanner { row_end = self.rows.end, input_rows = input, output_rows = self.mask.true_count(), - "v2 zone proof evaluated" + "v2 pruning plan evaluated" ); self.active = None; self.pieces.clear(); + self.pruned = true; } if self.mask.all_false() { self.done = true; return Ok(PlannerOutput::Done); } - if let Some(proof) = self.proofs.next() { + if let Some(plan) = self.pruning.next() { let scan = Scan::try_new( self.scan.session.clone(), - proof, + plan, vec![Split { rows: self.rows.clone(), mask: self.mask.clone(), @@ -189,19 +173,21 @@ impl Planner for PrunePlanner { self.active = Some(ScanAdaptor::new(scan, Arc::clone(&self.scan.locations))); return Ok(PlannerOutput::Continue); } - // Every proof has run: the planner this one continues into gets the kept rows. + // Every pruning plan has run: the planner this one continues into gets the kept rows. self.done = true; tracing::debug!( row_start = self.rows.start, row_end = self.rows.end, kept_rows = self.mask.true_count(), - "v2 zone proofs finished" + "v2 pruning finished" ); let next = self .next .take() .ok_or_else(|| vortex_err!("the prune planner has nothing to continue into"))?; - let kept = std::mem::replace(&mut self.mask, Mask::new_true(0)); + let kept = self + .pruned + .then(|| std::mem::replace(&mut self.mask, Mask::new_true(0))); Ok(PlannerOutput::Planner( WorkScope { file_ordinal: 0, diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 5a26e28ec97..c621bf19289 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -23,8 +23,8 @@ use crate::scan::execution::FilterPlans; /// over it, indexes every segment's byte range, and decodes the dictionary values it needs. A scan /// that finds an entry here reuses it, and one that misses builds the entry and keeps it for later /// scans. The lowered plan keeps the children it has initialized, each distinct projection keeps -/// its optimized plan, each distinct filter keeps its conjunct and zone-proof plans along with the -/// zones those proofs have already excluded, and dictionary values stay keyed by the cached plan +/// its optimized plan, each distinct filter keeps its conjunct and pruning plans along with the +/// zones those plans have already excluded, and dictionary values stay keyed by the cached plan /// that produced them. /// /// A cache describes one file's layout, segment locations, and segment source. Use a new cache @@ -40,8 +40,8 @@ struct Inner { /// them, so a projection must reuse its plan for its values to be found again. projections: Mutex>, ranges: Mutex>, - /// Plans of filters, keyed by the exact filter expression. A proof plan remembers the zones it - /// excluded, so sharing the plans proves a file's zones once for every scan of that filter. + /// Plans of filters, keyed by the exact filter expression. A pruning plan remembers the zones it + /// excluded, so sharing the plans prunes a file's zones once for every scan of that filter. filters: Mutex>>, values: ValueCache, } diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 2779739b281..0e690ebfde7 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -35,7 +35,7 @@ use crate::segments::SegmentId; use crate::segments::SegmentSource; /// Byte ranges of a file's segments, keyed by `(offset, length)`, back to their segment ids. -pub(super) type SegmentRanges = Arc>; +pub(crate) type SegmentRanges = Arc>; pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { let ranges = locations.iter().enumerate().map(|(id, location)| { @@ -46,13 +46,13 @@ pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { } /// Serves splits from any [`SegmentSource`], by segment id. -pub(super) struct SegmentScanIo { +pub(crate) struct SegmentScanIo { segments: Arc, ranges: SegmentRanges, } impl SegmentScanIo { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + pub(crate) fn new(segments: Arc, ranges: SegmentRanges) -> Self { Self { segments, ranges } } } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index d60f6b76289..4d2d40d7c38 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -3,62 +3,39 @@ //! The planner-driven scan executor. //! -//! A V2 scan lowers the file's layout once into a physical plan and runs every split of it -//! through the planning protocol: the [`Run`] driver advances -//! the split's planners and morsels, and every segment read is a protocol request served from the -//! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. +//! A V2 scan lowers the file's layout once into a physical plan and runs it through the planning +//! protocol: an engine composes planners over a [`PreparedScan`], a [`PrunePlanner`] into a +//! [`ChunkPlanner`], admits the first to a [`Run`](vortex_scan::planning::driver::Run) of its own +//! driver, and every segment read is a protocol request served from the file's +//! [`SegmentSource`]. The layout reader is never asked to evaluate anything, and this crate never +//! decides where compute runs or how IO is awaited: that is the engine's driver. //! -//! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] -//! directly and supports projection, filters, row ranges, and selections. Unsupported layouts return an -//! error; limits are not part of this builder's API. Work is partitioned into fixed +//! Import this module's [`ScanBuilder`] to prepare a V2 scan. It accepts a [`ScanFile`] directly +//! and supports projection, filters, row ranges, and selections. Unsupported layouts return an +//! error when the plan is prepared; limits are not part of this builder's API. Work is cut into //! chunks of up to 65,536 rows, or 1,048,576 rows for filtered scans, independently of the stored -//! layout's chunk boundaries. -//! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. +//! layout's chunk boundaries. Full-file projections use smaller 8,192-row multiples when +//! serialized bytes per row are high. //! -//! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and proved +//! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and pruned //! zones, segment index, and dictionary values with every other scan of that file, like a cached //! layout reader does for V1 scans. mod cache; -mod io; +pub(crate) mod io; #[cfg(test)] mod tests; -use std::future; use std::ops::Range; use std::sync::Arc; -use std::task::Poll; -use arrow_array::RecordBatch; -use arrow_schema::ArrowError; -use arrow_schema::Field; -use arrow_schema::SchemaRef; -use futures::SinkExt; -use futures::Stream; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::channel::mpsc; -use futures::stream; -use futures::stream::BoxStream; -use vortex_array::ArrayRef; -use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_io::request::IoService; -use vortex_io::request::IoSource; -use vortex_io::runtime::Handle; -use vortex_io::session::RuntimeSessionExt; -use vortex_scan::planning::driver::Progress; -use vortex_scan::planning::driver::Run; -use vortex_scan::planning::planner::Planner; -use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; pub use self::cache::ScanCache; pub use self::cache::WeakScanCache; @@ -66,15 +43,16 @@ use crate::LayoutRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; -use crate::scan::arrow::to_record_batch; +pub use crate::scan::execution::ChunkPlanner; use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; -use crate::scan::execution::FilterPlans; +pub use crate::scan::execution::FilterPlans; use crate::scan::execution::MAX_ROWS_PER_CHUNK; -use crate::scan::execution::PreparedScan; +pub use crate::scan::execution::PreparedScan; +pub use crate::scan::execution::PruneNext; +pub use crate::scan::execution::PrunePlanner; use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; use crate::scan::execution::rows_per_chunk; -use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::segments::SegmentSource; @@ -107,24 +85,21 @@ pub struct ScanFile { pub cache: Option, } -/// Builds a projection scan on the planner-driven V2 executor. +/// Prepares a scan on the planner-driven V2 executor. /// -/// This builder reads the file's layout and segments directly. Unsupported layouts return an -/// error when the stream is polled. Row limits are not supported. -pub struct ScanBuilder { +/// This builder reads the file's layout and segments directly and plans the scan without IO. +/// Driving the plan is the engine's: see [`PreparedScan`]. Row limits are not supported. +pub struct ScanBuilder { session: VortexSession, file: ScanFile, projection: BoundExpression, filter: Option, - ordered: bool, row_range: Option>, selection: Selection, - concurrency: usize, - map_fn: Arc VortexResult + Send + Sync>, row_offset: u64, } -impl ScanBuilder { +impl ScanBuilder { /// Creates a V2 scan over the file's layout and segment source. pub fn new(session: VortexSession, file: ScanFile) -> Self { Self { @@ -132,34 +107,13 @@ impl ScanBuilder { filter: None, session, file, - ordered: true, row_range: None, selection: Selection::default(), - concurrency: 4, - map_fn: Arc::new(Ok), row_offset: 0, } } - /// Converts each output batch to Arrow on the split's task. - pub fn into_record_batch_stream( - self, - schema: SchemaRef, - ) -> VortexResult> + Send + 'static> { - let struct_field = Field::new_struct("", schema.fields().clone(), false); - let session = self.session.clone(); - Ok(self - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - to_record_batch(chunk, &struct_field, &mut ctx) - }) - .into_stream()? - .map_err(|error| ArrowError::ExternalError(Box::new(error)))) - } -} - -impl ScanBuilder { - /// The output dtype after projection, before any output mapping. + /// The output dtype after projection. pub fn dtype(&self) -> &DType { self.projection.dtype() } @@ -176,12 +130,6 @@ impl ScanBuilder { self } - /// Configures whether batches are yielded in file order. - pub fn with_ordered(mut self, ordered: bool) -> Self { - self.ordered = ordered; - self - } - /// Restricts scanning to a contiguous file-local row range. pub fn with_row_range(mut self, row_range: Range) -> Self { self.row_range = Some(row_range); @@ -205,82 +153,25 @@ impl ScanBuilder { self } - /// Sets the number of chunks run at once per worker thread. Must be greater than zero. - pub fn with_concurrency(mut self, concurrency: usize) -> Self { - assert!(concurrency > 0); - self.concurrency = concurrency; - self - } - - /// Maps each output batch on its split's task. - pub fn map( - self, - map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, - ) -> ScanBuilder { - let old_map_fn = self.map_fn; - ScanBuilder { - session: self.session, - file: self.file, - projection: self.projection, - filter: self.filter, - ordered: self.ordered, - row_range: self.row_range, - selection: self.selection, - concurrency: self.concurrency, - row_offset: self.row_offset, - map_fn: Arc::new(move |array| old_map_fn(array).and_then(&map_fn)), - } - } - - /// Runs the scan on the session's runtime, preparing it on the first poll. - /// Preparation errors are returned as stream items. - pub fn into_stream(self) -> VortexResult>> { - let handle = self.session.handle(); - Ok(stream::once(async move { - handle - .spawn_cpu(move || ScanRun::try_new(self)?.execute_stream()) - .await - }) - .try_flatten() - .boxed()) - } -} - -/// A scan prepared over its file, ready to run. -struct ScanRun { - session: VortexSession, - scan: Arc, - /// Serves the scan's reads. - io: Arc, - /// The rows of the root row domain the scan covers. - rows: Range, - selection: Selection, - map_fn: Arc VortexResult + Send + Sync>, - ordered: bool, - concurrency: usize, - /// Held until the scan ends, so scans of the file that start meanwhile share it. - cache: ScanCache, -} - -impl ScanRun { /// Lowers the file's layout and plans the filter and the projection, from the file's cache - /// where an earlier scan planned them already. Touches no segment. - fn try_new(builder: ScanBuilder) -> VortexResult { - let file = builder.file; + /// where an earlier scan planned them already. Touches no segment. Lowering and optimizing a + /// wide file is real CPU, so a driver runs this where it runs compute. + pub fn prepare(self) -> VortexResult { + let file = self.file; let cache = file.cache.unwrap_or_default(); let root = match cache.root() { Some(root) => root, None => cache.insert_root(lower(&file.layout)?), }; let row_count = root.row_count(); - let filter = builder + let filter = self .filter .map(|filter| -> VortexResult<_> { match cache.filter_plans(&filter) { Some(plans) => Ok(plans), None => { let plans = - Arc::new(FilterPlans::new(filter.clone(), &root, &builder.session)?); + Arc::new(FilterPlans::new(filter.clone(), &root, &self.session)?); Ok(cache.insert_filter_plans(filter, plans)) } } @@ -288,12 +179,12 @@ impl ScanRun { .transpose()?; // Optimizing drops the zone tables a projection never reads and keeps the chunks of a // chunked column lazy, so an identity projection is optimized like any other. - let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype(); - let projection = match cache.projection(&builder.projection) { + let identity = self.projection.is_root() && self.projection.dtype() == root.dtype(); + let projection = match cache.projection(&self.projection) { Some(plan) => plan, None => { - let plan = optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; - cache.insert_projection(builder.projection, plan) + let plan = optimize(plan_row_idx_expression(self.projection.clone(), root)?)?; + cache.insert_projection(self.projection, plan) } }; let chunk_rows = if filter.is_some() { @@ -307,215 +198,20 @@ impl ScanRun { Some(ranges) => ranges, None => cache.insert_ranges(segment_ranges(&file.locations)), }; - let io = Arc::new(SegmentScanIo::new(file.segments, ranges)); - Ok(Self { - session: builder.session.clone(), - scan: Arc::new(PreparedScan { - session: builder.session, - locations: file.locations, - projection, - filter, - row_offset: builder.row_offset, - values: cache.values(), - chunk_rows, - }), - io, - rows: builder.row_range.unwrap_or(0..row_count), - selection: builder.selection, - map_fn: builder.map_fn, - ordered: builder.ordered, - concurrency: builder.concurrency, + let values = cache.values(); + Ok(PreparedScan { + session: self.session, + rows: self.row_range.unwrap_or(0..row_count), + selection: self.selection, + segments: file.segments, + ranges, cache, + locations: file.locations, + projection, + filter, + row_offset: self.row_offset, + values, + chunk_rows, }) } - - /// Runs the scan on planning runs the returned stream owns. - /// - /// The scan's first planner runs on one run the stream drives as it is polled, each of its - /// CPU steps a [`Handle::spawn_cpu`] task. Every chunk it spawns runs on a run and a session - /// of its own, on one task of the runtime that drives the run to its end and awaits only its - /// reads, up to `concurrency` chunks per worker thread at once; a chunk is started only when - /// the consumer has room for its batches. Every read is a `spawn_io` task of the segment - /// source. The engine thread polling the stream does no scan work itself. Batches of an - /// unordered scan are yielded as morsels finish; an ordered scan yields them in row order - /// once every morsel has. - fn execute_stream(self) -> VortexResult>> { - let handle = self.session.handle(); - let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); - let root = self.scan.planner(self.rows.clone(), self.selection); - let scope = WorkScope { - file_ordinal: 0, - rows: self.rows.clone(), - }; - let io = self.io; - let map_fn = self.map_fn; - let (chunk_handle, chunk_io, chunk_map) = - (handle.clone(), Arc::clone(&io), Arc::clone(&map_fn)); - let inline = in_flight <= 1; - let batches = run_stream(handle, root, scope, io.session(), map_fn, inline) - .map_ok(move |yielded| match yielded { - Yielded::Batch(start, batch) => { - stream::once(future::ready(Ok((start, batch)))).boxed() - } - Yielded::Spawn(scope, chunk) => chunk_stream( - &chunk_handle, - chunk, - scope, - chunk_io.session(), - Arc::clone(&chunk_map), - ), - }) - .try_flatten_unordered(in_flight); - let cache = self.cache; - let batches = batches.chain(stream::poll_fn(move |_| { - let _ = &cache; - Poll::Ready(None) - })); - if !self.ordered { - return Ok(batches.map_ok(|(_, batch)| batch).boxed()); - } - Ok(stream::once(async move { - let mut held: Vec<(u64, A)> = batches.try_collect().await?; - // Morsels finish in whatever order their reads arrive. - held.sort_by_key(|(start, _)| *start); - VortexResult::Ok(stream::iter(held.into_iter().map(|(_, batch)| Ok(batch)))) - }) - .try_flatten() - .boxed()) - } -} - -/// What the root run's stream yields: a mapped batch with the row it starts at, or a chunk a -/// planner of the run spawned. -enum Yielded { - Batch(u64, A), - Spawn(WorkScope, Box), -} - -/// One CPU step of a run, taken on a `spawn_cpu` task with the run. -enum Step { - Batch(u64, A), - Spawn(WorkScope, Box), - RootDone, - Idle, - Waiting, -} - -/// The state of the root run's stream. -struct RunState { - handle: Handle, - /// Taken while a CPU step runs on its task and put back with its result. - run: Option, - io: Arc, - map_fn: Arc VortexResult + Send + Sync>, - /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. - inline: bool, -} - -/// Runs `root` over `scope` on a planning run of its own, reading through `io`, each CPU step -/// on a `spawn_cpu` task, and yields its mapped batches and, unless `inline`, the chunks it -/// spawns. Inline, a spawned chunk runs on this run beside the others: the run's steps already -/// interleave every chunk's work, and a run per chunk would only add a hand-off per chunk for -/// parallelism the engine provides. -fn run_stream( - handle: Handle, - root: Box, - scope: WorkScope, - io: Arc, - map_fn: Arc VortexResult + Send + Sync>, - inline: bool, -) -> BoxStream<'static, VortexResult>> { - let mut run = Run::new(); - run.admit(root, scope, Arc::clone(&io)); - let state = RunState { - handle, - run: Some(run), - io, - map_fn, - inline, - }; - stream::try_unfold(state, |mut state| async move { - loop { - let run = state - .run - .take() - .ok_or_else(|| vortex_err!("the scan's run is missing"))?; - let map_fn = Arc::clone(&state.map_fn); - let (run, step) = state - .handle - .spawn_cpu(move || { - let mut run = run; - let step = run.advance().and_then(|progress| { - Ok(match progress { - Progress::Batch(batch) => { - Step::Batch(batch.scope.rows.start, map_fn(batch.array)?) - } - Progress::Spawn(scope, child) => Step::Spawn(scope, child), - Progress::RootDone(_) => Step::RootDone, - Progress::Idle => Step::Idle, - Progress::Waiting => Step::Waiting, - }) - }); - (run, step) - }) - .await; - let run = state.run.insert(run); - match step? { - Step::Batch(start, batch) => { - return Ok(Some((Yielded::Batch(start, batch), state))); - } - Step::Spawn(scope, child) if state.inline => { - run.admit(child, scope, Arc::clone(&state.io)); - } - Step::Spawn(scope, child) => { - return Ok(Some((Yielded::Spawn(scope, child), state))); - } - Step::RootDone if run.live_roots() > 0 => {} - Step::RootDone | Step::Idle => return Ok(None), - Step::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - }) - .boxed() -} - -/// Runs a spawned chunk's `root` over `scope` on a planning run of its own, on one task that -/// drives the run to its end and awaits only its reads, and yields its mapped batches with the -/// row each starts at. A planner the chunk spawns in turn runs as another root of the same run. -/// -/// One task per chunk, not one per step: a chunk's run waits for reads many times, and a task -/// hop at each wait is latency on the chunk's critical path that no other chunk's work fills. -fn chunk_stream( - handle: &Handle, - root: Box, - scope: WorkScope, - io: Arc, - map_fn: Arc VortexResult + Send + Sync>, -) -> BoxStream<'static, VortexResult<(u64, A)>> { - let (mut send, recv) = mpsc::channel(1); - let task = handle.spawn(async move { - let mut run = Run::new(); - run.admit(root, scope, Arc::clone(&io)); - loop { - match run.advance()? { - Progress::Batch(batch) => { - let start = batch.scope.rows.start; - let batch = map_fn(batch.array); - let failed = batch.is_err(); - if send.send(batch.map(|batch| (start, batch))).await.is_err() || failed { - return Ok(()); - } - } - Progress::Spawn(scope, child) => { - run.admit(child, scope, Arc::clone(&io)); - } - Progress::RootDone(_) if run.live_roots() > 0 => {} - Progress::RootDone(_) | Progress::Idle => return Ok(()), - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - }); - let errors = stream::once(task) - .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); - stream::select(recv, errors).boxed() } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index eabf863eb47..c6dccf7f12e 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -6,14 +6,11 @@ use std::ops::Range; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use std::time::Duration; use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; -use futures::FutureExt; +use arrow_schema::Field; use futures::TryStreamExt; -use futures::channel::oneshot; -use futures::future::Shared; use futures::stream; use rstest::rstest; use vortex_array::ArrayContext; @@ -44,8 +41,8 @@ use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; use vortex_scan::planning::driver::Batch; use vortex_scan::planning::driver::Progress; use vortex_scan::planning::driver::Run; @@ -55,12 +52,14 @@ use vortex_scan::planning::planner::WorkScope; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; +use super::ChunkPlanner; use super::MAX_ROWS_PER_CHUNK; +use super::PreparedScan; +use super::PrunePlanner; use super::ROW_BLOCK_SIZE; use super::ScanBuilder; use super::ScanCache; use super::ScanFile; -use super::ScanRun; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; @@ -71,6 +70,7 @@ use crate::layouts::row_idx::row_idx; use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; +use crate::scan::arrow::to_record_batch; use crate::scan::execution::SegmentLocation; use crate::scan::execution::row_chunks; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; @@ -188,34 +188,50 @@ fn builder( session: &VortexSession, segments: &Arc, layout: &LayoutRef, -) -> VortexResult> { +) -> VortexResult { Ok(ScanBuilder::new( session.clone(), scan_file(segments, layout)?, )) } -/// Runs `rows` of a prepared scan on a planning run of its own and returns its batches, in row -/// order, each with the rows it stands for. -async fn run_batches( - scan: &ScanRun, - rows: Range, -) -> VortexResult> { +/// Composes a scan's planners the way an engine does: the pruning plans first when there is a filter, +/// continuing into the chunks of the rows they keep. +fn compose(scan: &Arc, rows: Range) -> Box { + let chunks = { + let scan = Arc::clone(scan); + let rows = rows.clone(); + move |kept: Option| -> Box { + let selection = scan.selection().clone(); + Box::new(ChunkPlanner::new(scan, rows, selection, kept)) + } + }; + if scan.filter().is_some() { + Box::new(PrunePlanner::new(Arc::clone(scan), rows, Box::new(chunks))) + } else { + chunks(None) + } +} + +/// Runs `rows` of a prepared scan on one planning run, every spawned chunk admitted to it, and +/// returns its batches in row order, each with the rows it stands for. +async fn run_batches(scan: &Arc, rows: Range) -> VortexResult> { + let io = scan.io(); let mut run = Run::new(); run.admit( - scan.scan.planner(rows.clone(), scan.selection.clone()), + compose(scan, rows.clone()), WorkScope { file_ordinal: 0, rows, }, - scan.io.session(), + io.session(), ); let mut batches = Vec::new(); loop { match run.advance()? { Progress::Batch(batch) => batches.push(batch), Progress::Spawn(scope, chunk) => { - run.admit(chunk, scope, scan.io.session()); + run.admit(chunk, scope, io.session()); } Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => break, @@ -226,11 +242,8 @@ async fn run_batches( Ok(batches) } -/// Runs `rows` of a prepared scan on a planning run of its own and returns its arrays. -async fn run_rows( - scan: &ScanRun, - rows: Range, -) -> VortexResult> { +/// Runs `rows` of a prepared scan and returns its arrays. +async fn run_rows(scan: &Arc, rows: Range) -> VortexResult> { Ok(run_batches(scan, rows) .await? .into_iter() @@ -238,6 +251,13 @@ async fn run_rows( .collect()) } +/// Prepares and runs a whole scan, returning its arrays. +async fn run_scan(builder: ScanBuilder) -> VortexResult> { + let scan = Arc::new(builder.prepare()?); + let rows = scan.rows().clone(); + run_rows(&scan, rows).await +} + async fn collect( dtype: DType, stream: impl futures::Stream>, @@ -246,10 +266,15 @@ async fn collect( Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) } +/// Prepares and runs a whole scan, returning its rows as one chunked array. +async fn collect_scan(dtype: DType, builder: ScanBuilder) -> VortexResult { + Ok(ChunkedArray::try_new(run_scan(builder).await?, dtype)?.into_array()) +} + /// Runs `builder` on both executors and checks the planner-driven one returns the same rows. async fn assert_same_as_v1( session: &VortexSession, - make: impl Fn() -> VortexResult>, + make: impl Fn() -> VortexResult, ) -> VortexResult<()> { let builder = make()?; let dtype = builder.dtype().clone(); @@ -263,13 +288,12 @@ async fn assert_same_as_v1( .with_projection(builder.projection.clone()) .with_some_filter(builder.filter.clone()) .with_selection(builder.selection.clone()) - .with_row_offset(builder.row_offset) - .with_ordered(builder.ordered); + .with_row_offset(builder.row_offset); if let Some(rows) = builder.row_range.clone() { v1 = v1.with_row_range(rows); } let expected = collect(dtype.clone(), v1.into_stream()?).await?; - let actual = collect(dtype, builder.into_stream()?).await?; + let actual = collect_scan(dtype, builder).await?; assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); Ok(()) } @@ -420,8 +444,14 @@ async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { let scan = builder(&session, &segments, &layout)? .with_projection(projection) .with_row_range(990..1010); - let schema = Arc::new(session.arrow().to_arrow_schema(scan.dtype())?); - let batches: Vec<_> = scan.into_record_batch_stream(schema)?.try_collect().await?; + let schema = session.arrow().to_arrow_schema(scan.dtype())?; + let field = Field::new_struct("", schema.fields().clone(), false); + let mut ctx = session.create_execution_ctx(); + let batches = run_scan(scan) + .await? + .into_iter() + .map(|chunk| to_record_batch(chunk, &field, &mut ctx)) + .collect::>>()?; let values: Vec<_> = batches .iter() .flat_map(|batch| { @@ -518,8 +548,8 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; - let chunks = row_chunks(scan.rows.clone(), scan.scan.chunk_rows); + let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); + let chunks = row_chunks(scan.rows().clone(), scan.chunk_rows); assert_eq!( chunks.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) @@ -557,10 +587,7 @@ async fn list_layouts_return_an_error() -> VortexResult<()> { let session = new_session().with_tokio(); let strategy = TableStrategy::new(chunked_flat(), chunked_flat()).with_list_layout(); let (segments, layout) = write(&session, &strategy).await?; - let builder = builder(&session, &segments, &layout)?; - let error = builder - .into_stream()? - .try_next() + let error = run_scan(builder(&session, &segments, &layout)?) .await .expect_err("V2 must reject an unsupported layout"); assert!(error.to_string().contains("vortex.plan.list_pack")); @@ -580,112 +607,25 @@ async fn projection_skips_unsupported_unselected_layouts() -> VortexResult<()> { } #[tokio::test] -async fn unordered_returns_every_row() -> VortexResult<()> { +async fn every_row_is_returned() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = ScanRun::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; - let chunks: Vec = scan.execute_stream()?.try_collect().await?; + let chunks = run_scan(builder(&session, &segments, &layout)?).await?; let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); assert_eq!(rows, CHUNKS * CHUNK_ROWS); Ok(()) } -struct GatedSegments { - inner: Arc, - blocked: SegmentId, - release: Shared>, -} - -impl SegmentSource for GatedSegments { - fn request(&self, id: SegmentId) -> SegmentFuture { - let bytes = self.inner.request(id); - if id != self.blocked { - return bytes; - } - let release = self.release.clone(); - async move { - let _ = release.await; - bytes.await - } - .boxed() - } -} - -#[rstest] -#[timeout(Duration::from_secs(5))] -#[tokio::test] -async fn unordered_yields_before_the_split_finishes() -> VortexResult<()> { - let session = new_session().with_tokio(); - let segments = Arc::new(TestSegments::default()); - let (mut sequence, eof) = SequenceId::root().split(); - let arrays = (0..2) - .map(|index| { - let start = i32::try_from(index * MAX_ROWS_PER_CHUNK)?; - Ok(( - sequence.advance(), - PrimitiveArray::from_iter(start..start + i32::try_from(MAX_ROWS_PER_CHUNK)?) - .into_array(), - )) - }) - .collect::>>()?; - let dtype = arrays[0].1.dtype().clone(); - let layout = chunked_flat() - .write_stream( - ArrayContext::empty().into(), - Arc::::clone(&segments), - SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))) - .sendable(), - eof, - &session, - ) - .await?; - let blocked = *layout - .slot(1)? - .ok_or_else(|| vortex_err!("second chunk is missing"))? - .segment_ids() - .first() - .ok_or_else(|| vortex_err!("second chunk has no segment"))?; - let (release, gate) = oneshot::channel(); - let segments: Arc = Arc::new(GatedSegments { - inner: segments, - blocked, - release: gate.shared(), - }); - let mut scan = ScanRun::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; - // Put both projection morsels in one chunk, as a dense filtered scan would. - Arc::get_mut(&mut scan.scan) - .ok_or_else(|| vortex_err!("the prepared scan is not shared yet"))? - .chunk_rows = usize::try_from(layout.row_count())?; - let mut batches = scan.execute_stream()?; - let first = batches.try_next().await; - let _ = release.send(()); - let first = first?.ok_or_else(|| vortex_err!("first batch is missing"))?; - assert_arrays_eq!( - first, - PrimitiveArray::from_iter(0i32..i32::try_from(MAX_ROWS_PER_CHUNK)?), - &mut session.create_execution_ctx() - ); - let rest: Vec<_> = batches.try_collect().await?; - assert_eq!( - rest.iter().map(|array| array.len()).sum::(), - MAX_ROWS_PER_CHUNK - ); - Ok(()) -} - #[tokio::test] async fn primitive_column() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; let dtype = layout.dtype().clone(); - let scan = ScanRun::try_new( + let actual = collect_scan( + DType::Primitive(vortex_array::dtype::PType::I32, false.into()), builder(&session, &segments, &layout)? .with_projection(get_item("ints", root()).bind(&dtype)?) .with_row_range(990..1010), - )?; - let actual = collect( - DType::Primitive(vortex_array::dtype::PType::I32, false.into()), - scan.execute_stream()?, ) .await?; assert_arrays_eq!( @@ -704,12 +644,12 @@ async fn work_chunks_ignore_layout_boundaries( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, strategy.get().as_ref()).await?; - let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); assert_eq!( - row_chunks(scan.rows.clone(), scan.scan.chunk_rows), + row_chunks(scan.rows().clone(), scan.chunk_rows), vec![0..(CHUNKS * CHUNK_ROWS) as u64] ); - let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; + let arrays = run_rows(&scan, scan.rows().clone()).await?; assert_eq!(arrays.len(), 1); assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) @@ -769,7 +709,7 @@ async fn scan_cache_reuses_values_between_scans( let before = counted.reads.load(Ordering::Relaxed); let scan = ScanBuilder::new(session.clone(), file.clone()) .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![12u64, 15]))?); - let actual = collect(layout.dtype().clone(), scan.into_stream()?).await?; + let actual = collect_scan(layout.dtype().clone(), scan).await?; assert_arrays_eq!( actual, PrimitiveArray::from_iter([12i32, 15]), @@ -786,13 +726,17 @@ async fn scan_cache_reuses_values_between_scans( async fn selected_split_scope_is_bounded_before_projection() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = ScanRun::try_new( + let scan = Arc::new( builder(&session, &segments, &layout)? - .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 2001]))?), - )?; - let mut chunks = scan - .scan - .chunk_planner(scan.rows.clone(), scan.selection.clone(), None); + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 2001]))?) + .prepare()?, + ); + let mut chunks = ChunkPlanner::new( + Arc::clone(&scan), + scan.rows().clone(), + scan.selection().clone(), + None, + ); let PlannerOutput::Spawn(scope, _) = chunks.compute()? else { vortex_bail!("the selected rows were not planned"); }; @@ -825,7 +769,7 @@ async fn repeated_projections_reuse_their_values() -> VortexResult<()> { let scan = ScanBuilder::new(session.clone(), file.clone()) .with_projection(projection.clone()) .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![12u64, 18]))?); - let actual = collect(projection.dtype().clone(), scan.into_stream()?).await?; + let actual = collect_scan(projection.dtype().clone(), scan).await?; assert_arrays_eq!( actual, vortex_array::arrays::BoolArray::from_iter([false, true]), @@ -859,7 +803,7 @@ fn partition_rows_tile_the_file_on_row_blocks() { ); } -/// Zone proofs run once before any chunk is planned: rows every proof excludes are never +/// Pruning runs once before any chunk is planned: rows every pruning plan excludes are never /// given a chunk, and a scan whose zones all fail plans nothing. #[rstest] #[case::tail(2_999, Some(3_000..4_000))] @@ -871,12 +815,14 @@ async fn pruned_rows_get_no_chunk( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, zoned_table().as_ref()).await?; - let scan = ScanRun::try_new( - builder(&session, &segments, &layout)?.with_some_filter(Some( - gt(get_item("ints", root()), lit(above)).bind(layout.dtype())?, - )), - )?; - let batches = run_batches(&scan, scan.rows.clone()).await?; + let scan = Arc::new( + builder(&session, &segments, &layout)? + .with_some_filter(Some( + gt(get_item("ints", root()), lit(above)).bind(layout.dtype())?, + )) + .prepare()?, + ); + let batches = run_batches(&scan, scan.rows().clone()).await?; assert_eq!( batches .iter() @@ -887,10 +833,10 @@ async fn pruned_rows_get_no_chunk( Ok(()) } -/// Scans of one file that share a cache share their filter plans, so the zone proofs run once: +/// Scans of one file that share a cache share their filter plans, so the pruning plans run once: /// the second scan with the same filter re-reads the data it needs but never the zone maps. #[tokio::test] -async fn shared_cache_proves_zones_once() -> VortexResult<()> { +async fn shared_cache_prunes_zones_once() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, zoned_table().as_ref()).await?; let counted = Arc::new(CountingSegments { @@ -911,15 +857,15 @@ async fn shared_cache_proves_zones_once() -> VortexResult<()> { let scan = ScanBuilder::new(session.clone(), file.clone()) .with_projection(projection.clone()) .with_some_filter(Some(filter.clone())); - let actual = collect(projection.dtype().clone(), scan.into_stream()?).await?; + let actual = collect_scan(projection.dtype().clone(), scan).await?; assert_eq!(actual.len(), 1_000); reads.push(counted.reads.load(Ordering::Relaxed) - before); } let plans = cache .filter_plans(&filter) .expect("the first scan kept its filter plans"); - assert_eq!(plans.proofs().count(), 1); - // Each scan reads the one chunk the proof keeps; only the first reads the zone map. + assert_eq!(plans.pruning().count(), 1); + // Each scan reads the one chunk pruning keeps; only the first reads the zone map. assert_eq!(reads, vec![2, 1]); Ok(()) }