diff --git a/AGENTS.md b/AGENTS.md index c92cb4d2ff9..f38f7998935 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -223,7 +223,7 @@ These CI checks are the ones most often missed when adding files rather than edi default to Snappy. - parquet-rs: `Compression::ZSTD(ZstdLevel::try_new(3)?)` - DuckDB: `(FORMAT parquet, COMPRESSION zstd, COMPRESSION_LEVEL 3)` - - tpchgen-cli: `--parquet-compression='ZSTD(3)'` + - tpchgen-cli: `tpchgen-cli parquet --compression='ZSTD(3)'` - Benchmarks that deliberately compare Parquet codecs, such as the GPU compress benchmark's `--codec`, are exempt. diff --git a/Cargo.lock b/Cargo.lock index 22aace294d8..d64d2915fe5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9055,6 +9055,7 @@ version = "0.1.0" dependencies = [ "anyhow", "clap", + "codspeed-divan-compat", "indicatif", "lance-bench", "mimalloc", @@ -9063,6 +9064,7 @@ dependencies = [ "tabled", "tokio", "url", + "vortex", "vortex-bench", ] diff --git a/benchmarks/random-access-bench/Cargo.toml b/benchmarks/random-access-bench/Cargo.toml index cf359fc3993..5a502600232 100644 --- a/benchmarks/random-access-bench/Cargo.toml +++ b/benchmarks/random-access-bench/Cargo.toml @@ -26,6 +26,10 @@ tokio = { workspace = true, features = ["full"] } url = { workspace = true } vortex-bench = { workspace = true } +[dev-dependencies] +divan = { workspace = true } +vortex = { workspace = true } + [features] lance = ["dep:lance-bench"] @@ -33,5 +37,9 @@ lance = ["dep:lance-bench"] name = "random-access-bench" test = false +[[bench]] +name = "v2_take" +harness = false + [lints] workspace = true diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index 73f52f5b219..65d573d6100 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -22,12 +22,25 @@ untimed warm-up, then reuses the open file handle. Reopen mode includes file ope work in each timed iteration. CI drives the full matrix via [`scripts/random-access-split.py`](../../scripts/random-access-split.py). +Vortex measurements use the planner-driven V2 scan builder. + ## Running locally ```bash cargo run -p random-access-bench --profile release_debug --features lance ``` +## Micro benchmarks + +The `v2_take` divan benchmark runs each Vortex query with the default and V2 scans over the +files in the page cache, plus the per-scan cost of a V2 scan that selects no rows. Write the +files first: + +```bash +cargo run -p random-access-bench --profile release_debug -- --prepare-data --formats vortex +cargo bench -p random-access-bench --bench v2_take +``` + ## Running against S3 The same benchmark can read its data from an object store instead of local disk. The remote diff --git a/benchmarks/random-access-bench/benches/v2_take.rs b/benchmarks/random-access-bench/benches/v2_take.rs new file mode 100644 index 00000000000..e1c00b1995b --- /dev/null +++ b/benchmarks/random-access-bench/benches/v2_take.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Each random-access query with the default and V2 scans, over files in the page cache. +//! +//! Run `random-access-bench --prepare-data --formats vortex` first to write the files. + +#![expect(clippy::unwrap_used)] + +use std::hint::black_box; +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use random_access_bench::AccessPattern; +use random_access_bench::generate_indices; +use tokio::runtime::Runtime; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::buffer::Buffer; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; +use vortex::layout::scan::v2::ScanBuilder; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_bench::Format; +use vortex_bench::SESSION; +use vortex_bench::datasets::feature_vectors::FeatureVectorsData; +use vortex_bench::datasets::nested_structs::NestedStructsData; +use vortex_bench::datasets::taxi_data::TaxiData; +use vortex_bench::random_access::BenchDataset; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static RUNTIME: LazyLock = LazyLock::new(|| Runtime::new().unwrap()); + +/// The queries V2 can scan: nested lists use list layouts, which V2 does not support. +const QUERIES: &[&str] = &[ + "taxi/correlated", + "taxi/uniform", + "feature-vectors/correlated", + "feature-vectors/uniform", + "nested-structs/correlated", + "nested-structs/uniform", +]; + +fn main() { + divan::main(); +} + +struct Query { + file: VortexFile, + indices: Vec, +} + +impl Query { + fn open(query: &str) -> Self { + let (dataset, pattern) = query.split_once('/').unwrap(); + let dataset: Box = match dataset { + "taxi" => Box::new(TaxiData), + "feature-vectors" => Box::new(FeatureVectorsData), + _ => Box::new(NestedStructsData), + }; + let pattern = match pattern { + "correlated" => AccessPattern::Correlated, + _ => AccessPattern::Uniform, + }; + let indices = generate_indices(dataset.as_ref(), pattern); + let file = RUNTIME.block_on(async { + let path = dataset.path(Format::OnDiskVortex).await.unwrap(); + SESSION + .open_options() + .with_layout_reader_cache() + .open_path(path) + .await + .unwrap() + }); + Self { file, indices } + } + + fn indices(&self) -> StrictSortedBuffer { + StrictSortedBuffer::try_new(Buffer::from(self.indices.clone())).unwrap() + } + + fn take_v1(&self) { + let array = RUNTIME.block_on( + self.file + .scan() + .unwrap() + .with_row_indices(self.indices()) + .into_array_stream() + .unwrap() + .read_all(), + ); + black_box(array.unwrap()); + } + + fn scan_v2(&self, builder: impl FnOnce(ScanBuilder) -> ScanBuilder) { + let stream = builder(ScanBuilder::new(SESSION.clone(), self.file.scan_file())) + .into_stream() + .unwrap(); + let array = + RUNTIME.block_on(ArrayStreamAdapter::new(self.file.dtype().clone(), stream).read_all()); + black_box(array.unwrap()); + } +} + +/// The default scan, the target V2 should beat. +#[divan::bench(args = QUERIES)] +fn take_v1(bencher: Bencher, query: &str) { + let query = Query::open(query); + bencher.bench_local(|| query.take_v1()); +} + +#[divan::bench(args = QUERIES)] +fn take_v2(bencher: Bencher, query: &str) { + let query = Query::open(query); + bencher.bench_local(|| query.scan_v2(|scan| scan.with_row_indices(query.indices()))); +} + +/// A V2 scan that selects no rows: its per-scan cost before any split is planned. +#[divan::bench(args = QUERIES)] +fn setup_v2(bencher: Bencher, query: &str) { + let query = Query::open(query); + bencher.bench_local(|| query.scan_v2(|scan| scan.with_row_range(0..0))); +} diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 357e688ed53..8cf6bb590d2 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -72,7 +72,7 @@ const POISSON_EXPECTED_COUNT: usize = 100; const CACHED_WARMUP_DURATION: Duration = Duration::from_secs(1); /// Generate indices for the given dataset and access pattern. -fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec { +pub fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec { let row_count = dataset.row_count(); let mut rng = StdRng::seed_from_u64(42); diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 93a66ccd49e..095c4d4d5c5 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -29,10 +29,12 @@ use tokio::fs::File; use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; +use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; +use vortex::layout::scan::v2::ScanBuilder; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex::utils::aliases::hash_map::HashMap; @@ -126,7 +128,7 @@ impl RandomAccessor for ArrowIpcRandomAccessor { } } -/// Random accessor for Vortex format files. +/// Random accessor for Vortex format files, read with the V2 scan. /// /// The file handle is opened at construction time and reused across `take()` calls. pub struct VortexRandomAccessor { @@ -142,11 +144,7 @@ impl VortexRandomAccessor { name: impl Into, format: Format, ) -> anyhow::Result { - let file = SESSION - .open_options() - .with_layout_reader_cache() - .open_path(path.as_ref()) - .await?; + let file = SESSION.open_options().open_path(path.as_ref()).await?; Ok(Self { name: name.into(), format, @@ -163,7 +161,6 @@ impl VortexRandomAccessor { ) -> anyhow::Result { let file = SESSION .open_options() - .with_layout_reader_cache() .open_object_store(remote.store(), remote.key(path)?) .await?; Ok(Self { @@ -185,12 +182,11 @@ impl RandomAccessor for VortexRandomAccessor { } async fn take(&self, indices: &[u64]) -> anyhow::Result { - let indices_buf: Buffer = Buffer::from(indices.to_vec()); - let array = self - .file - .scan()? - .with_row_indices(StrictSortedBuffer::try_new(indices_buf)?) - .into_array_stream()? + 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 array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream) .read_all() .await?; diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 459429046c2..ecf614f092d 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -117,6 +117,11 @@ impl Layout { .into_typed() } + /// The first row of each chunk, followed by the layout's row count. + pub fn chunk_offsets(&self) -> &[u64] { + &self.data().chunk_offsets + } + /// Rebuild this layout with owned children. pub fn with_children(&self, children: Vec) -> Self { Self::new( diff --git a/vortex-layout/src/plan/physical/children.rs b/vortex-layout/src/plan/physical/children.rs index bfa1541b87c..21a1f8e73a1 100644 --- a/vortex-layout/src/plan/physical/children.rs +++ b/vortex-layout/src/plan/physical/children.rs @@ -30,7 +30,7 @@ impl PlanChildren { ) -> Self { Self { initializer: Some(Arc::new(initializer)), - cache: (0..len).map(|_| OnceCell::new()).collect::>().into(), + cache: (0..len).map(|_| OnceCell::new()).collect(), } } diff --git a/vortex-layout/src/plan/physical/lower.rs b/vortex-layout/src/plan/physical/lower.rs index 1ea284cc1ca..bba7266093c 100644 --- a/vortex-layout/src/plan/physical/lower.rs +++ b/vortex-layout/src/plan/physical/lower.rs @@ -83,22 +83,16 @@ fn lower_flat(layout: &FlatLayout) -> SegmentScanPlan { } fn lower_chunked(layout: &ChunkedLayout) -> VortexResult { - let mut row_offsets = Vec::with_capacity(layout.nchildren()); - let mut row_count = 0u64; - for index in 0..layout.nchildren() { - row_offsets.push(row_count); - row_count = row_count - .checked_add(layout.child_row_count(index)) - .ok_or_else(|| vortex_err!("Chunked row count overflow"))?; - } + let nchunks = layout.nchildren(); + let chunks = layout.to_layout(); // SAFETY: Chunked layout construction validates that every child has the parent dtype and - // that their row counts sum to the parent; offsets were computed with checked addition above. + // computes the chunk offsets with checked addition, ending at the parent's row count. Ok(unsafe { ConcatPlan::from_children_unchecked( layout.dtype().clone(), layout.row_count(), - row_offsets.into(), - lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), + layout.chunk_offsets()[..nchunks].into(), + PlanChildren::lazy(nchunks, move |index| lower_slot(&chunks, index)), ) }) } @@ -159,13 +153,17 @@ fn lazy_children(layout: LayoutRef, slots: Vec) -> PlanChildren { .get(index) .copied() .ok_or_else(|| vortex_err!("Missing plan child slot {index}"))?; - let child = layout - .slot(slot)? - .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?; - lower(&child) + lower_slot(&layout, slot) }) } +fn lower_slot(layout: &LayoutRef, slot: usize) -> VortexResult { + let child = layout + .slot(slot)? + .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?; + lower(&child) +} + fn lower_zoned(layout: &LayoutRef) -> VortexResult { // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary zones. let metadata = if let Some(layout) = layout.as_opt::() { diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 9d650ebbd26..8e9726fcd3a 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -671,15 +671,14 @@ impl Scan { 0, splits.iter().map(|split| split.rows.clone()).collect(), ); - // Only the rows some split reads can share a segment between readers. - let start = splits - .iter() - .map(|split| split.rows.start) - .min() - .unwrap_or(0); - let end = splits.iter().map(|split| split.rows.end).max().unwrap_or(0); + // Only the rows some split reads can share a segment between readers, so the rows + // between splits are never visited. match &root { - Root::Plan(plan) => core.shares.add(plan, start..end)?, + Root::Plan(plan) => { + for split in &splits { + core.shares.add(plan, split.rows.clone())?; + } + } } core.shares.retain_shared(); core.dirty_flags = vec![false; 1]; diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index fc6053d55ca..543a06f4ec1 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -28,6 +28,8 @@ use crate::segments::SegmentId; /// Rows a running plan produced, in row order. #[derive(Clone)] pub(super) struct Piece { + /// The scan split that produced the piece. Splits are in row order. + pub(super) split: usize, /// Orders the piece among the outputs of one scan. pub(super) rows: Range, /// The selected rows the piece holds. @@ -96,7 +98,7 @@ impl ScanAdaptor { } Ok(match self.step()? { Turn::Waiting | Turn::Done => AdaptorStep::Yield, - Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)), + Turn::Output(split, array) => AdaptorStep::Piece(self.piece(split, array)), Turn::Read(read) => { // Publish every read the scan has ready together, so the owner can coalesce them. let mut batch = vec![self.request(read.id, read.segment_id)?]; @@ -130,10 +132,11 @@ impl ScanAdaptor { Ok(turn) } - fn piece(&mut self, array: ArrayRef) -> Piece { + fn piece(&mut self, split: usize, array: ArrayRef) -> Piece { let start = self.next_piece; self.next_piece += 1; Piece { + split, rows: start..start + 1, array, } diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 9d8b8897ef3..4c054c090aa 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -15,13 +15,13 @@ 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_session::VortexSession; use self::projection::ProjectionPlanner; use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::Split; use crate::plan::physical::pipeline::ValueCache; /// Where a segment's bytes live in its source. @@ -66,11 +66,18 @@ pub(super) struct ScanPlans { struct SelectedRows { /// The split the rows belong to. scope: WorkScope, - /// Which rows of `scope.rows` are selected. - mask: Mask, + /// The pipeline splits producing the selected rows of `scope.rows`, in row order. + splits: Vec, } -/// The work for one split: project the rows of `scope` selected by `mask`. -pub(super) fn plan_projection(plans: ScanPlans, scope: WorkScope, mask: Mask) -> Box { - Box::new(ProjectionPlanner::new(plans, SelectedRows { scope, mask })) +/// The work for one split: project the rows of `scope` that `splits` select. +pub(super) fn plan_projection( + plans: ScanPlans, + scope: WorkScope, + splits: Vec, +) -> Box { + Box::new(ProjectionPlanner::new( + plans, + SelectedRows { scope, splits }, + )) } diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index b6180725c27..3f60aa482a4 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -17,7 +17,6 @@ use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; use crate::plan::physical::pipeline::Scan; -use crate::plan::physical::pipeline::Split; use crate::scan::execution::ScanPlans; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; @@ -55,7 +54,7 @@ impl Planner for ProjectionPlanner { let Some((plans, selected)) = self.work.take() else { return Ok(PlannerOutput::Done); }; - if selected.mask.all_false() { + if selected.splits.is_empty() { return Ok(PlannerOutput::Done); } Ok(PlannerOutput::Morsel( @@ -91,7 +90,9 @@ impl ProjectionMorsel { fn finish(&mut self) -> VortexResult { self.done = true; - self.pieces.sort_by_key(|piece| piece.rows.start); + // Splits run at once, so their pieces arrive interleaved. + self.pieces + .sort_by_key(|piece| (piece.split, piece.rows.start)); let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) .into_iter() .map(|piece| piece.array) @@ -135,16 +136,13 @@ impl Morsel for ProjectionMorsel { .plans .take() .vortex_expect("an uncompiled projection still has its plans"); - let scan = Scan::try_new( - plans.session, - plans.projection, - vec![Split { - rows: self.selected.scope.rows.clone(), - mask: self.selected.mask.clone(), - }], - )? - .with_row_offset(plans.row_offset) - .with_value_cache(plans.values); + let splits = std::mem::take(&mut self.selected.splits); + // Every split is active at once, so all of their reads are published together. + let active = splits.len(); + let scan = Scan::try_new(plans.session, plans.projection, splits)? + .with_max_active(active) + .with_row_offset(plans.row_offset) + .with_value_cache(plans.values); self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); return Ok(MorselOutput::Continue); }; diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index ce565846d7e..dfa413080a2 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -34,26 +34,73 @@ use crate::segments::SegmentFuture; 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(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { - let ranges = locations.iter().enumerate().map(|(id, location)| { - let id = SegmentId::from(u32::try_from(id).unwrap_or(u32::MAX)); - ((location.offset, location.length as usize), id) - }); - Arc::new(ranges.collect()) +/// Finds a file's segment by its byte range, to answer a range request with the segment. +#[derive(Clone)] +pub(super) enum SegmentRanges { + /// Every location, indexed by segment id and in offset order, as writers lay segments out. + /// Searched in place, so a scan builds nothing per segment. + Sorted(Arc<[SegmentLocation]>), + /// Segment ids keyed by `(offset, length)`, for locations in any other order. + Keyed(Arc>), +} + +impl SegmentRanges { + pub(super) fn new(locations: Arc<[SegmentLocation]>) -> Self { + if locations.is_sorted_by_key(|location| location.offset) { + return Self::Sorted(locations); + } + let ranges = locations.iter().enumerate().map(|(id, location)| { + let id = SegmentId::from(u32::try_from(id).unwrap_or(u32::MAX)); + ((location.offset, location.length as usize), id) + }); + Self::Keyed(Arc::new(ranges.collect())) + } + + fn find(&self, offset: u64, len: usize) -> Option { + match self { + Self::Sorted(locations) => { + let first = locations.partition_point(|location| location.offset < offset); + let index = locations[first..] + .iter() + .take_while(|location| location.offset == offset) + .position(|location| location.length as usize == len)?; + u32::try_from(first + index).ok().map(SegmentId::from) + } + Self::Keyed(ranges) => ranges.get(&(offset, len)).copied(), + } + } } /// Serves splits from any [`SegmentSource`], by segment id. pub(super) struct SegmentScanIo { segments: Arc, ranges: SegmentRanges, + shared: SharedReads, } impl SegmentScanIo { pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { - Self { segments, ranges } + Self { + segments, + ranges, + shared: SharedReads::default(), + } + } + + /// Expects one more session, the next planner's, to read each of `segments`, which the + /// session of the planner handed out before it reads too. Each segment is read once for all + /// the sessions expected to read it. + pub(super) fn share_with_next(&self, segments: impl IntoIterator) { + let mut shared = self.shared.lock(); + for segment in segments { + shared + .entry(segment) + .or_insert(SharedRead { + readers: 1, + read: None, + }) + .readers += 1; + } } } @@ -61,11 +108,22 @@ impl IoService for SegmentScanIo { fn session(&self) -> Arc { Arc::new(SegmentIoSource::new( Arc::clone(&self.segments), - Arc::clone(&self.ranges), + self.ranges.clone(), + Arc::clone(&self.shared), )) } } +/// Reads several sessions of one scan make of the same segment, by segment id. +type SharedReads = Arc>>; + +/// A segment read for several sessions: started by the first and dropped once each has taken it. +struct SharedRead { + /// Sessions still to take the read. + readers: usize, + read: Option, +} + /// Serves a split's range requests from a segment source, without blocking. /// /// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once, @@ -78,6 +136,7 @@ impl IoService for SegmentScanIo { pub(super) struct SegmentIoSource { segments: Arc, ranges: SegmentRanges, + shared: SharedReads, reads: Mutex>>, ready: Mutex>, /// Reads started by prefetches, consumed by the next fetch of their segment. @@ -90,16 +149,41 @@ pub(super) struct SegmentIoSource { type Prefetched = Shared>>>; impl SegmentIoSource { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + fn new(segments: Arc, ranges: SegmentRanges, shared: SharedReads) -> Self { Self { segments, ranges, + shared, reads: Mutex::new(FuturesUnordered::new()), ready: Mutex::default(), prefetched: Mutex::default(), announced: Mutex::default(), } } + + /// Requests `segment`, from the read other sessions share when the scan shares it. + fn request(&self, segment: SegmentId) -> SegmentFuture { + let mut shared = self.shared.lock(); + let Some(entry) = shared.get_mut(&segment) else { + return self.segments.request(segment); + }; + let read = entry + .read + .get_or_insert_with(|| { + self.segments + .request(segment) + .map_err(Arc::new) + .boxed() + .shared() + }) + .clone(); + entry.readers = entry.readers.saturating_sub(1); + if entry.readers == 0 { + shared.remove(&segment); + } + read.map_err(move |err| vortex_err!("shared read of segment {segment} failed: {err}")) + .boxed() + } } impl IoSource for SegmentIoSource { @@ -114,16 +198,16 @@ impl IoSource for SegmentIoSource { request.target ); }; - let segment = *self + let segment = self .ranges - .get(&(offset, len)) + .find(offset, len) .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; if request.intent == IoIntent::Announce { self.announced .lock() .entry(segment) .or_default() - .push_back(self.segments.request(segment)); + .push_back(self.request(segment)); continue; } if request.intent == IoIntent::Prefetch { @@ -133,7 +217,7 @@ impl IoSource for SegmentIoSource { .lock() .get_mut(&segment) .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) + .unwrap_or_else(|| self.request(segment)) .map_err(Arc::new) .boxed() .shared(); @@ -157,7 +241,7 @@ impl IoSource for SegmentIoSource { .lock() .get_mut(&segment) .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)), + .unwrap_or_else(|| self.request(segment)), }; let id = request.request; // In-memory and cached segments need no allocated asynchronous task. A pending @@ -243,11 +327,14 @@ mod tests { use vortex_io::request::IoOwnerId; use vortex_io::request::IoRequest; use vortex_io::request::IoRequestId; + use vortex_io::request::IoService; use vortex_io::request::IoSource; use vortex_io::request::IoTarget; use super::SegmentIoSource; - use super::segment_ranges; + use super::SegmentRanges; + use super::SegmentScanIo; + use super::SharedReads; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; use crate::segments::SegmentId; @@ -265,6 +352,29 @@ mod tests { } } + fn location(offset: u64, length: u32) -> SegmentLocation { + SegmentLocation { + offset, + length, + alignment: Alignment::none(), + } + } + + /// Ranges are found by search when locations are in offset order, where a zero-length + /// segment can share its offset with the next, and by key otherwise. + #[test] + fn segment_ranges_find_each_segment() { + let sorted = SegmentRanges::new([location(0, 4), location(4, 0), location(4, 8)].into()); + assert!(matches!(sorted, SegmentRanges::Sorted(_))); + let unsorted = SegmentRanges::new([location(4, 0), location(4, 8), location(0, 4)].into()); + assert!(matches!(unsorted, SegmentRanges::Keyed(_))); + for (ranges, ids) in [(sorted, [0, 1, 2]), (unsorted, [2, 0, 1])] { + let found = [(0, 4), (4, 0), (4, 8)].map(|(offset, len)| ranges.find(offset, len)); + assert_eq!(found, ids.map(|id| Some(SegmentId::from(id)))); + assert_eq!(ranges.find(4, 4), None); + } + } + fn request(intent: IoIntent, id: u32) -> IoRequest { IoRequest { intent, @@ -301,7 +411,11 @@ mod tests { length: 4, alignment: Alignment::none(), }); - let io = SegmentIoSource::new(segments, segment_ranges(&locations)); + let io = SegmentIoSource::new( + segments, + SegmentRanges::new(locations.into()), + SharedReads::default(), + ); let mut pending = request(IoIntent::Fetch, 1); pending.target = IoTarget::range(4, 4); io.submit(IoOwnerId(0), vec![pending, request(IoIntent::Fetch, 0)])?; @@ -324,6 +438,26 @@ mod tests { Ok(()) } + /// Sessions of a scan that share a segment read it once, and a later session reads it again + /// once every expected reader has taken it. + #[tokio::test] + async fn sessions_share_a_segment_read() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let io = SegmentScanIo::new( + Arc::clone(&segments) as _, + SegmentRanges::new([location(0, 4)].into()), + ); + io.share_with_next([SegmentId::from(0)]); + for reads in [1, 1, 2] { + let session = io.session(); + session.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0)])?; + let completion = std::future::poll_fn(|cx| session.poll_completion(cx)).await?; + assert!(completion.result.is_ok()); + assert_eq!(segments.0.load(Ordering::Relaxed), reads); + } + Ok(()) + } + /// A fetch of a prefetched segment waits on the prefetch's read rather than reading again, /// and only the fetch is delivered. #[tokio::test] @@ -334,7 +468,11 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let io = SegmentIoSource::new( + Arc::clone(&segments) as _, + SegmentRanges::new([location].into()), + SharedReads::default(), + ); io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; assert_eq!(segments.0.load(Ordering::Relaxed), 1); @@ -362,7 +500,11 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let io = SegmentIoSource::new( + Arc::clone(&segments) as _, + SegmentRanges::new([location].into()), + SharedReads::default(), + ); io.submit( IoOwnerId(0), vec![ diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 46cbde987fc..4f8c6f5dc0e 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -4,21 +4,26 @@ //! 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 -//! 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. +//! through the planning protocol. A single file planner hands out the scan's units one at a time +//! to a [`Pool`] of workers, each advancing the planners and morsels of the units it claimed on a +//! [`Run`](vortex_scan::planning::driver::Run) of its own, and every segment read is a protocol +//! request served from the file's [`SegmentSource`]. The layout reader is never asked to evaluate +//! anything. //! //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] -//! directly and supports projection, row ranges, and selections. Unsupported layouts return an -//! error; filters and limits are not part of this builder's API. Work is partitioned into fixed -//! chunks of up to 65,536 rows, independently of the stored layout's chunk boundaries. -//! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. +//! directly and supports projection, row ranges, and selections. Layouts without a pipeline +//! implementation return an error; filters and limits are not part of this builder's API. Work is +//! partitioned into fixed chunks of up to 65,536 rows, independently of the stored layout's chunk +//! boundaries. Full-file projections use smaller 8,192-row multiples when serialized bytes per row +//! are high. A partial selection is instead planned by the stored segments its rows read, so each +//! segment is read once and rows between selected runs are never visited. mod io; -mod split; #[cfg(test)] mod tests; +use std::cmp::Ordering; +use std::collections::VecDeque; use std::ops::Range; use std::sync::Arc; @@ -29,6 +34,7 @@ use arrow_schema::SchemaRef; use futures::Stream; use futures::StreamExt; use futures::TryStreamExt; +use futures::future; use futures::stream; use futures::stream::BoxStream; use vortex_array::ArrayRef; @@ -36,37 +42,47 @@ use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_io::request::IoService; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::AllOr; +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::planning::pool::Pool; +use vortex_scan::planning::pool::PoolOptions; +use vortex_scan::planning::pool::Spawn; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; -use crate::layouts::chunked::Chunked; -use crate::layouts::dict::Dict; -use crate::layouts::flat::Flat; -use crate::layouts::struct_::Struct; -use crate::layouts::zoned::LegacyStats; -use crate::layouts::zoned::Zoned; use crate::plan::physical::lower; use crate::plan::physical::optimize; +use crate::plan::physical::pipeline::Reach; +use crate::plan::physical::pipeline::Split; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; use crate::scan::execution::ScanPlans; pub use crate::scan::execution::SegmentLocation; use crate::scan::execution::plan_projection; +use crate::scan::v2::io::SegmentRanges; use crate::scan::v2::io::SegmentScanIo; -use crate::scan::v2::io::segment_ranges; -use crate::scan::v2::split::SplitPlan; +use crate::segments::SegmentId; use crate::segments::SegmentSource; const ROW_BLOCK_SIZE: usize = 8192; const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; +/// A chunk selected in more runs than this is planned whole: finding the segments of each run +/// would cost more than the chunks it could skip. +const MAX_RUNS: usize = 64; +/// A planner grows to share stored segments only while it spans at most this many rows, so a +/// segment spanning much of the file cannot join every selected run into one task. +const MAX_PLANNER_ROWS: u64 = 4 * MAX_ROWS_PER_CHUNK as u64; /// Limit estimated bytes per task without changing the writer-aligned row granularity. fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { @@ -82,6 +98,22 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE } +/// The rows of one planner's pipeline splits, in row order. +type PlannerSplits = Vec>; + +/// Whether two sorted segment lists have a segment in common. +fn shares_any(left: &[SegmentId], right: &[SegmentId]) -> bool { + let (mut left, mut right) = (left.iter().peekable(), right.iter().peekable()); + while let (Some(l), Some(r)) = (left.peek(), right.peek()) { + match l.cmp(r) { + Ordering::Less => _ = left.next(), + Ordering::Greater => _ = right.next(), + Ordering::Equal => return true, + } + } + false +} + /// 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() { @@ -109,8 +141,9 @@ pub struct ScanFile { /// Builds a projection 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. Filters and row limits are not supported. +/// This builder reads the file's layout and segments directly. A layout with no pipeline +/// implementation returns an error when the first split that reads it is planned. Filters and row +/// limits are not supported. pub struct ScanBuilder { session: VortexSession, file: ScanFile, @@ -223,165 +256,297 @@ impl ScanBuilder { } } - /// Runs the scan on the session's runtime, preparing it on the first poll. - /// Preparation errors are returned as stream items. + /// Runs the scan on a [`Pool`] of workers on the session's runtime, starting it on the first + /// poll. Errors, including those preparing the scan, are returned as stream items. + /// + /// The scan's file planner hands out its units one at a time, and each is run by the next + /// spare worker while the file planner plans the next. Up to `concurrency` units run at once + /// per worker thread, and their batches are yielded in row order when the scan is ordered. pub fn into_stream(self) -> VortexResult>> { let handle = self.session.handle(); - Ok(stream::once(async move { - handle - .spawn_cpu(move || PreparedScan::try_new(self)?.execute_stream()) - .await - }) - .try_flatten() - .boxed()) - } -} - -/// Whether every required layout lowers to plans with a pipeline implementation. -/// -/// List layouts lower to a list assembly, which cannot compile before its offsets are read. -fn supports_layout(layout: &LayoutRef, identity: bool) -> VortexResult { - if !(layout.is::() - || layout.is::() - || layout.is::() - || layout.is::() - || layout.is::() - || layout.is::()) - { - return Ok(false); - } - // An identity projection cannot use zone statistics. Leave their layouts uninitialized. - if identity && (layout.is::() || layout.is::()) { - return match layout.slot(0)? { - Some(data) => supports_layout(&data, identity), - None => Ok(false), + let spawn: Spawn = Arc::new(move |worker| handle.spawn(worker).detach()); + let io = Arc::new(SegmentScanIo::new( + Arc::clone(&self.file.segments), + SegmentRanges::new(Arc::clone(&self.file.locations)), + )); + let options = PoolOptions { + max_workers: get_available_parallelism().unwrap_or(1), + units_per_worker: self.concurrency, + ordered: self.ordered, }; + let map = self.map_fn; + let root = FilePlanner { + scan: Some(ScanRequest { + session: self.session, + file: self.file, + projection: self.projection, + row_range: self.row_range, + selection: self.selection, + row_offset: self.row_offset, + }), + io: Arc::clone(&io), + planners: None, + }; + let scope = WorkScope { + file_ordinal: 0, + rows: 0..0, + }; + Ok(stream::once(future::lazy(move |_| { + Pool::stream(Box::new(root), scope, io, map, spawn, options) + })) + .flatten() + .boxed()) } - for child in layout.children()? { - if !supports_layout(&child, identity)? { - return Ok(false); - } - } - Ok(true) } -/// A scan planned over its file, ready to run its splits. -struct PreparedScan { +/// What a scan reads, before it is planned. +struct ScanRequest { session: VortexSession, - plans: ScanPlans, - /// Serves the splits' reads. - io: Arc, - /// The scan's splits, over the root row domain, in row order. - splits: Vec>, + file: ScanFile, + projection: BoundExpression, + row_range: Option>, selection: Selection, - map_fn: Arc VortexResult + Send + Sync>, - ordered: bool, - concurrency: usize, + row_offset: u64, } -impl PreparedScan { +impl ScanRequest { /// Lowers the file's layout, plans the projection, and partitions the selected range into - /// fixed-size row chunks. - fn try_new(builder: ScanBuilder) -> VortexResult { - let file = builder.file; - vortex_ensure!( - supports_layout(&file.layout, builder.projection.is_root())?, - "V2 scans do not support this file's layout" - ); + /// fixed-size row chunks, to plan units over. + fn plan(self, io: Arc) -> VortexResult { + let file = self.file; let root = lower(&file.layout)?; let row_count = root.row_count(); // An identity projection needs no rewrites. Keep the lowered children lazy, including // statistics that a projection-only scan never reads. - let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype(); + let identity = self.projection.is_root() && self.projection.dtype() == root.dtype(); let projection = if identity { root } else { - optimize(plan_row_idx_expression(builder.projection, root)?)? + optimize(plan_row_idx_expression(self.projection, root)?)? }; let chunk_rows = if identity { rows_per_chunk(row_count, &file.locations) } else { MAX_ROWS_PER_CHUNK }; - let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); - let io = Arc::new(SegmentScanIo::new( - file.segments, - segment_ranges(&file.locations), - )); - Ok(Self { - session: builder.session.clone(), + let chunks = row_chunks(self.row_range.unwrap_or(0..row_count), chunk_rows); + Ok(Planners { plans: ScanPlans { - session: builder.session, + session: self.session, locations: file.locations, projection, - row_offset: builder.row_offset, + row_offset: self.row_offset, values: Default::default(), }, + selection: self.selection, io, - splits, - selection: builder.selection, - map_fn: builder.map_fn, - ordered: builder.ordered, - concurrency: builder.concurrency, + chunks: chunks.into_iter(), + runs: VecDeque::new(), + whole: None, + splits: Vec::new(), + last: Vec::new(), + reached: Vec::new(), }) } +} - /// 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 row_mask = self.selection.row_mask(rows); - if row_mask.mask().all_false() { - return None; - } - let scope = WorkScope { - file_ordinal: 0, - rows: row_mask.row_range(), - }; - Some(SplitPlan { - root: plan_projection( - self.plans.clone(), - scope.clone(), - row_mask.mask().clone(), - ), - scope, - }) - }) - .collect() +/// The scan's root planner: plans the scan on its first step, and hands out one unit per step +/// after that, in row order. Each unit is a planner that projects its splits. +struct FilePlanner { + /// The scan, until it is planned. + scan: Option, + io: Arc, + /// The scan's units, once it is planned, until every one has been handed out. + planners: Option, +} + +impl IoConsumer for FilePlanner { + // The file planner publishes no requests. + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for FilePlanner { + fn state(&self) -> State { + if self.scan.is_some() || self.planners.is_some() { + State::NeedsCompute + } else { + State::Done + } } - /// Runs every split as a task on the session's runtime and yields their 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. - 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; - // 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::>>() + fn compute(&mut self) -> VortexResult { + if let Some(scan) = self.scan.take() { + self.planners = Some(scan.plan(Arc::clone(&self.io))?); + } + let Some(planners) = &mut self.planners else { + return Ok(PlannerOutput::Done); + }; + match planners.next_plan()? { + Some((scope, root)) => Ok(PlannerOutput::Planner(scope, root)), + None => { + self.planners = None; + Ok(PlannerOutput::Done) + } + } + } +} + +/// What the next fixed chunk holds for planning. +enum NextRun { + /// A selected run to group into a planner. + Run(Range), + /// The selected span of a chunk planned whole. + Whole(Range), + /// Every chunk has been visited. + Done, +} + +/// Plans a scan's units one at a time, in row order, so each can be handed to a worker as soon +/// as it is planned, while the rest are still being planned. +/// +/// A fully selected chunk, or one selected in more than [`MAX_RUNS`] runs, is one split of its +/// own planner. Otherwise each selected run joins the planner before it, even across chunk +/// boundaries, while the two read a stored segment in common and the planner spans at most +/// [`MAX_PLANNER_ROWS`], so a stored segment is decoded by one planner unless it spans more rows +/// than that. Within a planner, a run extends the split before it when both read exactly the +/// same segments: every row between them then lies in those segments' chunks, so the split visits +/// no chunk the selection skips. +/// +/// A column's segments hold contiguous rows, so a run can only share a segment with an earlier +/// run of its planner if it shares one with the run just before it, and a segment several +/// planners read is read by neighbouring planners. So a planner is complete once the run after it +/// is known, and the segments it shares with that run are all it shares with later planners: the +/// scan's IO reads each of them once for both. +struct Planners { + plans: ScanPlans, + selection: Selection, + io: Arc, + /// The fixed chunks not yet visited. + chunks: std::vec::IntoIter>, + /// The selected runs of the last chunk visited, not yet grouped. + runs: VecDeque>, + /// A chunk planned whole, waiting for the planner before it to be handed out. + whole: Option>, + /// The pipeline splits of the planner being grouped. + splits: PlannerSplits, + /// The segments the last run of the planner reads, and those all its runs read. + last: Vec, + reached: Vec, +} + +impl Planners { + /// The next unit, with the rows it covers. + fn next_plan(&mut self) -> VortexResult)>> { + let Some(ranges) = self.next_splits()? else { + return Ok(None); + }; + let rows = ranges[0].start..ranges[ranges.len() - 1].end; + let splits = ranges + .into_iter() + .map(|rows| Split { + mask: self.selection.row_mask(&rows).mask().clone(), + rows, }) - }); - let batches = if self.ordered { - tasks.buffered(concurrency).boxed() - } else { - tasks.buffer_unordered(concurrency).boxed() + .collect(); + let scope = WorkScope { + file_ordinal: 0, + rows, }; - Ok(batches - .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) - .try_flatten() - .boxed()) + let root = plan_projection(self.plans.clone(), scope.clone(), splits); + Ok(Some((scope, root))) + } + + /// The pipeline splits of the next planner, in row order. + fn next_splits(&mut self) -> VortexResult> { + if let Some(whole) = self.whole.take() { + return Ok(Some(vec![whole])); + } + loop { + match self.next_run() { + NextRun::Done => { + return Ok((!self.splits.is_empty()).then(|| self.close(&[]))); + } + NextRun::Whole(span) => { + if self.splits.is_empty() { + return Ok(Some(vec![span])); + } + self.whole = Some(span); + return Ok(Some(self.close(&[]))); + } + NextRun::Run(run) => { + let mut segments = Vec::new(); + self.plans.projection.reach( + run.clone(), + &Reach::Offset(0), + &mut |segment, _| segments.push(segment), + )?; + segments.sort_unstable(); + segments.dedup(); + let joins = self.splits.first().is_some_and(|first| { + run.end - first.start <= MAX_PLANNER_ROWS + && shares_any(&self.last, &segments) + }); + if joins { + match self.splits.last_mut() { + Some(split) if segments == self.last => split.end = run.end, + _ => self.splits.push(run), + } + self.reached.extend_from_slice(&segments); + self.last = segments; + continue; + } + let planner = (!self.splits.is_empty()).then(|| self.close(&segments)); + self.splits.push(run); + self.reached.extend_from_slice(&segments); + self.last = segments; + if planner.is_some() { + return Ok(planner); + } + } + } + } + } + + /// Hands out the planner being grouped. The segments it shares with `next`, those the run + /// after it reads, are read once for it and the planner that run starts. + fn close(&mut self, next: &[SegmentId]) -> PlannerSplits { + self.reached.sort_unstable(); + self.reached.dedup(); + self.io.share_with_next( + self.reached + .iter() + .filter(|segment| next.binary_search(segment).is_ok()) + .copied(), + ); + self.reached.clear(); + self.last.clear(); + std::mem::take(&mut self.splits) + } + + fn next_run(&mut self) -> NextRun { + loop { + if let Some(run) = self.runs.pop_front() { + return NextRun::Run(run); + } + let Some(rows) = self.chunks.next() else { + return NextRun::Done; + }; + let row_mask = self.selection.row_mask(&rows); + let mask = row_mask.mask(); + match mask.slices() { + AllOr::None => {} + AllOr::Some(slices) if slices.len() <= MAX_RUNS => self.runs.extend( + slices + .iter() + .map(|&(start, end)| rows.start + start as u64..rows.start + end as u64), + ), + AllOr::All | AllOr::Some(_) => { + let span = mask.first().unwrap_or(0)..mask.last().map_or(0, |last| last + 1); + return NextRun::Whole( + rows.start + span.start as u64..rows.start + span.end as u64, + ); + } + } + } } } diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs deleted file mode 100644 index b9371735d39..00000000000 --- a/vortex-layout/src/scan/v2/split.rs +++ /dev/null @@ -1,47 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::future; -use std::sync::Arc; - -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 { - /// 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 a910172f87b..b098d6089bb 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -33,14 +33,21 @@ use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; 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::Planner; +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::Planners; use super::ScanBuilder; use super::ScanFile; +use super::io::SegmentRanges; +use super::io::SegmentScanIo; use super::row_chunks; use super::rows_per_chunk; use crate::LayoutRef; @@ -58,6 +65,51 @@ use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::segments::TestSegments; + +/// Plans `builder`'s scan, returning its units' planner and the IO they read through. +fn plan(builder: ScanBuilder) -> VortexResult<(Planners, Arc)> { + let io = Arc::new(SegmentScanIo::new( + Arc::clone(&builder.file.segments), + SegmentRanges::new(Arc::clone(&builder.file.locations)), + )); + let scan = super::ScanRequest { + session: builder.session, + file: builder.file, + projection: builder.projection, + row_range: builder.row_range, + selection: builder.selection, + row_offset: builder.row_offset, + }; + Ok((scan.plan(Arc::clone(&io))?, io)) +} + +/// Every unit of a planned scan, in row order. +fn units(planners: &mut Planners) -> VortexResult)>> { + let mut units = Vec::new(); + while let Some(unit) = planners.next_plan()? { + units.push(unit); + } + Ok(units) +} + +/// Runs one unit on a run of its own and returns its batches in row order. +async fn run_unit( + (scope, root): (WorkScope, Box), + io: &SegmentScanIo, +) -> VortexResult> { + let mut run = Run::new(); + run.admit(root, scope, io.session()); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::Waiting => std::future::poll_fn(|cx| run.poll_completion(cx)).await?, + _ => break, + } + } + batches.sort_by_key(|batch| batch.scope.rows.start); + Ok(batches.into_iter().map(|batch| batch.array).collect()) +} use crate::sequence::SequenceId; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt as _; @@ -384,16 +436,16 @@ 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 (mut planners, io) = plan(builder(&session, &segments, &layout)?)?; + let splits = units(&mut planners)?; assert_eq!( splits.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(); + let rows = split.0.rows.clone(); let before = counted.reads.load(Ordering::Relaxed); - let arrays = split.run(Arc::clone(&scan.io)).await?; + let arrays = run_unit(split, &io).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; @@ -416,7 +468,97 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { Ok(()) } -/// V2 rejects list layouts: their element ranges are only known once their offsets are read. +/// Three stored chunks of 70,000 consecutive integers, read through a counting source. +async fn counted_ints(session: &VortexSession) -> VortexResult<(Arc, LayoutRef)> { + const STORED_CHUNK_ROWS: i32 = 70_000; + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let arrays = (0..3) + .map(|index| { + let start = index * STORED_CHUNK_ROWS; + Ok(( + sequence.advance(), + PrimitiveArray::from_iter(start..start + STORED_CHUNK_ROWS).into_array(), + )) + }) + .collect::>>(); + let dtype = DType::Primitive(vortex_array::dtype::PType::I32, false.into()); + let layout = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(arrays)).sendable(), + eof, + session, + ) + .await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + Ok((counted, layout)) +} + +#[rstest] +#[case::one_stored_chunk(vec![5, 6, 30_000, 140_100], vec![5..30_001, 140_100..140_101], 2)] +// Rows in one work chunk but different stored chunks plan nothing between them. +#[case::two_stored_chunks(vec![66_000, 75_000], vec![66_000..66_001, 75_000..75_001], 2)] +// Rows in different work chunks but one stored chunk read it once. +#[case::two_work_chunks(vec![60_000, 66_000], vec![60_000..66_001], 1)] +#[tokio::test] +async fn splits_read_each_stored_chunk_once( + #[case] rows: Vec, + #[case] scopes: Vec>, + #[case] reads: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (counted, layout) = counted_ints(&session).await?; + let segments = Arc::clone(&counted) as Arc; + let scan = || { + plan( + builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), + ) + }; + // Runs reading the same stored chunk are one pipeline split, which visits no other chunk. + let (mut planners, _) = scan()?; + let mut pipeline_splits = Vec::new(); + while let Some(splits) = planners.next_splits()? { + pipeline_splits.extend(splits); + } + assert_eq!( + pipeline_splits, scopes, + "each planner here reads one stored chunk" + ); + let (mut planners, io) = scan()?; + let splits = units(&mut planners)?; + assert_eq!( + splits + .iter() + .map(|(scope, _)| scope.rows.clone()) + .collect::>(), + scopes + ); + let mut arrays = Vec::new(); + for split in splits { + arrays.extend(run_unit(split, &io).await?); + } + // Splits read only the stored chunks holding their selected rows, each chunk once. + assert_eq!(counted.reads.load(Ordering::Relaxed), reads); + let expected = rows + .iter() + .map(|&row| i32::try_from(row)) + .collect::, _>>()?; + assert_arrays_eq!( + ChunkedArray::try_new(arrays, layout.dtype().clone())?.into_array(), + expected.into_array(), + &mut session.create_execution_ctx() + ); + Ok(()) +} + +/// V2 rejects list layouts when it plans them: their element ranges are only known once their +/// offsets are read. #[tokio::test] async fn list_layouts_return_an_error() -> VortexResult<()> { let session = new_session().with_tokio(); @@ -428,7 +570,7 @@ async fn list_layouts_return_an_error() -> VortexResult<()> { .try_next() .await .expect_err("V2 must reject an unsupported layout"); - assert!(error.to_string().contains("V2 scans do not support")); + assert!(error.to_string().contains("has no pipeline implementation")); Ok(()) } @@ -436,8 +578,8 @@ async fn list_layouts_return_an_error() -> 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 chunks: Vec = scan.execute_stream()?.try_collect().await?; + let scan = builder(&session, &segments, &layout)?.with_ordered(false); + let chunks: Vec = scan.into_stream()?.try_collect().await?; let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); assert_eq!(rows, CHUNKS * CHUNK_ROWS); Ok(()) @@ -448,14 +590,12 @@ 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( - builder(&session, &segments, &layout)? - .with_projection(get_item("ints", root()).bind(&dtype)?) - .with_row_range(990..1010), - )?; + let scan = 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()?, + scan.into_stream()?, ) .await?; assert_arrays_eq!( @@ -474,9 +614,13 @@ 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 arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; + let (planners, _) = plan(builder(&session, &segments, &layout)?)?; + assert_eq!( + planners.chunks.collect::>(), + vec![0..(CHUNKS * CHUNK_ROWS) as u64] + ); + let scan = builder(&session, &segments, &layout)?; + let arrays: Vec<_> = scan.into_stream()?.try_collect().await?; assert_eq!(arrays.len(), 1); assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) diff --git a/vortex-scan/Cargo.toml b/vortex-scan/Cargo.toml index d12c207b049..a54e2ad63a9 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -25,11 +25,11 @@ vortex-utils = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } +parking_lot = { workspace = true } roaring = { workspace = true } tracing = { workspace = true } [dev-dependencies] -parking_lot = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs index 8cb78099272..329ddf964b6 100644 --- a/vortex-scan/src/planning/driver/mod.rs +++ b/vortex-scan/src/planning/driver/mod.rs @@ -15,6 +15,11 @@ //! for IO, and the caller decides how to wait, by awaiting [`Run::poll_completion`] or by handing //! a completion it took from a source to [`Run::complete`]. [`Driver::run`] is the blocking //! convenience for one root. +//! +//! A root admitted with [`Run::admit_distributing`] runs nothing it emits: each child is handed +//! back as a [`Progress::Handoff`], for another run to admit with [`Run::admit_child`]. A +//! [`Pool`](crate::planning::pool::Pool) uses this to hand a file planner's units to whichever +//! worker is spare. use std::collections::BTreeMap; use std::collections::VecDeque; @@ -59,10 +64,30 @@ pub struct Batch { pub array: ArrayRef, } +/// A planner or morsel, run as the root of a run. +pub enum Child { + /// A planner. + Planner(Box), + /// A morsel. + Morsel(Box), +} + +/// A child a distributing root emitted, for another run to admit. +pub struct Handoff { + /// The distributing root that emitted the child. + pub root: RootId, + /// The scope the child covers. + pub scope: WorkScope, + /// The child, which has not been visited. + pub child: Child, +} + /// What a call to [`Run::advance`] left behind. pub enum Progress { /// A morsel's batch. One morsel's batches arrive in the order it emitted them. Batch(Batch), + /// A child of a root admitted with [`Run::admit_distributing`], in emission order. + Handoff(Handoff), /// Every item descended from the root has retired, and its IO session was cleared. RootDone(RootId), /// Every live item waits for IO. Await [`Run::poll_completion`], or pass the next completion @@ -72,11 +97,6 @@ pub enum Progress { Idle, } -enum Item { - Planner(Box), - Morsel(Box), -} - /// What one item's `compute()` produced. enum Output { Planner(PlannerOutput), @@ -89,7 +109,7 @@ struct Work { id: IoOwnerId, root: RootId, scope: WorkScope, - item: Item, + item: Child, /// Stable child-ordinal path, beginning with the root's admission ordinal. path: Vec, /// The parent continuation sorts after all previously emitted child subtrees. @@ -115,15 +135,15 @@ impl Work { fn state(&self) -> State { match &self.item { - Item::Planner(planner) => planner.state(), - Item::Morsel(morsel) => morsel.state(), + Child::Planner(planner) => planner.state(), + Child::Morsel(morsel) => morsel.state(), } } fn consumer(&mut self) -> &mut dyn IoConsumer { match &mut self.item { - Item::Planner(planner) => &mut **planner, - Item::Morsel(morsel) => &mut **morsel, + Child::Planner(planner) => &mut **planner, + Child::Morsel(morsel) => &mut **morsel, } } @@ -179,6 +199,8 @@ const IO_POLL_INTERVAL: usize = 64; struct Root { io: Arc, live: usize, + /// Whether the children the root emits are handed off instead of run. + distributes: bool, } /// A set of root planners advanced by its caller on one thread at a time. @@ -229,11 +251,44 @@ impl Run { root: Box, scope: WorkScope, io: Arc, + ) -> RootId { + self.admit_child(Child::Planner(root), scope, io) + } + + /// Admits a planner or morsel, such as one handed off by another run, as a root. + pub fn admit_child(&mut self, child: Child, scope: WorkScope, io: Arc) -> RootId { + self.admit_root(child, scope, io, false) + } + + /// Admits `root` without running what it emits: each child it emits is handed back by + /// [`advance`](Self::advance) as a [`Progress::Handoff`], in emission order. + pub fn admit_distributing( + &mut self, + root: Box, + scope: WorkScope, + io: Arc, + ) -> RootId { + self.admit_root(Child::Planner(root), scope, io, true) + } + + fn admit_root( + &mut self, + root: Child, + scope: WorkScope, + io: Arc, + distributes: bool, ) -> RootId { let id = RootId(self.next_root); self.next_root += 1; - self.roots.insert(id, Root { io, live: 0 }); - self.spawn(id, scope, Item::Planner(root), vec![id.0]); + self.roots.insert( + id, + Root { + io, + live: 0, + distributes, + }, + ); + self.spawn(id, scope, root, vec![id.0]); id } @@ -375,7 +430,7 @@ impl Run { Ok(()) } - fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec) { + fn spawn(&mut self, root: RootId, scope: WorkScope, item: Child, path: Vec) { let id = IoOwnerId(self.next_owner); self.next_owner += 1; if let Some(root) = self.roots.get_mut(&root) { @@ -431,11 +486,11 @@ impl Run { /// Visits one item: computes once if it is ready, and settles it otherwise. fn visit(&mut self, mut work: Box) -> VortexResult<()> { let output = match &mut work.item { - Item::Planner(planner) => match planner.state() { + Child::Planner(planner) => match planner.state() { State::NeedsCompute => Output::Planner(planner.compute()?), _ => return self.settle(work), }, - Item::Morsel(morsel) => match morsel.state() { + Child::Morsel(morsel) => match morsel.state() { State::NeedsCompute => Output::Morsel(morsel.compute()?), _ => return self.settle(work), }, @@ -453,13 +508,11 @@ impl Run { self.settle(work)?; } Output::Planner(PlannerOutput::Planner(scope, child)) => { - let path = work.child_path(); - self.spawn(work.root, scope, Item::Planner(child), path); + self.emit(&mut work, scope, Child::Planner(child)); self.queue.insert(work); } Output::Planner(PlannerOutput::Morsel(scope, morsel)) => { - let path = work.child_path(); - self.spawn(work.root, scope, Item::Morsel(morsel), path); + self.emit(&mut work, scope, Child::Morsel(morsel)); self.queue.insert(work); } Output::Morsel(MorselOutput::Batch(array)) => { @@ -477,6 +530,25 @@ impl Run { Ok(()) } + /// Spawns a child `work` emitted, or hands it off if `work` is a distributing root. + fn emit(&mut self, work: &mut Work, scope: WorkScope, child: Child) { + let distributes = work.path.len() == 1 + && self + .roots + .get(&work.root) + .is_some_and(|root| root.distributes); + if distributes { + self.events.push_back(Progress::Handoff(Handoff { + root: work.root, + scope, + child, + })); + } else { + let path = work.child_path(); + self.spawn(work.root, scope, child, path); + } + } + /// Registers the requests an item published with its root's session. fn submit(&mut self, work: &mut Work, batch: IoBatch) -> VortexResult<()> { if batch.is_empty() { @@ -540,6 +612,7 @@ impl Driver { loop { match run.advance()? { Progress::Batch(batch) => batches.push(batch), + Progress::Handoff(_) => vortex_bail!("a driver root cannot hand off children"), 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..7437fe92336 100644 --- a/vortex-scan/src/planning/driver/tests.rs +++ b/vortex-scan/src/planning/driver/tests.rs @@ -503,6 +503,9 @@ 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::Handoff(handoff) => { + format!("handoff {:?} {:?}", handoff.root, handoff.scope.rows) + } Progress::Waiting => { events.push("waiting".to_string()); return Ok(events); @@ -515,6 +518,59 @@ fn drain(run: &mut Run) -> VortexResult> { } } +/// A distributing root hands off the planners and morsels it emits, unvisited and in emission +/// order, and another run admits each as a root of its own. +#[test] +fn a_distributing_root_hands_off_its_children() -> VortexResult<()> { + let log = Log::default(); + let file = ScriptedPlanner::boxed( + "file", + vec![ + PlannerStep::Planner( + scope(0..1), + vec![one_batch_morsel(scope(0..1), 1), PlannerStep::Done], + ), + one_batch_morsel(scope(1..3), 2), + PlannerStep::Done, + ], + &log, + ); + let source = source(); + let mut distributing = Run::new(); + distributing.admit_distributing(file, scope(0..3), Arc::clone(&source) as Arc); + let mut handoffs = Vec::new(); + loop { + match distributing.advance()? { + Progress::Handoff(handoff) => handoffs.push(handoff), + Progress::RootDone(_) => break, + _ => vortex_bail!("expected a handoff or the root to finish"), + } + } + let scopes: Vec<_> = handoffs.iter().map(|h| h.scope.rows.clone()).collect(); + assert_eq!(scopes, vec![0..1, 1..3]); + assert_eq!(log.events(), vec!["compute file"; 3]); + + let mut run = Run::new(); + for handoff in handoffs { + run.admit_child( + handoff.child, + handoff.scope, + Arc::clone(&source) as Arc, + ); + } + assert_eq!( + drain(&mut run)?, + vec![ + "batch RootId(0) 0..1", + "done RootId(0)", + "batch RootId(1) 1..3", + "done RootId(1)", + "idle", + ] + ); + Ok(()) +} + /// Roots run side by side: a root with ready work finishes while another waits for its read, and /// each root reads through, and finally clears, its own session. #[test] diff --git a/vortex-scan/src/planning/mod.rs b/vortex-scan/src/planning/mod.rs index 3095712f161..3ac13f9c334 100644 --- a/vortex-scan/src/planning/mod.rs +++ b/vortex-scan/src/planning/mod.rs @@ -10,6 +10,7 @@ pub mod driver; pub mod morsel; pub mod next; pub mod planner; +pub mod pool; #[cfg(test)] mod tests; diff --git a/vortex-scan/src/planning/pool.rs b/vortex-scan/src/planning/pool.rs new file mode 100644 index 00000000000..5c8577ae124 --- /dev/null +++ b/vortex-scan/src/planning/pool.rs @@ -0,0 +1,519 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A pool of workers that run the units a single root planner hands out. +//! +//! The root, such as a file planner, is admitted with [`Run::admit_distributing`] to a run of its +//! own, and each step of it hands out at most one unit. Any worker may step it, and the units it +//! hands out wait in a shared queue until a spare worker claims one and admits it to its own +//! [`Run`]. Everything a unit emits stays on the run that claimed it, so a unit's planners and +//! morsels, and the arrays they decode, stay with one worker. +//! +//! A worker steps the root before its own work whenever fewer units wait than workers are +//! parked, so a unit is handed out as soon as some worker can take it, while the root goes on +//! to plan the next. Workers are spawned as units are handed out with no worker parked to take +//! them, up to a limit, and a worker claims more units while the ones it holds wait for IO. +//! +//! At most a fixed number of units are handed out and not yet consumed, so the root stops +//! planning while the consumer falls behind. + +use std::collections::BTreeMap; +use std::collections::VecDeque; +use std::future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use futures::Stream; +use futures::future::BoxFuture; +use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::IoService; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::planning::driver::Child; +use crate::planning::driver::Progress; +use crate::planning::driver::RootId; +use crate::planning::driver::Run; +use crate::planning::planner::Planner; +use crate::planning::planner::WorkScope; + +/// Runs a worker future, typically by spawning it on a runtime. +pub type Spawn = Arc) + Send + Sync>; + +/// Maps each batch to the pool's output, on the worker that produced it. +pub type MapFn = Arc VortexResult + Send + Sync>; + +/// How a [`Pool`] runs its units. +#[derive(Clone, Copy, Debug)] +pub struct PoolOptions { + /// The most workers the pool spawns. + pub max_workers: usize, + /// The most units one worker runs at once. + pub units_per_worker: usize, + /// Whether output follows the order the root handed out its units in. Otherwise each batch + /// is yielded as soon as it is produced. + pub ordered: bool, +} + +/// Runs the units a root planner hands out on a pool of workers. See the [module](self) docs. +pub struct Pool; + +impl Pool { + /// Admits `root` and returns the stream of its units' mapped batches. + /// + /// The first worker is spawned now. Within a unit, ordered output is sorted by the start of + /// each batch's scope, since a unit's morsels finish in whatever order their reads arrive. + /// Dropping the stream stops the workers. + pub fn stream( + root: Box, + scope: WorkScope, + io: Arc, + map: MapFn, + spawn: Spawn, + options: PoolOptions, + ) -> PoolStream { + let max_workers = options.max_workers.max(1); + let units_per_worker = options.units_per_worker.max(1); + let mut file = Run::new(); + file.admit_distributing(root, scope, io.session()); + let shared = Arc::new(Shared { + io, + map, + spawn, + max_workers, + units_per_worker, + max_in_flight: max_workers * units_per_worker, + state: Mutex::new(State { + file: Some(file), + file_waiting: false, + file_done: false, + pending: VecDeque::new(), + next_seq: 0, + in_flight: 0, + workers: 0, + live_workers: 0, + parked: Vec::new(), + stopped: false, + events: VecDeque::new(), + consumer: None, + }), + }); + shared.spawn_worker(&mut shared.state.lock()); + PoolStream { + shared, + ordered: options.ordered, + units: None, + next_unit: 0, + buffered: BTreeMap::new(), + ready: VecDeque::new(), + } + } +} + +/// A unit handed out by the root and not yet claimed. +struct Unit { + /// The unit's position in the order the root handed out its units. + seq: u64, + scope: WorkScope, + child: Child, +} + +/// What workers report to the stream. +enum Event { + /// A mapped batch of a unit, with the start of the scope of the morsel that produced it. + Batch { + seq: u64, + start: u64, + value: T, + }, + /// Every batch of the unit has been reported. + UnitDone(u64), + /// The root has finished, having handed out this many units. + RootDone(u64), + Error(VortexError), +} + +struct Shared { + io: Arc, + map: MapFn, + spawn: Spawn, + max_workers: usize, + units_per_worker: usize, + max_in_flight: usize, + state: Mutex>, +} + +struct State { + /// The root's run, while no worker steps or polls it. + file: Option, + /// Whether the root waits for IO. + file_waiting: bool, + file_done: bool, + pending: VecDeque, + next_seq: u64, + /// Units handed out and not yet consumed. + in_flight: usize, + /// Workers spawned, and those that have not exited. + workers: usize, + live_workers: usize, + /// Wakers of parked workers. + parked: Vec, + /// Set by an error or by dropping the stream. + stopped: bool, + events: VecDeque>, + consumer: Option, +} + +impl State { + /// Whether the root's run can be stepped now. + fn file_ready(&self, max_in_flight: usize) -> bool { + self.file.is_some() && !self.file_waiting && self.in_flight < max_in_flight + } + + fn wake_workers(&mut self) { + for waker in self.parked.drain(..) { + waker.wake(); + } + } + + fn report(&mut self, event: Event) { + self.events.push_back(event); + if let Some(consumer) = self.consumer.take() { + consumer.wake(); + } + } +} + +impl Shared { + fn spawn_worker(self: &Arc, state: &mut State) { + state.workers += 1; + state.live_workers += 1; + let shared = Arc::clone(self); + (self.spawn)(Box::pin(async move { + let exit = WorkerExit(Arc::clone(&shared)); + if let Err(error) = Worker::new(&shared).run().await { + shared.fail(error); + } + drop(exit); + })); + } + + fn fail(&self, error: VortexError) { + let mut state = self.state.lock(); + if !state.stopped { + state.stopped = true; + state.report(Event::Error(error)); + } + state.wake_workers(); + } + + /// Takes the root's run to step it: whenever it is ready if `always`, and otherwise only + /// while fewer units wait than workers are parked. + fn take_file(&self, always: bool) -> Option { + let mut state = self.state.lock(); + if state.stopped || !state.file_ready(self.max_in_flight) { + return None; + } + if !always && state.pending.len() > state.parked.len() { + return None; + } + state.file.take() + } + + /// Steps the root's run until it hands out a unit, waits, or finishes. + fn step_file(self: &Arc, mut file: Run) -> VortexResult<()> { + let progress = file.advance(); + let mut state = self.state.lock(); + match progress? { + Progress::Handoff(handoff) => { + let seq = state.next_seq; + state.next_seq += 1; + state.in_flight += 1; + state.pending.push_back(Unit { + seq, + scope: handoff.scope, + child: handoff.child, + }); + state.file = Some(file); + if state.pending.len() > state.parked.len() && state.workers < self.max_workers { + self.spawn_worker(&mut state); + } + } + Progress::Waiting => { + state.file_waiting = true; + state.file = Some(file); + } + Progress::RootDone(_) | Progress::Idle => { + state.file_done = true; + let units = state.next_seq; + state.report(Event::RootDone(units)); + } + Progress::Batch(_) => vortex_bail!("the pool's root produced a batch"), + } + state.wake_workers(); + Ok(()) + } + + fn claim(&self) -> Option { + let mut state = self.state.lock(); + if state.stopped { + return None; + } + state.pending.pop_front() + } +} + +impl Shared { + fn release(&self, units: usize) { + let mut state = self.state.lock(); + state.in_flight -= units; + state.wake_workers(); + } +} + +/// Reports a worker's exit, including by panic, so the stream ends once every worker has. +struct WorkerExit(Arc>); + +impl Drop for WorkerExit { + fn drop(&mut self) { + let mut state = self.0.state.lock(); + state.live_workers -= 1; + if let Some(consumer) = state.consumer.take() { + consumer.wake(); + } + } +} + +/// One worker's run, holding the units it claimed. +struct Worker<'a, T> { + shared: &'a Arc>, + run: Run, + seqs: HashMap, +} + +impl<'a, T: Send + 'static> Worker<'a, T> { + fn new(shared: &'a Arc>) -> Self { + Self { + shared, + run: Run::new(), + seqs: HashMap::default(), + } + } + + async fn run(mut self) -> VortexResult<()> { + loop { + if self.shared.state.lock().stopped { + return Ok(()); + } + if let Some(file) = self.shared.take_file(false) { + self.shared.step_file(file)?; + continue; + } + if self.run.live_roots() > 0 { + match self.run.advance()? { + Progress::Batch(batch) => { + let seq = self.seq(batch.root)?; + let value = (self.shared.map)(batch.array)?; + self.shared.state.lock().report(Event::Batch { + seq, + start: batch.scope.rows.start, + value, + }); + continue; + } + Progress::RootDone(root) => { + let seq = self.seq(root)?; + self.seqs.remove(&root); + self.shared.state.lock().report(Event::UnitDone(seq)); + continue; + } + Progress::Handoff(_) => vortex_bail!("a unit cannot hand off children"), + Progress::Waiting | Progress::Idle => {} + } + } + if self.run.live_roots() < self.shared.units_per_worker + && let Some(unit) = self.shared.claim() + { + let root = self + .run + .admit_child(unit.child, unit.scope, self.shared.io.session()); + self.seqs.insert(root, unit.seq); + continue; + } + if let Some(file) = self.shared.take_file(true) { + self.shared.step_file(file)?; + continue; + } + if self.run.live_roots() == 0 { + let state = self.shared.state.lock(); + if state.file_done && state.pending.is_empty() { + return Ok(()); + } + } + future::poll_fn(|cx| self.park(cx)).await?; + } + } + + fn seq(&self, root: RootId) -> VortexResult { + self.seqs + .get(&root) + .copied() + .ok_or_else(|| vortex_err!("no unit for {root:?}")) + } + + /// Waits until there is something to do: a unit to claim, the root to step, a delivery to + /// the worker's own run or to the root's, or the pool to stop or finish. + fn park(&mut self, cx: &mut Context<'_>) -> Poll> { + let poll_file = { + let mut state = self.shared.state.lock(); + let can_claim = + !state.pending.is_empty() && self.run.live_roots() < self.shared.units_per_worker; + let finished = state.file_done && state.pending.is_empty(); + if state.stopped + || can_claim + || state.file_ready(self.shared.max_in_flight) + || (finished && self.run.live_roots() == 0) + { + return Poll::Ready(Ok(())); + } + if !state.parked.iter().any(|waker| waker.will_wake(cx.waker())) { + state.parked.push(cx.waker().clone()); + } + if state.file_waiting { + state.file.take() + } else { + None + } + }; + if let Some(mut file) = poll_file { + let delivered = file.poll_completion(cx); + let mut state = self.shared.state.lock(); + state.file = Some(file); + if delivered.is_ready() { + state.file_waiting = false; + state.wake_workers(); + return delivered; + } + } + if self.run.live_roots() > 0 { + return self.run.poll_completion(cx); + } + Poll::Pending + } +} + +/// The mapped batches of a [`Pool`]'s units. +pub struct PoolStream { + shared: Arc>, + ordered: bool, + /// How many units the root handed out, once it has finished. + units: Option, + /// Units fully consumed, in order, when ordered; otherwise units finished. + next_unit: u64, + /// Batches of units not yet yielded, and whether each unit is done. + buffered: BTreeMap, bool)>, + ready: VecDeque, +} + +impl PoolStream { + fn finished(&self) -> bool { + self.units == Some(self.next_unit) + } + + fn handle(&mut self, event: Event) -> VortexResult<()> { + match event { + Event::Batch { seq, start, value } => { + if self.ordered { + self.buffered.entry(seq).or_default().0.push((start, value)); + } else { + self.ready.push_back(value); + } + } + Event::UnitDone(seq) => { + if self.ordered { + self.buffered.entry(seq).or_default().1 = true; + let mut released = 0; + while let Some(entry) = self.buffered.first_entry() + && *entry.key() == self.next_unit + && entry.get().1 + { + let (mut batches, _) = entry.remove(); + batches.sort_by_key(|(start, _)| *start); + self.ready + .extend(batches.into_iter().map(|(_, value)| value)); + self.next_unit += 1; + released += 1; + } + if released > 0 { + self.shared.release(released); + } + } else { + self.next_unit += 1; + self.shared.release(1); + } + } + Event::RootDone(units) => self.units = Some(units), + Event::Error(error) => return Err(error), + } + Ok(()) + } +} + +// Nothing in the stream is pinned. +impl Unpin for PoolStream {} + +impl Stream for PoolStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + loop { + if let Some(value) = this.ready.pop_front() { + return Poll::Ready(Some(Ok(value))); + } + if this.finished() { + return Poll::Ready(None); + } + let events = { + let mut state = this.shared.state.lock(); + if state.events.is_empty() { + if state.stopped { + return Poll::Ready(None); + } + if state.live_workers == 0 { + state.stopped = true; + return Poll::Ready(Some(Err(vortex_err!( + "every pool worker exited before the scan finished" + )))); + } + state.consumer = Some(cx.waker().clone()); + return Poll::Pending; + } + std::mem::take(&mut state.events) + }; + for event in events { + if let Err(error) = this.handle(event) { + return Poll::Ready(Some(Err(error))); + } + } + } + } +} + +impl Drop for PoolStream { + fn drop(&mut self) { + let mut state = self.shared.state.lock(); + state.stopped = true; + state.wake_workers(); + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-scan/src/planning/pool/tests.rs b/vortex-scan/src/planning/pool/tests.rs new file mode 100644 index 00000000000..ef1400224f2 --- /dev/null +++ b/vortex-scan/src/planning/pool/tests.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::future::Future; +use std::pin::pin; +use std::sync::Arc; +use std::task::Wake; +use std::thread; +use std::thread::Thread; + +use futures::StreamExt; +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; + +use super::*; +use crate::planning::morsel::MorselOutput; +use crate::planning::tests::scripted::Log; +use crate::planning::tests::scripted::MorselStep; +use crate::planning::tests::scripted::PlannerStep; +use crate::planning::tests::scripted::RecordingIoSource; +use crate::planning::tests::scripted::ScriptedPlanner; +use crate::planning::tests::scripted::array_of; +use crate::planning::tests::scripted::scope; + +struct ThreadWaker(Thread); + +impl Wake for ThreadWaker { + fn wake(self: Arc) { + self.0.unpark(); + } +} + +fn block_on(future: F) -> F::Output { + let mut future = pin!(future); + let waker = Waker::from(Arc::new(ThreadWaker(thread::current()))); + let mut cx = Context::from_waker(&waker); + loop { + if let Poll::Ready(output) = future.as_mut().poll(&mut cx) { + return output; + } + thread::park(); + } +} + +/// Runs each worker on a thread of its own. +fn threads() -> Spawn { + Arc::new(|worker| { + thread::spawn(move || block_on(worker)); + }) +} + +/// Gives every session a source of its own, with the bytes at offset 0. +struct Sessions; + +impl IoService for Sessions { + fn session(&self) -> Arc { + let source = RecordingIoSource::default(); + source.canned_bytes(0, buffer![0u8, 1, 2, 3].into_byte_buffer()); + Arc::new(source) + } +} + +fn lengths() -> MapFn { + Arc::new(|array: ArrayRef| Ok(array.len())) +} + +/// A unit that reads a segment, then emits a morsel with `rows` rows over `start..start + 1`. +fn unit(start: u64, rows: u64) -> PlannerStep { + PlannerStep::Planner( + scope(start..start + 1), + vec![ + PlannerStep::Io(vec![IoRequest { + intent: IoIntent::Fetch, + request: IoRequestId(0), + target: IoTarget::range(0, 4), + }]), + PlannerStep::Morsel( + scope(start..start + 1), + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(rows))), + MorselStep::Compute(MorselOutput::Done), + ], + ), + PlannerStep::Done, + ], + ) +} + +fn run(root: Vec, options: PoolOptions) -> Vec> { + let log = Log::default(); + let stream = Pool::stream( + ScriptedPlanner::boxed("file", root, &log), + scope(0..0), + Arc::new(Sessions), + lengths(), + threads(), + options, + ); + block_on(stream.collect()) +} + +fn options(max_workers: usize, ordered: bool) -> PoolOptions { + PoolOptions { + max_workers, + units_per_worker: 2, + ordered, + } +} + +#[rstest] +#[case::one_worker(1)] +#[case::several_workers(3)] +fn ordered_output_follows_hand_out_order(#[case] workers: usize) -> VortexResult<()> { + let mut root: Vec<_> = (1..=8).map(|rows| unit(rows, rows)).collect(); + root.push(PlannerStep::Done); + let lengths = run(root, options(workers, true)) + .into_iter() + .collect::>>()?; + assert_eq!(lengths, (1..=8).collect::>()); + Ok(()) +} + +#[test] +fn unordered_output_yields_every_batch() -> VortexResult<()> { + let mut root: Vec<_> = (1..=8).map(|rows| unit(rows, rows)).collect(); + root.push(PlannerStep::Done); + let mut lengths = run(root, options(3, false)) + .into_iter() + .collect::>>()?; + lengths.sort_unstable(); + assert_eq!(lengths, (1..=8).collect::>()); + Ok(()) +} + +#[test] +fn a_root_with_no_units_ends_the_stream() { + assert!(run(vec![PlannerStep::Done], options(2, true)).is_empty()); +} + +#[test] +fn a_failing_unit_ends_the_stream_with_its_error() { + let root = vec![ + unit(0, 1), + PlannerStep::Planner(scope(1..2), vec![PlannerStep::Wait]), + PlannerStep::Done, + ]; + let outputs = run(root, options(2, true)); + let error = outputs + .iter() + .find_map(|output| output.as_ref().err()) + .map(|error| error.to_string()); + assert!( + error + .as_deref() + .is_some_and(|message| message.contains("no outstanding fetch")), + "{error:?}" + ); +} diff --git a/vortex-sqllogictest/slt/tpch/generate_data.sh b/vortex-sqllogictest/slt/tpch/generate_data.sh index 29a0beda69a..9779d3d083c 100755 --- a/vortex-sqllogictest/slt/tpch/generate_data.sh +++ b/vortex-sqllogictest/slt/tpch/generate_data.sh @@ -20,7 +20,8 @@ DATA_DIR="${SCRIPT_DIR}/data" mkdir -p "${DATA_DIR}" echo "Generating TPC-H data (SF=0.1)..." -uvx tpchgen-cli -s 0.1 --format=parquet --parquet-compression='ZSTD(3)' --output-dir "${DATA_DIR}/" +# Pinned: tpchgen-cli 4.0.0 replaced --format with a subcommand, breaking the unpinned call. +uvx tpchgen-cli@4.0.0 parquet -s 0.1 --compression='ZSTD(3)' --output-dir "${DATA_DIR}/" for f in "${DATA_DIR}"/*.parquet; do echo "Converting $(basename "$f") to Vortex..."