From 8d3a7d9ba5c81691782704ba6b3e98f6c356ff34 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 18:12:06 +0000 Subject: [PATCH 01/14] docs(scan): handoff for a minimal planner-driven decode slice Describe how to land the vortex-scan planning driver running develop's pipeline::Scan decode path, with IO through the planning protocol, as a small PR stack extracted from 42b29cc459. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01P1oQtzTgSzkvrGDKisoNko Signed-off-by: Joe Isaacs Signed-off-by: Claude --- vortex-scan/design/planner-decode-slice.md | 194 +++++++++++++++++++++ 1 file changed, 194 insertions(+) create mode 100644 vortex-scan/design/planner-decode-slice.md diff --git a/vortex-scan/design/planner-decode-slice.md b/vortex-scan/design/planner-decode-slice.md new file mode 100644 index 00000000000..6c3cb157a92 --- /dev/null +++ b/vortex-scan/design/planner-decode-slice.md @@ -0,0 +1,194 @@ + + + +# Handoff: minimal planner-driven decode slice + +Status: handoff for a new implementation session. Verified against `origin/develop` at +`18aebeaac8` (which includes #10462) and source commit `42b29cc459`. + +## Goal + +Make the `vortex-scan` planning driver (`Run` / `Planner` / `Morsel`) execute develop's sans-IO +`pipeline::Scan` decode/read path, with all IO going through the planning protocol +(`IoRequest` / `Completion`), as the first mergeable step towards planner-driven V2 scans. + +Source material: + +- `42b29cc459`, PR [#10464](https://github.com/vortex-data/vortex/pull/10464) "feat(scan): + planner-driven V2 scans on the physical plan pipelines" (+11367 lines). Not an ancestor of + develop; its merge base is `1a9ed685d1`. +- Full-change reference: PR [#10466](https://github.com/vortex-data/vortex/pull/10466), branch + `origin/ji/vibrant-sagan-2w7pov`. +- PR [#10462](https://github.com/vortex-data/vortex/pull/10462) (Zoned pruning, `Share`, + `ValueCache`) is **already merged into develop** as `18aebeaac8`. +- PR [#10463](https://github.com/vortex-data/vortex/pull/10463) (`3c37da1a2a`, coalescing IO + service, `IoIntent::Forget`, `vortex_io::request::trace`) is **not** on develop and not needed. + +## Scope + +In scope: driver, a protocol adapter from `pipeline::Scan` to `IoRequest`s, a projection-only +planner, segment IO via the existing `SegmentSource`, and an opt-in hook behind +`VORTEX_SCAN_V2=1` for scans with no filter and no list columns. + +Non-goals (later PRs): filter evaluation (`FilterPlanner`, conjunct ordering), zone pruning in +the planner, IO announcements and prefetch, the coalescing IO service, plan caching, repeated +scans, DuckDB integration, DataFusion diagnostics, and any performance work. + +## What develop already has + +| Piece | Location on develop | +| --- | --- | +| `WorkScope`, `State`, `PlannerOutput`, `Planner` | `vortex-scan/src/planning/planner.rs` (81 lines) | +| `Morsel`, `MorselOutput` | `vortex-scan/src/planning/morsel.rs` (54) | +| `Next`, `next_fn` | `vortex-scan/src/planning/next.rs` (74); module root `planning/mod.rs` (11) | +| `IoRequest`, `IoBatch`, `IoIntent::{Fetch,Prefetch,Announce}`, `IoTarget`, `IoResult`, `Completion`, `IoSource`, `IoService`, `IoConsumer` | `vortex-io/src/request/mod.rs` (223); no `Forget`, no `trace` module | +| `pipeline::Scan` with `try_new`, `with_value_cache` (L712), `step() -> Turn::{Read,Output,Waiting,Done}` (L718), `deliver(ReadId, BufferHandle)` (L755) | `vortex-layout/src/plan/physical/pipeline/scan.rs`; `Turn` at L82 | +| `ValueCache` | `vortex-layout/src/plan/physical/pipeline/values.rs` L22 | +| `lower(&LayoutRef) -> PlanRef`, public, handles Flat/Chunked/Struct/Dict/List/Zoned/LegacyStats | `vortex-layout/src/plan/physical/lower.rs` L50 | +| `optimize(PlanRef)` | `vortex-layout/src/plan/physical/optimize.rs` L28 | +| `plan_row_idx_expression` | `vortex-layout/src/plan/physical/plans/row_idx.rs` L135 | +| `attempt_split_ranges` | `vortex-layout/src/scan/splits.rs` L33 | + +## What lives only in `42b29cc459` + +Driver (`vortex-scan`): + +- `src/planning/driver/mod.rs` (644): `RootId` L56, `Batch` L59, `Progress` L69, `Run` L194 with + `new` L222, `admit` L241, `cancel` L262, `advance` L282, `complete` L331, `poll_completion` + L372; blocking `Driver` L600-644 (`run` L624). +- `src/planning/driver/trace.rs` (75), `driver/tests.rs` (814), `planning/tests/mod.rs` (6), + `planning/tests/scripted.rs` (418). +- `planner.rs` +12 / `morsel.rs` +6: only `trace_name()` / `trace_selection()` default methods, + used solely by tracing. +- `Cargo.toml` +6: `vortex-utils` dependency; dev-deps `parking_lot`, `rstest`, + `vortex-array` (`_test-harness`). + +Hidden dependency on #10463: `driver/mod.rs` L41 imports `vortex_io::request::trace::timestamp_ns` +and L254 calls `IoSource::trace_id()`; `trace.rs` imports `trace::{next_id, timestamp_ns}`. + +## PR plan + +Extract each file with `git show 42b29cc459: > ` and then trim as described. + +### PR A: driver in `vortex-scan` (~600 + ~1.2k tests) + +Standalone on develop. + +1. Copy `vortex-scan/src/planning/driver/{mod.rs,tests.rs}`, + `vortex-scan/src/planning/tests/{mod.rs,scripted.rs}`; add `pub mod driver;` and + `#[cfg(test)] mod tests;` to `planning/mod.rs`. +2. Strip tracing: drop `trace.rs`, the `trace` field on `Run` (L195) and its setup in `new` + (L224), every `tracing::debug!(ts_ns = timestamp_ns(), ...)` block (admit, cancel, completion, + selection ~L499-530), and the `timestamp_ns` import. About 60 lines. Do not add + `trace_name`/`trace_selection` to `Planner`/`Morsel`. + Alternative: cherry-pick the ~40 lines of `vortex-io` from `3c37da1a2a` + (`request/trace.rs` +24, `request/mod.rs` `trace_id()`); not recommended for this slice. +3. Add the `Cargo.toml` deps listed above. +4. Tests: `driver/tests.rs` and `scripted.rs` exercise the driver with scripted planners and a + fake `IoSource`; drop any test that asserts trace output. + +### PR B: protocol adapter in `vortex-layout` (~700 + ~300 tests) + +Depends on A. New `vortex-layout/src/scan/planning/` and `vortex-layout/src/scan/v2/` modules. + +- `scan/planning/graph.rs` (257): `ScanGraph` wraps `pipeline::Scan`; `ProtocolGraph` maps + `Turn::Read` to a `Fetch` `IoRequest` (L207-216) and delivers `IoResult`s back. Keep + `.with_value_cache(values)` (L73) and the `ValueCache` import (L27): they compile on develop + now that #10462 landed. +- `scan/planning/projection.rs` (568): keep `ProjectionPlanner` with `new` and `cut` (L104) and + `ProjectionMorsel` (L251-~362). Trim `compute` (L202) to: take `selected`, `cut`, pop a split, + emit `PlannerOutput::Morsel`. Drop `prefetch` (L70), `forget_pruned` (L131, needs + `IoIntent::Forget` from #10463), the `VORTEX_SCAN_IO_FORGET_PRUNED_PROJECTION` switch, the + `plan_segments` import from `v2::prefetch`, and the tests at L382+ that cover them. +- `scan/planning/mod.rs` (171): keep `SegmentLocation` (L44) and `ScanPlans` (L66, **including** + `values: ValueCache`). Move `SelectedRows` here from `filter.rs` L62. Drop `announce`, + `filter`, `plan_split`/`plan_sparse_projection` (they wrap `AnnouncePlanner`); replace with a + `plan_projection(plans, scope, mask) -> Box` returning `ProjectionPlanner`. +- `scan/v2/io.rs` (321): keep `segment_ranges`, `SegmentScanIo`, `SegmentIoSource` (L36-212), + adapting `SegmentSource` to `IoService`. Delete the `IoIntent::Forget` branch (L111). Keep the + two tests (L213+), rewriting any that build `Forget` requests. +- `scan/v2/split.rs` (90): keep `SplitPlan::run`, the `admit` / `advance` / + `poll_completion` loop (~L76-85). Drop `EARLY_ANNOUNCE` (L44-60). +- `scan/v2/splits.rs` (343): keep `projection_splits` (L252), `chunk_starts` (L135) / + `collect_starts`, `layout_chunk_starts` (L54) / `collect_layout_starts`, `cut_at` (~110 lines). + Drop `max_split_rows`, `MAX_SPARSE_ROWS`, the filter/sparse boundary functions. +- `scan/v2/lower.rs` (207): **drop entirely**. Develop's `plan::physical::lower` now handles + Zoned/LegacyStats, lowering them to their data when unoptimized. +- Add `pub(crate) mod planning; pub mod v2;` to `scan/mod.rs`. +- Tests: an in-memory `SegmentSource` plus layouts built like + `vortex-layout/src/plan/physical/pipeline/tests.rs` (`lower(&layout)`), driven by + `Driver::run` from PR A; compare with `assert_arrays_eq!` against the default scan. + +### PR C: opt-in hook (~550 + tests) + +- `scan/v2/mod.rs` (113): `ENV_VAR = "VORTEX_SCAN_V2"`, `enabled()`, `supports()` (no list + columns), `ScanFile { layout, locations, segments }`. Drop `io` and `plans` fields and the + `conjuncts`, `file`, `prefetch`, `pruning`, `repeated_scan`, `share` modules. Gate + `supports_scan` on `builder.filter.is_none()` (scans with a filter use the default executor). +- `scan/v2/scan_builder.rs` (333): `ScanBuilder::from_default` and `into_stream`; replace + `prepare()` / `RepeatedScanV2` with a small local prepare: lower the projection with + `plan::physical::lower`, `optimize`, compute splits via `chunk_starts` + `projection_splits`, + build one `SplitPlan` per split. +- `scan/v2/stream.rs` (141): `into_stream` and `LazyScanStream`; drop the + `pruning::file_pruning_enabled` import and branch. +- `vortex-layout/src/scan/scan_builder.rs`: add `into_parts()` and `ScanParts` (+44 in the + source commit; drop `has_filter_and_limit`, which is subsumed by the no-filter gate). +- `vortex-file/src/planning/mod.rs`: only `segment_locations` and `scan_file(&VortexFile)` + (source L51-78, ~30 lines). It must be in `vortex-file` because + `Footer::segment_specs_with_metadata` (`vortex-file/src/footer/mod.rs` L210) is `pub(crate)`. + Use `VortexFile::segment_source()` (`vortex-file/src/file.rs` L151); drop `file.scan_io()` and + `scan_file_cache` (both from #10463 / plan caching). +- Tests: `v2/tests.rs` cases without filters, plus `vortex-file` round trips like the + `v2::into_stream(file.scan()?, scan_file(&file))` test in `vortex-file/src/open.rs` (~L571-608 + in the source commit). + +### PR D (optional): first caller + +- `vortex-datafusion/src/persistent/opener.rs` on develop, L454-468: keep the builder chain, + then `if scan::v2::enabled() && scan::v2::supports_scan(&b, &file) { v2::into_stream(...) } + else { b.into_stream() }` (~+15). Skip `diagnostics.rs` and `benchmark_segment_cache`. +- Non-engine alternative: next to `VortexFile::scan()` (`vortex-file/src/file.rs` L217). + +### Later + +1. `FilterPlanner` (`scan/planning/filter.rs` 511 + `filter/tests.rs` 190) and + `scan/v2/conjuncts.rs` (226); lift the no-filter gate. +2. Zone pruning in the planner (`scan/v2/pruning.rs` 275; `ZonedPlan` is already on develop). +3. #10463 coalescing service, `FileScanIo` (`vortex-file/src/segments/{io_service,scan_io}.rs`), + `IoIntent::Forget`, tracing. +4. `announce.rs` (409), `prefetch.rs` (145), `share.rs` (147), split heuristics (sparse, + `max_split_rows`), `file.rs` plan caching (102), `repeated_scan.rs` (606). +5. `vortex-file/src/planning/{footer_open,footer_prune,tests}` (~1.8k incl. fixtures). +6. DuckDB (`file_prefetch.rs`, `ffi.rs`, C++ reader), DataFusion diagnostics. + +## Size + +Slice: ~3.3k lines including tests (~1.85k non-test). Full stack: ~16.3k (#10462 +1622, +#10463 +3301, #10464 +11367); #10462 is already merged. + +## Verification + +Run only when requested, scoped per crate (`vortex-scan`, `vortex-layout`, `vortex-file`, +`vortex-datafusion`): + +```bash +cargo clippy -p --all-targets --all-features -- -D warnings +cargo +nightly-2026-09-10 fmt -p # NIGHTLY_TOOLCHAIN in .github/workflows/ci.yml +cargo nextest run -p +``` + +Then a parity check with `VORTEX_SCAN_V2=1` on an unfiltered DataFusion query. + +## Open risks + +- **CI regression, cause unknown.** On CI's 96-core ARM runners the full V2 stack regressed + (Vortex queries up to 12x, FineWeb vortex-compact 1.25-1.8x, TPC-H SF10 q06 +28%) but not on + 4-core local runs. Per-split IO locks and the file read driver were ruled out. This slice is a + good base to bisect it on CI. +- **Local q01 +10%:** extra decimal to Arrow conversion and bit-packed unpack versus V1's + filter-on-compressed path. Expect the slice to be slower than V1 until filters land. +- `SegmentScanIo` issues one `SegmentSource` request per segment with no coalescing. +- Develop's `plan/physical` changed since the source commit's merge base (`1a9ed685d1`, ~2.2k + lines via #10462). `Scan::try_new(session, plan, splits)` is unchanged, but recheck imports + such as `pipeline::Split` and `ValueCache` when porting `graph.rs`. +- List columns cannot compile before their offsets are read; such files stay on V1. From 90aab885b4c8da8a5b3782a7146ca24dadb404c9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 18:52:14 +0000 Subject: [PATCH 02/14] feat(scan): planner-driven V2 executor for projection scans behind one fork Port the planning driver (`Run`/`Driver`) into vortex-scan and a projection-only V2 executor into vortex-layout that runs develop's pipeline `Scan` through the planning protocol, with every segment read issued as an `IoRequest` served from the file's `SegmentSource`. `ScanBuilder`'s lazy prepare step is the single place a streaming scan picks its executor: `v2::select` returns V1 unless `VORTEX_SCAN_V2=1`, the builder carries a `ScanFile`, and the scan has no filter, no limit and no list layout. `VortexFile::scan()` attaches the `ScanFile`; no caller changes. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017cUKVPYqfN7SLcXetK59i6 --- Cargo.lock | 3 + vortex-file/src/file.rs | 30 +- vortex-layout/src/scan/mod.rs | 2 + vortex-layout/src/scan/planning/graph.rs | 249 ++++++ vortex-layout/src/scan/planning/mod.rs | 86 ++ vortex-layout/src/scan/planning/projection.rs | 195 +++++ vortex-layout/src/scan/scan_builder.rs | 47 +- vortex-layout/src/scan/v2/io.rs | 315 +++++++ vortex-layout/src/scan/v2/mod.rs | 248 ++++++ vortex-layout/src/scan/v2/split.rs | 47 + vortex-layout/src/scan/v2/splits.rs | 86 ++ vortex-layout/src/scan/v2/tests.rs | 339 ++++++++ vortex-scan/Cargo.toml | 6 + vortex-scan/src/planning/driver/mod.rs | 551 ++++++++++++ vortex-scan/src/planning/driver/tests.rs | 813 ++++++++++++++++++ vortex-scan/src/planning/mod.rs | 4 + vortex-scan/src/planning/tests/mod.rs | 6 + vortex-scan/src/planning/tests/scripted.rs | 418 +++++++++ 18 files changed, 3430 insertions(+), 15 deletions(-) create mode 100644 vortex-layout/src/scan/planning/graph.rs create mode 100644 vortex-layout/src/scan/planning/mod.rs create mode 100644 vortex-layout/src/scan/planning/projection.rs create mode 100644 vortex-layout/src/scan/v2/io.rs create mode 100644 vortex-layout/src/scan/v2/mod.rs create mode 100644 vortex-layout/src/scan/v2/split.rs create mode 100644 vortex-layout/src/scan/v2/splits.rs create mode 100644 vortex-layout/src/scan/v2/tests.rs create mode 100644 vortex-scan/src/planning/driver/mod.rs create mode 100644 vortex-scan/src/planning/driver/tests.rs create mode 100644 vortex-scan/src/planning/tests/mod.rs create mode 100644 vortex-scan/src/planning/tests/scripted.rs diff --git a/Cargo.lock b/Cargo.lock index 009891d5e94..1d5c3a887f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12378,7 +12378,9 @@ version = "0.1.0" dependencies = [ "async-trait", "futures", + "parking_lot", "roaring", + "rstest", "tracing", "vortex-array", "vortex-buffer", @@ -12386,6 +12388,7 @@ dependencies = [ "vortex-io", "vortex-mask", "vortex-session", + "vortex-utils", ] [[package]] diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 72a5e8787e0..50aa797c4a5 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -21,6 +21,8 @@ use vortex_layout::LayoutReader; use vortex_layout::scan::layout::LayoutReaderDataSource; use vortex_layout::scan::scan_builder::ScanBuilder; use vortex_layout::scan::split_by::SplitBy; +use vortex_layout::scan::v2::ScanFile; +use vortex_layout::scan::v2::SegmentLocation; use vortex_layout::segments::SegmentSource; use vortex_scan::DataSourceRef; use vortex_session::VortexSession; @@ -215,10 +217,30 @@ impl VortexFile { /// Initiate a scan of the file, returning a builder for projection, filtering, selection, and /// execution options. pub fn scan(&self) -> VortexResult> { - Ok(ScanBuilder::new( - self.session.clone(), - self.layout_reader()?, - )) + Ok( + ScanBuilder::new(self.session.clone(), self.layout_reader()?) + .with_scan_file(self.scan_file()), + ) + } + + /// The parts of this file the planner-driven scan executor reads through: its root layout, + /// where each of its segments lives, and its segment source. + pub fn scan_file(&self) -> ScanFile { + let locations = self + .footer + .segment_specs_with_metadata() + .iter() + .map(|spec| SegmentLocation { + offset: spec.offset, + length: spec.length, + alignment: spec.alignment, + }) + .collect(); + ScanFile { + layout: Arc::clone(self.footer.layout()), + locations, + segments: Arc::clone(&self.segment_source), + } } /// Returns `true` if file-level statistics prove the expression cannot diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index e08e4cb3f31..c8a02ed74d9 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -6,6 +6,7 @@ mod filter; pub mod layout; mod limit; pub mod multi; +pub(crate) mod planning; pub mod repeated_scan; pub mod scan_builder; pub mod split_by; @@ -13,6 +14,7 @@ mod splits; mod tasks; #[cfg(test)] mod test; +pub mod v2; /// A heuristic for an ideal split size. /// diff --git a/vortex-layout/src/scan/planning/graph.rs b/vortex-layout/src/scan/planning/graph.rs new file mode 100644 index 00000000000..679fc6c6e44 --- /dev/null +++ b/vortex-layout/src/scan/planning/graph.rs @@ -0,0 +1,249 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use rustc_hash::FxHashMap; +use vortex_array::ArrayRef; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_io::request; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +use vortex_scan::planning::planner::State; +use vortex_session::VortexSession; + +use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::ReadId; +use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::Turn; +use crate::plan::physical::pipeline::ValueCache; +use crate::scan::planning::SegmentLocation; +use crate::segments::SegmentId; + +/// Rows a running plan produced, in row order. +#[derive(Clone)] +pub(crate) struct Piece { + /// Orders the piece among the pieces of one graph. + pub(crate) rows: Range, + /// The selected rows the piece holds. + pub(crate) array: ArrayRef, +} + +/// A plan running over one row range on the pipeline [`Scan`]. +/// +/// The scan produces the selected rows in row order and leaves its reads to its owner. +pub(crate) struct ScanGraph { + scan: Scan, + /// A step the scan already took that the owner has not seen. + pending: Option, + done: bool, + /// Orders the pieces handed out. + next_piece: u64, +} + +/// What a graph wants from its owner. +enum GraphState { + NeedsCompute, + Waiting, + Done, +} + +impl ScanGraph { + /// Compiles nothing yet: the scan compiles `plan` over `rows`, restricted to `mask`, on its + /// first step. `row_offset` is the global row index of the plan's first row, and `values` the + /// whole values shared with the split's other scans. + pub(crate) fn try_new( + session: VortexSession, + plan: &PlanRef, + rows: Range, + mask: Mask, + row_offset: u64, + values: ValueCache, + ) -> VortexResult { + let scan = Scan::try_new(session, plan.clone(), vec![Split { rows, mask }])? + .with_row_offset(row_offset) + .with_value_cache(values); + Ok(Self { + scan, + pending: None, + done: false, + next_piece: 0, + }) + } + + fn state(&self) -> GraphState { + match &self.pending { + _ if self.done => GraphState::Done, + Some(Turn::Waiting) => GraphState::Waiting, + _ => GraphState::NeedsCompute, + } + } + + fn step(&mut self) -> VortexResult { + let turn = match self.pending.take() { + Some(Turn::Waiting) | None => self.scan.step()?, + Some(turn) => turn, + }; + match &turn { + Turn::Done => self.done = true, + Turn::Waiting => self.pending = Some(Turn::Waiting), + Turn::Read(_) | Turn::Output(..) => {} + } + Ok(turn) + } + + fn piece(&mut self, array: ArrayRef) -> Piece { + let start = self.next_piece; + self.next_piece += 1; + Piece { + rows: start..start + 1, + array, + } + } + + fn deliver(&mut self, read: ReadId, bytes: BufferHandle) -> VortexResult<()> { + self.scan.deliver(read, bytes)?; + if matches!(self.pending, Some(Turn::Waiting)) { + self.pending = None; + } + Ok(()) + } +} + +/// What one [`ProtocolGraph::compute`] produced. +pub(crate) enum GraphStep { + /// Control returns to the owner with CPU work still ready. + Yield, + /// Requests to publish, each returned once. + NeedsIO(IoBatch), + /// Rows the plan produced. + Piece(Piece), +} + +/// A [`ScanGraph`] whose segment reads are published as protocol requests. +/// +/// The scan names segments by id; the protocol names byte ranges. Each scan read becomes a +/// `Fetch` of the segment's range at the segment's alignment, and each delivery is handed back to +/// the scan under its own id. +pub(crate) struct ProtocolGraph { + graph: ScanGraph, + locations: Arc<[SegmentLocation]>, + /// Published and undelivered requests, with the scan read and segment each one answers. + outstanding: FxHashMap, + /// A delivery the scan could not take, reported by the next compute. + failed: Option, + /// The protocol id the next published request gets. + next_id: u32, +} + +impl ProtocolGraph { + /// Wraps `graph`, numbering its requests from `first_id`. Request ids must be unique over an + /// owner's lifetime. + pub(crate) fn new(graph: ScanGraph, locations: Arc<[SegmentLocation]>, first_id: u32) -> Self { + Self { + graph, + locations, + outstanding: FxHashMap::default(), + failed: None, + next_id: first_id, + } + } + + /// The owner's protocol state while the graph is running. + pub(crate) fn state(&self) -> State { + if self.failed.is_some() { + return State::NeedsCompute; + } + match self.graph.state() { + GraphState::Done => State::Done, + GraphState::NeedsCompute => State::NeedsCompute, + GraphState::Waiting => State::Waiting, + } + } + + pub(crate) fn compute(&mut self) -> VortexResult { + if let Some(error) = self.failed.take() { + return Err(error); + } + Ok(match self.graph.step()? { + Turn::Waiting | Turn::Done => GraphStep::Yield, + Turn::Output(_, array) => GraphStep::Piece(self.graph.piece(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)?]; + loop { + match self.graph.step()? { + Turn::Read(read) => batch.push(self.request(read.id, read.segment_id)?), + turn => { + if !matches!(turn, Turn::Waiting) { + self.graph.pending = Some(turn); + self.graph.done = false; + } + break; + } + } + } + GraphStep::NeedsIO(batch) + } + }) + } + + fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { + let location = self.location(segment_id)?; + let id = self.next_id; + self.next_id = self + .next_id + .checked_add(1) + .ok_or_else(|| vortex_err!("ProtocolGraph ran out of request ids"))?; + self.outstanding.insert(id, (read, segment_id)); + Ok(request::IoRequest { + intent: IoIntent::Fetch, + request: request::IoRequestId(id), + target: location.target(), + }) + } + + /// Hands a delivered read to the scan. The driver only delivers requests this graph + /// published, so anything else is a driver bug. Bytes the scan cannot take fail the owner's + /// next compute. + pub(crate) fn set_io_result(&mut self, id: request::IoRequestId, result: IoResult) { + let Some((read, segment_id)) = self.outstanding.remove(&id.0) else { + vortex_panic!("ProtocolGraph: delivery of {id:?}, which is not outstanding"); + }; + let IoResult::Bytes(bytes) = result else { + vortex_panic!("ProtocolGraph: segment {segment_id} answered with a size"); + }; + if let Err(error) = self.deliver(read, segment_id, bytes) + && self.failed.is_none() + { + self.failed = Some(error.with_context(format!("delivering segment {segment_id}"))); + } + } + + fn deliver( + &mut self, + read: ReadId, + segment_id: SegmentId, + bytes: BufferHandle, + ) -> VortexResult<()> { + // A source that honours the target's alignment makes this a check, not a copy. + let alignment = self.location(segment_id)?.alignment; + let bytes = BufferHandle::new_host(bytes.try_into_host_sync()?.aligned(alignment)); + self.graph.deliver(read, bytes) + } + + fn location(&self, id: SegmentId) -> VortexResult { + self.locations + .get(*id as usize) + .copied() + .ok_or_else(|| vortex_err!("segment {id} has no known location")) + } +} diff --git a/vortex-layout/src/scan/planning/mod.rs b/vortex-layout/src/scan/planning/mod.rs new file mode 100644 index 00000000000..b632e55e139 --- /dev/null +++ b/vortex-layout/src/scan/planning/mod.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Layout-level planners and morsels for the planning protocol. +//! +//! A split is evaluated over physical plans: a [`ProjectionPlanner`] emits a [`ProjectionMorsel`] +//! per projection split, which runs the projection plan over the split's selected rows through +//! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published +//! as a protocol request, so the planning driver decides when IO happens. + +pub(crate) mod graph; +mod projection; + +use std::sync::Arc; + +pub use projection::ProjectionPlanner; +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 crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::ValueCache; + +/// Where a segment's bytes live in its source. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct SegmentLocation { + /// Byte offset of the segment from the start of the source. + pub offset: u64, + /// Length of the segment in bytes. + pub length: u32, + /// Alignment the segment's bytes must have once delivered. + pub alignment: Alignment, +} + +impl SegmentLocation { + /// The byte range a read of the segment asks for. + pub fn target(&self) -> IoTarget { + IoTarget::Range { + offset: self.offset, + len: self.length as usize, + alignment: self.alignment, + } + } +} + +/// What every split of a scan shares: the projection plan and where its segments live. +#[derive(Clone)] +pub struct ScanPlans { + /// The session used for decoding and expression evaluation. + pub session: VortexSession, + /// The location of every segment, indexed by segment id. + pub locations: Arc<[SegmentLocation]>, + /// The optimized projection plan, over the scan's root row domain. + pub projection: PlanRef, + /// Where the chunks the projection reads start, over the root row domain, sorted. A split's + /// projection is cut at each of these, so every projection split reads at most one chunk of + /// every column. + pub projection_starts: Arc<[u64]>, + /// The global row index of the root row domain's first row. + pub row_offset: u64, + /// Whole values, such as dictionary values, kept between the scans of one split. Each split + /// starts with its own, dropped when the split finishes. + pub values: ValueCache, +} + +/// The rows of a split selected for projection. +pub struct SelectedRows { + /// The split the rows belong to. + pub scope: WorkScope, + /// Which rows of `scope.rows` are selected. + pub mask: Mask, +} + +/// The work for one split: project the rows of `scope` selected by `mask`. +pub fn plan_projection(plans: ScanPlans, scope: WorkScope, mask: Mask) -> Box { + Box::new(ProjectionPlanner::new( + ScanPlans { + values: ValueCache::default(), + ..plans + }, + SelectedRows { scope, mask }, + )) +} diff --git a/vortex-layout/src/scan/planning/projection.rs b/vortex-layout/src/scan/planning/projection.rs new file mode 100644 index 00000000000..e230a3789a6 --- /dev/null +++ b/vortex-layout/src/scan/planning/projection.rs @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_scan::planning::morsel::Morsel; +use vortex_scan::planning::morsel::MorselOutput; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; + +use crate::scan::planning::ScanPlans; +use crate::scan::planning::SelectedRows; +use crate::scan::planning::graph::GraphStep; +use crate::scan::planning::graph::Piece; +use crate::scan::planning::graph::ProtocolGraph; +use crate::scan::planning::graph::ScanGraph; +use crate::scan::v2::splits::projection_splits; + +/// Turns a split's selected rows into the morsels that project them. +/// +/// The rows are cut into projection splits where the projection's chunks start, and each split +/// with a selected row becomes one [`ProjectionMorsel`], so a morsel reads at most one chunk of +/// every column. The planner finishes once it has handed out every morsel. +pub struct ProjectionPlanner { + plans: ScanPlans, + /// The split's selected rows, until the first compute cuts them. + selected: Option, + /// The projection splits not yet handed out, last first. + pending: Vec, +} + +impl ProjectionPlanner { + /// Creates a planner for `selected`. + pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + Self { + plans, + selected: Some(selected), + pending: Vec::new(), + } + } + + /// The projection splits of `selected` that have a selected row, last first. + fn cut(&self, selected: SelectedRows) -> Vec { + let SelectedRows { scope, mask } = selected; + let start = scope.rows.start; + let index = |row: u64| usize::try_from(row - start).vortex_expect("split row fits usize"); + let mut splits = projection_splits(&self.plans.projection_starts, scope.rows.clone()) + .into_iter() + .filter_map(|rows| { + let mask = mask.slice(index(rows.start)..index(rows.end)); + (!mask.all_false()).then(|| SelectedRows { + scope: WorkScope { + file_ordinal: scope.file_ordinal, + rows, + }, + mask, + }) + }) + .collect::>(); + splits.reverse(); + splits + } +} + +impl IoConsumer for ProjectionPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for ProjectionPlanner { + fn state(&self) -> State { + if self.selected.is_none() && self.pending.is_empty() { + State::Done + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + if let Some(selected) = self.selected.take() { + self.pending = self.cut(selected); + } + // No projection split has a selected row. + let Some(selected) = self.pending.pop() else { + return Ok(PlannerOutput::Done); + }; + let scope = selected.scope.clone(); + Ok(PlannerOutput::Morsel( + scope, + Box::new(ProjectionMorsel::new(self.plans.clone(), selected)), + )) + } +} + +/// Runs the projection plan over a split's selected rows and returns them as one batch, in row +/// order. A selection that projects to no rows finishes without a batch. +pub struct ProjectionMorsel { + plans: ScanPlans, + selected: SelectedRows, + graph: Option, + pieces: Vec, + done: bool, +} + +impl ProjectionMorsel { + /// Creates a morsel projecting `selected`. + pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + Self { + plans, + selected, + graph: None, + pieces: Vec::new(), + done: false, + } + } + + fn finish(&mut self) -> VortexResult { + self.done = true; + self.pieces.sort_by_key(|piece| piece.rows.start); + let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) + .into_iter() + .map(|piece| piece.array) + .filter(|array| !array.is_empty()) + .collect(); + let array = match arrays.len() { + 0 => return Ok(MorselOutput::Done), + 1 => arrays.remove(0), + _ => ChunkedArray::try_new(arrays, self.plans.projection.dtype().clone())?.into_array(), + }; + Ok(MorselOutput::Batch(array)) + } +} + +impl IoConsumer for ProjectionMorsel { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some(graph) = self.graph.as_mut() { + graph.set_io_result(request, result); + } + } +} + +impl Morsel for ProjectionMorsel { + fn state(&self) -> State { + match &self.graph { + _ if self.done => State::Done, + None => State::NeedsCompute, + Some(graph) => match graph.state() { + State::Done => State::NeedsCompute, + state => state, + }, + } + } + + fn compute(&mut self) -> VortexResult { + if self.done { + vortex_bail!("ProjectionMorsel: compute called after Done"); + } + let Some(graph) = self.graph.as_mut() else { + let graph = ScanGraph::try_new( + self.plans.session.clone(), + &self.plans.projection, + self.selected.scope.rows.clone(), + self.selected.mask.clone(), + self.plans.row_offset, + self.plans.values.clone(), + )?; + self.graph = Some(ProtocolGraph::new( + graph, + Arc::clone(&self.plans.locations), + 0, + )); + return Ok(MorselOutput::Continue); + }; + if graph.state() == State::Done { + return self.finish(); + } + Ok(match graph.compute()? { + GraphStep::Yield => MorselOutput::Continue, + GraphStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), + GraphStep::Piece(piece) => { + self.pieces.push(piece); + MorselOutput::Continue + } + }) + } +} diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 814c0759453..47b4d80a80d 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -42,6 +42,9 @@ use crate::scan::repeated_scan::RepeatedScan; use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; use crate::scan::splits::attempt_split_ranges; +use crate::scan::v2; +use crate::scan::v2::Executor; +use crate::scan::v2::ScanFile; /// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs. /// @@ -54,32 +57,34 @@ use crate::scan::splits::attempt_split_ranges; /// Projection and filter expressions must be bound against the reader dtype. Work is divided by /// the configured [`SplitBy`] strategy or by explicit selection ranges. pub struct ScanBuilder { - session: VortexSession, + pub(super) session: VortexSession, layout_reader: LayoutReaderRef, - projection: BoundExpression, - filter: Option, + pub(super) projection: BoundExpression, + pub(super) filter: Option, /// Whether the scan needs to return splits in the order they appear in the file. - ordered: bool, + pub(super) ordered: bool, /// Optionally read a subset of the rows in the file. - row_range: Option>, + pub(super) row_range: Option>, /// The selection mask to apply to the selected row range. // TODO(joe): replace this is usage of row_id selection, see - selection: Selection, + pub(super) selection: Selection, /// How to split the file for concurrent processing. split_by: SplitBy, /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare) /// uses them instead of walking the layout. natural_splits: Option>, /// The number of splits to make progress on concurrently **per-thread**. - concurrency: usize, + pub(super) concurrency: usize, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. - map_fn: Arc VortexResult + Send + Sync>, + pub(super) map_fn: Arc VortexResult + Send + Sync>, metrics_registry: Option>, /// Maximal number of rows to read after filtering. - limit: Option, + pub(super) limit: Option, /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression, /// but not by the scan [`Selection`] which remains relative. - row_offset: u64, + pub(super) row_offset: u64, + /// Where the scanned file's segments live, which lets the planner-driven executor read it. + pub(super) scan_file: Option, } impl ScanBuilder { @@ -103,6 +108,7 @@ impl ScanBuilder { metrics_registry: None, limit: None, row_offset: 0, + scan_file: None, } } @@ -266,6 +272,14 @@ impl ScanBuilder { self } + /// Attach the segment map of the file the layout reader was opened over. + /// + /// Scans with a segment map may run on the planner-driven executor, see [`v2`](crate::scan::v2). + pub fn with_scan_file(mut self, file: ScanFile) -> Self { + self.scan_file = Some(file); + self + } + /// The [`DType`] returned by the scan, after applying the projection. pub fn dtype(&self) -> VortexResult { Ok(self.projection.dtype().clone()) @@ -296,6 +310,7 @@ impl ScanBuilder { metrics_registry: self.metrics_registry, limit: self.limit, row_offset: self.row_offset, + scan_file: self.scan_file, map_fn: Arc::new(move |a| old_map_fn(a).and_then(&map_fn)), } } @@ -363,6 +378,16 @@ impl ScanBuilder { )) } + /// Prepares the scan and builds its stream on the executor [`v2::select`] picks. + /// + /// This is the only place a streaming scan chooses between executors. + fn prepare_stream(self) -> VortexResult>> { + match v2::select(self)? { + Executor::V1(builder) => builder.prepare()?.execute_stream(None), + Executor::V2(scan) => scan.execute_stream(), + } + } + /// Constructs a task per row split of the scan, returned as a vector of futures. /// /// Scans with row limits must use [`Self::into_stream`] or [`Self::into_iter`] so the scan @@ -426,7 +451,7 @@ impl Stream for LazyScanStream { // task: poll_next must only wait for and poll an already-constructed stream. // This also keeps construction errors on the Preparing -> Error path rather // than running construction on the caller's executor. - let task = handle.spawn_cpu(move || builder.prepare()?.execute_stream(None)); + let task = handle.spawn_cpu(move || builder.prepare_stream()); self.state = LazyScanState::Preparing(task); } LazyScanState::Preparing(task) => match ready!(Pin::new(task).poll(cx)) { diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs new file mode 100644 index 00000000000..66c4ee6d4a7 --- /dev/null +++ b/vortex-layout/src/scan/v2/io.rs @@ -0,0 +1,315 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::VecDeque; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use futures::FutureExt; +use futures::StreamExt; +use futures::TryFutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; +use futures::stream::FuturesUnordered; +use parking_lot::Mutex; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoResult; +use vortex_io::request::IoService; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::scan::planning::SegmentLocation; +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()) +} + +/// Serves splits from any [`SegmentSource`], by segment id. +pub(super) struct SegmentScanIo { + segments: Arc, + ranges: SegmentRanges, +} + +impl SegmentScanIo { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self { segments, ranges } + } +} + +impl IoService for SegmentScanIo { + fn session(&self) -> Arc { + Arc::new(SegmentIoSource::new( + Arc::clone(&self.segments), + Arc::clone(&self.ranges), + )) + } +} + +/// 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, +/// [`poll`](IoSource::poll) takes a finished one if there is any, and the split's future awaits +/// the next with [`poll_completion`](IoSource::poll_completion) when the driver has nothing else +/// to run. [`wait`](IoSource::wait) is never called. +/// +/// An announcement requests the segment and holds the request, never polled, until a prefetch or +/// fetch consumes it, so a source that coalesces requests can fold it into a nearby read. +pub(super) struct SegmentIoSource { + segments: Arc, + ranges: SegmentRanges, + reads: Mutex>>, + /// Reads started by prefetches, consumed by the next fetch of their segment. + prefetched: Mutex>, + /// Announced segments, requested and never polled. + announced: Mutex>>, +} + +/// A read started ahead of the fetch that needs it. +type Prefetched = Shared>>>; + +impl SegmentIoSource { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self { + segments, + ranges, + reads: Mutex::new(FuturesUnordered::new()), + prefetched: Mutex::default(), + announced: Mutex::default(), + } + } +} + +impl IoSource for SegmentIoSource { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + let reads = self.reads.lock(); + let mut prefetched = self.prefetched.lock(); + for request in batch { + let IoTarget::Range { offset, len, .. } = request.target else { + vortex_bail!( + "SegmentIoSource only serves byte ranges, not {:?}", + request.target + ); + }; + let segment = *self + .ranges + .get(&(offset, len)) + .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; + if request.intent == IoIntent::Announce { + self.announced + .lock() + .entry(segment) + .or_default() + .push_back(self.segments.request(segment)); + continue; + } + if request.intent == IoIntent::Prefetch { + prefetched.entry(segment).or_insert_with(|| { + let read = self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared(); + // Polling registers the read as wanted, so the source starts it now. + drop( + read.clone() + .poll_unpin(&mut Context::from_waker(Waker::noop())), + ); + read + }); + continue; + } + let bytes = match prefetched.remove(&segment) { + Some(read) => read + .map_err(move |err| { + vortex_err!("prefetched read of segment {segment} failed: {err}") + }) + .boxed(), + None => self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)), + }; + let id = request.request; + reads.push( + async move { + Completion { + owner, + request: id, + result: bytes.await.map(IoResult::Bytes), + } + } + .boxed(), + ); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + // The split's future polls again with its own waker before it waits, so a read that + // finishes after this check is not missed. + let mut cx = Context::from_waker(Waker::noop()); + match self.reads.lock().poll_next_unpin(&mut cx) { + Poll::Ready(completion) => Ok(completion), + Poll::Pending => Ok(None), + } + } + + fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + match self.reads.lock().poll_next_unpin(cx) { + Poll::Ready(Some(completion)) => Poll::Ready(Ok(completion)), + Poll::Ready(None) => Poll::Ready(Err(vortex_err!( + "The split's driver is waiting with no reads in flight" + ))), + Poll::Pending => Poll::Pending, + } + } + + fn wait(&self) -> VortexResult { + vortex_bail!("SegmentIoSource is awaited through poll_completion, not waited on") + } + + fn release(&self, _owner: IoOwnerId) {} + + fn clear(&self) { + self.reads.lock().clear(); + self.prefetched.lock().clear(); + self.announced.lock().clear(); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use futures::FutureExt; + use futures::future; + use vortex_array::buffer::BufferHandle; + use vortex_buffer::Alignment; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use vortex_io::request::IoIntent; + use vortex_io::request::IoOwnerId; + use vortex_io::request::IoRequest; + use vortex_io::request::IoRequestId; + use vortex_io::request::IoSource; + use vortex_io::request::IoTarget; + + use super::SegmentIoSource; + use super::segment_ranges; + use crate::scan::planning::SegmentLocation; + use crate::segments::SegmentFuture; + use crate::segments::SegmentId; + use crate::segments::SegmentSource; + + /// Serves four bytes per segment and counts the reads it is asked for. + #[derive(Default)] + struct CountingSegments(AtomicUsize); + + impl SegmentSource for CountingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.0.fetch_add(1, Ordering::Relaxed); + let bytes = ByteBuffer::from(vec![u8::try_from(*id).unwrap_or(u8::MAX); 4]); + future::ready(Ok(BufferHandle::new_host(bytes))).boxed() + } + } + + fn request(intent: IoIntent, id: u32) -> IoRequest { + IoRequest { + intent, + request: IoRequestId(id), + target: IoTarget::range(0, 4), + } + } + + /// A fetch of a prefetched segment waits on the prefetch's read rather than reading again, + /// and only the fetch is delivered. + #[tokio::test] + async fn fetch_reuses_a_prefetched_read() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + + io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 1)])?; + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(1)); + assert!(completion.result.is_ok()); + assert!(io.poll()?.is_none()); + + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 2)])?; + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(2)); + assert!(completion.result.is_ok()); + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + Ok(()) + } + + #[tokio::test] + async fn announcements_keep_independent_consumers() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + io.submit( + IoOwnerId(0), + vec![ + request(IoIntent::Announce, 0), + request(IoIntent::Announce, 1), + ], + )?; + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + for id in 2..5 { + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + assert!( + std::future::poll_fn(|cx| io.poll_completion(cx)) + .await? + .result + .is_ok() + ); + assert_eq!( + segments.0.load(Ordering::Relaxed), + if id < 4 { 2 } else { 3 } + ); + } + Ok(()) + } +} diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs new file mode 100644 index 00000000000..16003898523 --- /dev/null +++ b/vortex-layout/src/scan/v2/mod.rs @@ -0,0 +1,248 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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. +//! +//! [`select`] is the one place a streaming [`ScanBuilder`] picks between this executor and the +//! default one. It picks this one when `VORTEX_SCAN_V2=1`, the builder carries a [`ScanFile`], +//! and the scan only uses what this executor implements so far: no filter, no limit, and a layout +//! that lowers to plans with a pipeline implementation. Every other scan runs on the default +//! executor, unchanged. + +mod io; +mod split; +pub(crate) mod splits; +#[cfg(test)] +mod tests; + +use std::env; +use std::ops::Range; +use std::sync::Arc; +use std::sync::LazyLock; + +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use futures::stream::BoxStream; +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::request::IoService; +use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; +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::plan_row_idx_expression; +use crate::scan::planning::ScanPlans; +pub use crate::scan::planning::SegmentLocation; +use crate::scan::planning::plan_projection; +use crate::scan::scan_builder::ScanBuilder; +use crate::scan::v2::io::SegmentScanIo; +use crate::scan::v2::io::segment_ranges; +use crate::scan::v2::split::SplitPlan; +use crate::scan::v2::splits::chunk_starts; +use crate::scan::v2::splits::projection_splits; +use crate::segments::SegmentSource; + +/// Set to `1` to run supported scans on the planner-driven executor. +pub const ENV_VAR: &str = "VORTEX_SCAN_V2"; + +/// Whether supported scans run on the planner-driven executor. +pub fn enabled() -> bool { + static ENABLED: LazyLock = LazyLock::new(|| env::var(ENV_VAR).is_ok_and(|v| v == "1")); + *ENABLED +} + +/// The parts of a file the planner-driven executor reads through. +#[derive(Clone)] +pub struct ScanFile { + /// The root layout the scan builder's reader was opened over. + pub layout: LayoutRef, + /// The location of every segment, indexed by segment id. + pub locations: Arc<[SegmentLocation]>, + /// Fetches segment bytes by id. + pub segments: Arc, +} + +/// The executor a streaming scan runs on. +pub(crate) enum Executor { + /// The default executor, over the layout reader. + V1(ScanBuilder), + /// The planner-driven executor, prepared. + V2(PreparedScan), +} + +/// Picks the executor for `builder` and prepares the scan when it is this one. +pub(crate) fn select(builder: ScanBuilder) -> VortexResult> { + if !enabled() { + return Ok(Executor::V1(builder)); + } + select_supported(builder) +} + +/// Prepares `builder` on this executor when it supports the scan. +fn select_supported(builder: ScanBuilder) -> VortexResult> { + let declined = if builder.filter.is_some() { + Some("filter") + } else if builder.limit.is_some() { + Some("limit") + } else { + match &builder.scan_file { + None => Some("no scan file"), + Some(file) if !supports_layout(&file.layout)? => Some("layout"), + Some(_) => None, + } + }; + if let Some(reason) = declined { + tracing::debug!(target: "vortex_layout::scan::v2", reason, "scan runs on v1"); + return Ok(Executor::V1(builder)); + } + tracing::debug!(target: "vortex_layout::scan::v2", "scan runs on v2"); + Ok(Executor::V2(PreparedScan::try_new(builder)?)) +} + +/// Whether every layout under `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) -> VortexResult { + if !(layout.is::() + || layout.is::() + || layout.is::() + || layout.is::() + || layout.is::() + || layout.is::()) + { + return Ok(false); + } + for child in layout.children()? { + if !supports_layout(&child)? { + return Ok(false); + } + } + Ok(true) +} + +/// A scan planned over its file, ready to run its splits. +pub(crate) struct PreparedScan { + session: VortexSession, + plans: ScanPlans, + /// Serves the splits' reads. + io: Arc, + /// The scan's splits, over the root row domain, in row order. + splits: Vec>, + selection: Selection, + map_fn: Arc VortexResult + Send + Sync>, + ordered: bool, + concurrency: usize, +} + +impl PreparedScan { + /// Lowers the file's layout, plans the projection over it, and cuts the scan's row range + /// where the projection's chunks start. + fn try_new(builder: ScanBuilder) -> VortexResult { + let file = builder + .scan_file + .ok_or_else(|| vortex_err!("a V2 scan needs a scan file"))?; + let root = lower(&file.layout)?; + let projection = optimize(plan_row_idx_expression(builder.projection, root.clone())?)?; + let projection_starts: Arc<[u64]> = chunk_starts([&projection])?.into(); + let rows = builder.row_range.unwrap_or(0..root.row_count()); + let splits = projection_splits(&projection_starts, rows); + let io = Arc::new(SegmentScanIo::new( + file.segments, + segment_ranges(&file.locations), + )); + Ok(Self { + session: builder.session.clone(), + plans: ScanPlans { + session: builder.session, + locations: file.locations, + projection, + projection_starts, + row_offset: builder.row_offset, + values: Default::default(), + }, + io, + splits, + selection: builder.selection, + map_fn: builder.map_fn, + ordered: builder.ordered, + concurrency: builder.concurrency, + }) + } + + /// The splits with selected rows, each ready to be admitted to a planning run. + fn split_plans(&self) -> Vec { + self.splits + .iter() + .filter_map(|rows| { + let 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() + } + + /// 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. + pub(crate) 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::>>() + }) + }); + let batches = if self.ordered { + tasks.buffered(concurrency).boxed() + } else { + tasks.buffer_unordered(concurrency).boxed() + }; + Ok(batches + .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) + .try_flatten() + .boxed()) + } +} diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs new file mode 100644 index 00000000000..b9371735d39 --- /dev/null +++ b/vortex-layout/src/scan/v2/split.rs @@ -0,0 +1,47 @@ +// 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/splits.rs b/vortex-layout/src/scan/v2/splits.rs new file mode 100644 index 00000000000..7537dfc2534 --- /dev/null +++ b/vortex-layout/src/scan/v2/splits.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splits a scan from its plans rather than from the layout reader. +//! +//! A scan is cut where the chunks of the columns its projection reads start, so each split reads +//! at most one chunk of every column. + +use std::iter; +use std::ops::Range; + +use itertools::Itertools; +use vortex_error::VortexResult; + +use crate::plan::physical::Concat; +use crate::plan::physical::Eval; +use crate::plan::physical::Filter; +use crate::plan::physical::Pack; +use crate::plan::physical::PlanRef; +use crate::plan::physical::Take; + +/// The row positions where the chunks of `plans` start, over their shared row domain, sorted and +/// without duplicates. +/// +/// This follows plans whose children cover the same rows, or known parts of them, and stops at any +/// other plan: a take's values cover other rows. +pub(super) fn chunk_starts<'a>( + plans: impl IntoIterator, +) -> VortexResult> { + let mut starts = Vec::new(); + for plan in plans { + collect_starts(plan, 0, &mut starts)?; + } + starts.sort_unstable(); + starts.dedup(); + Ok(starts) +} + +fn collect_starts(plan: &PlanRef, offset: u64, starts: &mut Vec) -> VortexResult<()> { + if let Some(concat) = plan.as_opt::() { + for (index, &child_offset) in concat.row_offsets().iter().enumerate() { + starts.push(offset + child_offset); + collect_starts(&plan.child_required(index)?, offset + child_offset, starts)?; + } + } else if plan.is::() { + collect_starts(&plan.child_required(0)?, offset, starts)?; + } else if plan.is::() || plan.is::() || plan.is::() { + for child in plan.children().iter() { + collect_starts(&child?, offset, starts)?; + } + } + Ok(()) +} + +/// The projection splits of `rows`: cut at every one of `starts` inside it. +pub(crate) fn projection_splits(starts: &[u64], rows: Range) -> Vec> { + let first = starts.partition_point(|&start| start <= rows.start); + let last = starts.partition_point(|&start| start < rows.end); + iter::once(rows.start) + .chain(starts[first..last].iter().copied()) + .chain(iter::once(rows.end)) + .tuple_windows() + .map(|(start, end)| start..end) + .filter(|range| !range.is_empty()) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::whole(&[], 10..20, vec![10..20])] + #[case::cut(&[0, 12, 15, 30], 10..20, vec![10..12, 12..15, 15..20])] + #[case::start_on_boundary(&[10, 15], 10..20, vec![10..15, 15..20])] + #[case::empty(&[10], 5..5, vec![])] + fn projection( + #[case] starts: &[u64], + #[case] rows: Range, + #[case] expected: Vec>, + ) { + assert_eq!(projection_splits(starts, rows), expected); + } +} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs new file mode 100644 index 00000000000..9c71dc0f990 --- /dev/null +++ b/vortex-layout/src/scan/v2/tests.rs @@ -0,0 +1,339 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::num::NonZeroUsize; +use std::ops::Range; +use std::sync::Arc; + +use futures::TryStreamExt; +use futures::stream; +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::expr::pack; +use vortex_array::expr::root; +use vortex_array::validity::Validity; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_io::session::RuntimeSessionExt; +use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_session::VortexSession; + +use super::Executor; +use super::ScanFile; +use super::select_supported; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::chunked::writer::ChunkedLayoutStrategy; +use crate::layouts::flat::writer::FlatLayoutStrategy; +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::planning::SegmentLocation; +use crate::scan::scan_builder::ScanBuilder; +use crate::segments::SegmentSource; +use crate::segments::TestSegments; +use crate::sequence::SequenceId; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt as _; +use crate::test::new_session; + +const CHUNK_ROWS: usize = 1000; +const CHUNKS: usize = 4; + +/// A table chunk of `CHUNK_ROWS` rows starting at row `start`: an integer, a nullable string and a +/// list column. +fn table_chunk(start: usize) -> VortexResult { + let rows = start..start + CHUNK_ROWS; + let ints = Buffer::from_iter( + rows.clone() + .map(|row| i32::try_from(row).unwrap_or(i32::MAX)), + ); + let strings = VarBinArray::from_iter( + rows.clone() + .map(|row| (row % 3 != 0).then(|| format!("row {row}"))), + DType::Utf8(true.into()), + ); + let elements = Buffer::from_iter((0..2 * CHUNK_ROWS).map(|i| i64::try_from(i).unwrap_or(0))); + let offsets = Buffer::from_iter((0..=CHUNK_ROWS).map(|i| u32::try_from(2 * i).unwrap_or(0))); + let lists = ListArray::try_new( + elements.into_array(), + offsets.into_array(), + Validity::NonNullable, + )?; + Ok(StructArray::from_fields(&[ + ("ints", ints.into_array()), + ("strings", strings.into_array()), + ("lists", lists.into_array()), + ])? + .into_array()) +} + +/// Writes `CHUNKS` table chunks with `strategy`. +async fn write( + session: &VortexSession, + strategy: &dyn LayoutStrategy, +) -> VortexResult<(Arc, LayoutRef)> { + let segments = Arc::new(TestSegments::default()); + let (mut sequence_id, eof) = SequenceId::root().split(); + let chunks = (0..CHUNKS) + .map(|chunk| Ok((sequence_id.advance(), table_chunk(chunk * CHUNK_ROWS)?))) + .collect::>(); + let dtype = table_chunk(0)?.dtype().clone(); + let layout = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(chunks)).sendable(), + eof, + session, + ) + .await?; + Ok((segments, layout)) +} + +fn chunked_flat() -> Arc { + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())) +} + +/// Every column is shredded into its own chunked leaves. +fn table() -> Arc { + Arc::new(TableStrategy::new(chunked_flat(), chunked_flat())) +} + +/// Every column is shredded into chunked leaves with zone statistics. +fn zoned_table() -> Arc { + let zoned = ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: NonZeroUsize::new(CHUNK_ROWS).unwrap_or(NonZeroUsize::MIN), + ..Default::default() + }, + ); + Arc::new(TableStrategy::new(chunked_flat(), Arc::new(zoned))) +} + +/// A [`ScanFile`] over test segments. They have no byte offsets, and the scan fetches them by id, +/// so each location is a placeholder whose offset is its id. +fn scan_file(segments: &Arc, layout: &LayoutRef) -> VortexResult { + let mut count = 0; + for layout in layout.depth_first_traversal() { + for id in layout?.segment_ids() { + count = count.max(*id as usize + 1); + } + } + let locations = (0..count as u64) + .map(|offset| SegmentLocation { + offset, + length: 0, + alignment: Alignment::none(), + }) + .collect(); + Ok(ScanFile { + layout: Arc::clone(layout), + locations, + segments: Arc::clone(segments), + }) +} + +fn builder( + session: &VortexSession, + segments: &Arc, + layout: &LayoutRef, +) -> VortexResult> { + let reader = layout.new_reader( + "".into(), + Arc::clone(segments), + session, + &Default::default(), + )?; + Ok(ScanBuilder::new(session.clone(), reader).with_scan_file(scan_file(segments, layout)?)) +} + +async fn collect( + dtype: DType, + stream: impl futures::Stream>, +) -> VortexResult { + let chunks: Vec = stream.try_collect().await?; + Ok(ChunkedArray::try_new(chunks, 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>, +) -> VortexResult<()> { + let dtype = make()?.dtype()?; + let expected = collect(dtype.clone(), make()?.prepare()?.execute_stream(None)?).await?; + let Executor::V2(scan) = select_supported(make()?)? else { + panic!("the scan should run on the planner-driven executor"); + }; + let actual = collect(dtype, scan.execute_stream()?).await?; + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) +} + +#[derive(Clone, Copy, Debug)] +enum Strategy { + Flat, + Table, + ZonedTable, +} + +impl Strategy { + fn get(self) -> Arc { + match self { + Strategy::Flat => chunked_flat(), + Strategy::Table => table(), + Strategy::ZonedTable => zoned_table(), + } + } +} + +#[rstest] +#[case::whole(None, None)] +#[case::row_range(Some(250..3_100), None)] +#[case::empty_row_range(Some(10..10), None)] +#[case::row_indices(None, Some(&[0, 1, 999, 1000, 2500, 3999][..]))] +#[case::row_range_and_indices(Some(900..2600), Some(&[0, 999, 1000, 2500, 3999][..]))] +#[tokio::test] +async fn matches_v1( + #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, + #[case] row_range: Option>, + #[case] rows: Option<&'static [u64]>, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, strategy.get().as_ref()).await?; + assert_same_as_v1(&session, || { + let mut builder = builder(&session, &segments, &layout)?; + if let Some(row_range) = row_range.clone() { + builder = builder.with_row_range(row_range); + } + if let Some(rows) = rows { + builder = + builder.with_row_indices(StrictSortedBuffer::try_new(Buffer::copy_from(rows))?); + } + Ok(builder) + }) + .await +} + +#[rstest] +#[tokio::test] +async fn projection_matches_v1( + #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, strategy.get().as_ref()).await?; + let dtype = layout.dtype().clone(); + assert_same_as_v1(&session, || { + let projection = pack( + [ + ("strings", get_item("strings", root())), + ("idx", row_idx()), + ("big", gt(get_item("ints", root()), lit(1500i32))), + ], + false.into(), + ); + Ok(builder(&session, &segments, &layout)? + .with_projection(projection.bind(&dtype)?) + .with_row_offset(100)) + }) + .await +} + +#[tokio::test] +async fn filters_and_limits_run_on_v1() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let dtype = layout.dtype().clone(); + let filter = gt(get_item("ints", root()), lit(10i32)).bind(&dtype)?; + let filtered = builder(&session, &segments, &layout)?.with_filter(filter); + assert!(matches!(select_supported(filtered)?, Executor::V1(_))); + let limited = builder(&session, &segments, &layout)?.with_limit(10); + assert!(matches!(select_supported(limited)?, Executor::V1(_))); + Ok(()) +} + +#[tokio::test] +async fn scans_without_a_scan_file_run_on_v1() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let builder = ScanBuilder::new(session.clone(), reader); + assert!(matches!(select_supported(builder)?, Executor::V1(_))); + Ok(()) +} + +/// List layouts are scanned by the default executor: their element ranges are only known once +/// their offsets are read, after a pipeline would have been compiled. +#[tokio::test] +async fn list_layouts_run_on_v1() -> 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)?; + assert!(matches!(select_supported(builder)?, Executor::V1(_))); + Ok(()) +} + +#[tokio::test] +async fn unordered_returns_every_row() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let Executor::V2(scan) = + select_supported(builder(&session, &segments, &layout)?.with_ordered(false))? + else { + panic!("the scan should run on the planner-driven executor"); + }; + let chunks: Vec = scan.execute_stream()?.try_collect().await?; + let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); + assert_eq!(rows, CHUNKS * CHUNK_ROWS); + 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 Executor::V2(scan) = select_supported( + builder(&session, &segments, &layout)? + .with_projection(get_item("ints", root()).bind(&dtype)?) + .with_row_range(990..1010), + )? + else { + panic!("the scan should run on the planner-driven executor"); + }; + let actual = collect( + DType::Primitive(vortex_array::dtype::PType::I32, false.into()), + scan.execute_stream()?, + ) + .await?; + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter(990i32..1010), + &mut session.create_execution_ctx() + ); + Ok(()) +} diff --git a/vortex-scan/Cargo.toml b/vortex-scan/Cargo.toml index d706ec20e22..d12c207b049 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -21,11 +21,17 @@ vortex-error = { workspace = true } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } +vortex-utils = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } roaring = { workspace = true } tracing = { workspace = true } +[dev-dependencies] +parking_lot = { workspace = true } +rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } + [lints] workspace = true diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs new file mode 100644 index 00000000000..8cb78099272 --- /dev/null +++ b/vortex-scan/src/planning/driver/mod.rs @@ -0,0 +1,551 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A single-worker driver for the protocol. +//! +//! A [`Run`] drives any number of root planners, each with its own IO session, and hands their +//! batches to its caller as they are produced. Runnable items are prioritised in DFS preorder, +//! with roots ordered by admission and children by emission. Each visit takes at most one CPU +//! step. A parent's continuation follows its emitted child's entire subtree. Waiting work is +//! skipped and regains its priority when IO makes it runnable. Published IO is registered +//! immediately and state is inspected again before parking. Output follows emission order, +//! not global row order. An earlier runnable branch can starve later branches. +//! +//! A run never blocks: [`Run::advance`] returns [`Progress::Waiting`] when every live item waits +//! 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. + +use std::collections::BTreeMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_utils::aliases::hash_map::Entry; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::planning::morsel::Morsel; +use crate::planning::morsel::MorselOutput; +use crate::planning::planner::Planner; +use crate::planning::planner::PlannerOutput; +use crate::planning::planner::State; +use crate::planning::planner::WorkScope; + +/// Identifies one root admitted to a [`Run`]. Unique within the run, in admission order. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct RootId(pub u64); + +/// One array produced by a morsel, with the scope the morsel was created for. +pub struct Batch { + /// The root the morsel descends from. + pub root: RootId, + /// Scope of the morsel that produced the array. + pub scope: WorkScope, + /// The non-empty array. + pub array: ArrayRef, +} + +/// 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), + /// 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 + /// to [`Run::complete`]. + Waiting, + /// Nothing is live: every admitted root has finished. + Idle, +} + +enum Item { + Planner(Box), + Morsel(Box), +} + +/// What one item's `compute()` produced. +enum Output { + Planner(PlannerOutput), + Morsel(MorselOutput), +} + +/// One live planner or morsel. Boxed in the run queue and the parked map, so scheduling moves a +/// pointer. +struct Work { + id: IoOwnerId, + root: RootId, + scope: WorkScope, + item: Item, + /// Stable child-ordinal path, beginning with the root's admission ordinal. + path: Vec, + /// The parent continuation sorts after all previously emitted child subtrees. + next_child: u64, + /// Published `Fetch` requests not yet delivered, with the target each must be answered with. + outstanding: HashMap, +} + +impl Work { + fn priority(&self) -> Vec { + let mut priority = self.path.clone(); + priority.push(self.next_child); + priority + } + + // Child k takes the continuation's old key as its path. Its subtree then sorts before + // the parent's new continuation at k + 1, even if descendants are discovered after IO. + fn child_path(&mut self) -> Vec { + let path = self.priority(); + self.next_child += 1; + path + } + + fn state(&self) -> State { + match &self.item { + Item::Planner(planner) => planner.state(), + Item::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, + } + } + + /// Checks `result` against the outstanding request it answers and hands it to the item. + fn deliver(&mut self, request: IoRequestId, result: IoResult) -> VortexResult<()> { + let Some(target) = self.outstanding.remove(&request) else { + vortex_bail!( + "completion for {request:?} of {:?}, which is not outstanding", + self.id + ); + }; + if !result.matches(&target) { + vortex_bail!("IO source answered {target:?} with {}", result.kind()); + } + self.consumer().set_io_result(request, result); + Ok(()) + } +} + +/// Keeps scheduling order separate from lookup for deliveries to runnable work. +#[derive(Default)] +struct ReadyQueue { + priorities: BTreeMap, IoOwnerId>, + work: HashMap>, +} + +impl ReadyQueue { + fn insert(&mut self, work: Box) { + self.priorities.insert(work.priority(), work.id); + self.work.insert(work.id, work); + } + + fn pop(&mut self) -> Option> { + let (_, owner) = self.priorities.pop_first()?; + self.work.remove(&owner) + } + + fn is_empty(&self) -> bool { + self.priorities.is_empty() + } + + fn cancel(&mut self, root: RootId) { + self.work.retain(|_, work| work.root != root); + self.priorities + .retain(|_, owner| self.work.contains_key(owner)); + } +} + +/// Poll IO even when one high-priority branch keeps producing CPU work. +const IO_POLL_INTERVAL: usize = 64; + +/// One admitted root: the session its descendants read through, and how many of them are live. +struct Root { + io: Arc, + live: usize, +} + +/// A set of root planners advanced by its caller on one thread at a time. +/// +/// Dropping a run clears the IO session of every root still live, even when it ends in an error. +#[derive(Default)] +pub struct Run { + roots: HashMap, + next_root: u64, + next_owner: u64, + queue: ReadyQueue, + parked: HashMap>, + /// Batches and finished roots not yet returned by [`advance`](Self::advance). + events: VecDeque, + /// Visits left before the sources are polled for completions again. + visits_until_poll: usize, + step_limit: Option, + steps: usize, +} + +impl Drop for Run { + fn drop(&mut self) { + for root in self.roots.values() { + root.io.clear(); + } + } +} + +impl Run { + /// Creates a run with nothing admitted. + pub fn new() -> Self { + Self::default() + } + + /// Fails with an error naming the limit if more than `steps` visits are needed. Tests use it + /// to catch loops. + pub fn with_step_limit(mut self, steps: usize) -> Self { + self.step_limit = Some(steps); + self + } + + /// Admits `root` and everything it spawns, reading through `io`. + /// + /// The session belongs to the root: the run submits the requests of every descendant to it, + /// and clears it once they have all retired or the root is cancelled. + pub fn admit( + &mut self, + root: Box, + scope: WorkScope, + io: Arc, + ) -> 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]); + id + } + + /// Drops every live item descended from `root` and clears its session. Batches the root + /// already produced are still returned, followed by no [`Progress::RootDone`]. + pub fn cancel(&mut self, root: RootId) { + let Some(cancelled) = self.roots.remove(&root) else { + return; + }; + self.queue.cancel(root); + self.parked.retain(|_, work| work.root != root); + cancelled.io.clear(); + } + + /// The number of admitted roots that have not finished. + pub fn live_roots(&self) -> usize { + self.roots.len() + } + + /// Runs ready work, taking completions the sources already have, until there is something to + /// hand back or everything left waits for IO. Never blocks. + pub fn advance(&mut self) -> VortexResult { + loop { + if let Some(event) = self.events.pop_front() { + return Ok(event); + } + if self.visits_until_poll == 0 || self.queue.is_empty() { + self.take_completions()?; + self.visits_until_poll = IO_POLL_INTERVAL; + if self.queue.is_empty() { + return Ok(if self.parked.is_empty() { + Progress::Idle + } else { + Progress::Waiting + }); + } + } + self.visits_until_poll -= 1; + let Some(work) = self.queue.pop() else { + continue; + }; + self.steps += 1; + if let Some(limit) = self.step_limit + && self.steps > limit + { + vortex_bail!("driver exceeded its step limit of {limit}"); + } + self.visit(work)?; + } + } + + /// Delivers a completion taken from the session of one of the run's roots. + pub fn complete(&mut self, completion: Completion) -> VortexResult<()> { + let Completion { + owner, + request, + result, + } = completion; + let result = result?; + if let Entry::Occupied(mut parked) = self.parked.entry(owner) { + let work = parked.get_mut(); + work.deliver(request, result)?; + // An item that still waits for other requests stays parked. + if work.state() != State::Waiting { + let work = parked.remove(); + self.queue.insert(work); + } else if work.outstanding.is_empty() { + vortex_bail!( + "work for scope {:?} waits with no outstanding fetch", + work.scope + ); + } + return Ok(()); + } + // An item that published fetches and had CPU work left is still in the run queue. + let Some(work) = self.queue.work.get_mut(&owner) else { + vortex_bail!("completion for unknown work {owner:?}"); + }; + work.deliver(request, result) + } + + /// Waits for a completion from any root's session, registering `cx` to be woken by the next + /// one, and delivers every completion that is ready. + /// + /// Call it after [`advance`](Self::advance) returned [`Progress::Waiting`]. + pub fn poll_completion(&mut self, cx: &mut Context<'_>) -> Poll> { + let sources: Vec<_> = self + .roots + .values() + .map(|root| Arc::clone(&root.io)) + .collect(); + if sources.is_empty() { + return Poll::Ready(Err(vortex_err!( + "the run is waiting for IO with no live root" + ))); + } + let mut delivered = false; + for io in sources { + match io.poll_completion(cx) { + Poll::Ready(Ok(completion)) => { + if let Err(error) = self.complete(completion) { + return Poll::Ready(Err(error)); + } + delivered = true; + } + Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), + Poll::Pending => {} + } + } + if delivered { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + } + + /// Delivers the completions every root's session has ready, without blocking. + fn take_completions(&mut self) -> VortexResult<()> { + if self.roots.len() == 1 { + // The common case, without collecting the sessions. + let Some(io) = self.roots.values().next().map(|root| Arc::clone(&root.io)) else { + return Ok(()); + }; + while let Some(completion) = io.poll()? { + self.complete(completion)?; + } + return Ok(()); + } + let sources: Vec<_> = self + .roots + .values() + .map(|root| Arc::clone(&root.io)) + .collect(); + for io in sources { + while let Some(completion) = io.poll()? { + self.complete(completion)?; + } + } + Ok(()) + } + + fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec) { + let id = IoOwnerId(self.next_owner); + self.next_owner += 1; + if let Some(root) = self.roots.get_mut(&root) { + root.live += 1; + } + let work = Box::new(Work { + id, + root, + scope, + item, + path, + next_child: 0, + outstanding: HashMap::default(), + }); + self.queue.insert(work); + } + + /// Retires a finished item, and its root once nothing descended from it is live. + fn retire(&mut self, work: &Work) { + let Entry::Occupied(mut root) = self.roots.entry(work.root) else { + return; + }; + root.get().io.release(work.id); + root.get_mut().live -= 1; + if root.get().live == 0 { + root.remove().io.clear(); + self.events.push_back(Progress::RootDone(work.root)); + } + } + + /// Parks an item that waits for IO, which is a protocol error with no fetch outstanding. + fn park(&mut self, work: Box) -> VortexResult<()> { + if work.outstanding.is_empty() { + vortex_bail!( + "work for scope {:?} waits with no outstanding fetch", + work.scope + ); + } + self.parked.insert(work.id, work); + Ok(()) + } + + /// Requeues, parks, or retires an item according to the state it reports. + fn settle(&mut self, work: Box) -> VortexResult<()> { + match work.state() { + State::Done => self.retire(&work), + State::NeedsCompute => self.queue.insert(work), + State::Waiting => self.park(work)?, + } + Ok(()) + } + + /// 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() { + State::NeedsCompute => Output::Planner(planner.compute()?), + _ => return self.settle(work), + }, + Item::Morsel(morsel) => match morsel.state() { + State::NeedsCompute => Output::Morsel(morsel.compute()?), + _ => return self.settle(work), + }, + }; + match output { + Output::Planner(PlannerOutput::Done) | Output::Morsel(MorselOutput::Done) => { + self.retire(&work); + } + Output::Planner(PlannerOutput::Continue) | Output::Morsel(MorselOutput::Continue) => { + self.queue.insert(work); + } + Output::Planner(PlannerOutput::NeedsIO(batch)) + | Output::Morsel(MorselOutput::NeedsIO(batch)) => { + self.submit(&mut work, batch)?; + self.settle(work)?; + } + Output::Planner(PlannerOutput::Planner(scope, child)) => { + let path = work.child_path(); + self.spawn(work.root, scope, Item::Planner(child), path); + 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.queue.insert(work); + } + Output::Morsel(MorselOutput::Batch(array)) => { + if array.is_empty() { + vortex_bail!("morsel for scope {:?} returned an empty batch", work.scope); + } + self.events.push_back(Progress::Batch(Batch { + root: work.root, + scope: work.scope.clone(), + array, + })); + self.queue.insert(work); + } + } + Ok(()) + } + + /// 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() { + vortex_bail!("compute() published an empty batch"); + } + for request in &batch { + if request.intent == IoIntent::Fetch + && work + .outstanding + .insert(request.request, request.target) + .is_some() + { + vortex_bail!( + "fetch {:?} of {:?} was published while it is outstanding", + request.request, + work.id + ); + } + } + let Some(root) = self.roots.get(&work.root) else { + vortex_bail!("work {:?} outlived its root {:?}", work.id, work.root); + }; + root.io.submit(work.id, batch) + } +} + +/// Drives one root to completion on the calling thread, blocking while it waits for IO. +pub struct Driver { + io: Arc, + step_limit: Option, +} + +impl Driver { + /// Creates a driver that submits every request to `io`. + pub fn new(io: Arc) -> Self { + Self { + io, + step_limit: None, + } + } + + /// Fails with an error naming the limit if more than `steps` visits are needed. Tests use it + /// to catch loops. + pub fn with_step_limit(self, steps: usize) -> Self { + Self { + step_limit: Some(steps), + ..self + } + } + + /// Runs `root` and everything it spawns, returning batches in emission order. + pub fn run(&self, root: Box) -> VortexResult> { + let mut run = Run::new(); + run.step_limit = self.step_limit; + let scope = WorkScope { + file_ordinal: 0, + rows: 0..0, + }; + run.admit(root, scope, Arc::clone(&self.io)); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::RootDone(_) | Progress::Idle => return Ok(batches), + Progress::Waiting => run.complete(self.io.wait()?)?, + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs new file mode 100644 index 00000000000..389ad24372d --- /dev/null +++ b/vortex-scan/src/planning/driver/tests.rs @@ -0,0 +1,813 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Waker; + +use rstest::rstest; +use vortex_array::assert_arrays_eq; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_io::request::IoRequest; + +use super::*; +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::ctx; +use crate::planning::tests::scripted::scope; + +fn request(id: u32, offset: u64) -> IoRequest { + IoRequest { + intent: IoIntent::Fetch, + request: IoRequestId(id), + target: IoTarget::range(offset, 4), + } +} + +fn driver(source: &Arc) -> Driver { + Driver::new(Arc::clone(source) as Arc).with_step_limit(1_000) +} + +fn source() -> Arc { + let source = RecordingIoSource::default(); + source.canned_bytes(0, buffer![0u8, 1, 2, 3].into_byte_buffer()); + source.canned_bytes(4, buffer![4u8, 5, 6, 7].into_byte_buffer()); + Arc::new(source) +} + +fn lifo_source() -> Arc { + let source = RecordingIoSource::lifo(); + source.canned_bytes(0, buffer![0u8, 1, 2, 3].into_byte_buffer()); + source.canned_bytes(4, buffer![4u8, 5, 6, 7].into_byte_buffer()); + Arc::new(source) +} + +fn one_batch_morsel(scope: WorkScope, rows: u64) -> PlannerStep { + PlannerStep::Morsel( + scope, + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(rows))), + MorselStep::Compute(MorselOutput::Done), + ], + ) +} + +#[test] +fn root_done_produces_nothing() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed("root", vec![PlannerStep::Done], &log); + let batches = driver(&source()).run(root)?; + assert!(batches.is_empty()); + assert_eq!(log.events(), vec!["compute root"]); + Ok(()) +} + +#[test] +fn one_morsel_one_batch() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![one_batch_morsel(scope(10..20), 3), PlannerStep::Done], + &log, + ); + let batches = driver(&source()).run(root)?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].scope, scope(10..20)); + assert_arrays_eq!(batches[0].array, array_of(3), &mut ctx()); + Ok(()) +} + +#[test] +fn two_morsels_emit_in_order() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + one_batch_morsel(scope(0..1), 1), + one_batch_morsel(scope(1..3), 2), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source()).run(root)?; + let scopes: Vec<_> = batches.iter().map(|batch| batch.scope.clone()).collect(); + assert_eq!(scopes, vec![scope(0..1), scope(1..3)]); + assert_arrays_eq!(batches[0].array, array_of(1), &mut ctx()); + assert_arrays_eq!(batches[1].array, array_of(2), &mut ctx()); + Ok(()) +} + +#[test] +fn dfs_finishes_child_subtree_before_parent_continues() -> VortexResult<()> { + let log = Log::default(); + let child = PlannerStep::Planner( + scope(0..5), + vec![one_batch_morsel(scope(0..5), 1), PlannerStep::Done], + ); + let root = ScriptedPlanner::boxed( + "parent", + vec![ + child, + PlannerStep::Continue, + one_batch_morsel(scope(5..10), 2), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source()).run(root)?; + let scopes: Vec<_> = batches.iter().map(|batch| batch.scope.clone()).collect(); + assert_eq!(scopes, vec![scope(0..5), scope(5..10)]); + assert_eq!( + log.events(), + vec![ + "compute parent", + "compute child", + "compute morsel", + "compute morsel", + "compute child", + "compute parent", + "compute parent", + "compute morsel", + "compute morsel", + "compute parent", + ] + ); + Ok(()) +} + +#[test] +fn needs_io_delivers_every_request_in_order() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Io(vec![request(0, 4), request(1, 0)]), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source).run(root)?; + assert!(batches.is_empty()); + assert_eq!( + source.performed(), + vec![IoTarget::range(4, 4), IoTarget::range(0, 4)] + ); + assert_eq!( + log.events(), + vec![ + "publish root", + "deliver root 0 bytes", + "deliver root 1 bytes", + "compute root", + ] + ); + Ok(()) +} + +/// A fetch published while CPU work is left is performed once, and delivered when the planner +/// gets round to waiting for it. +#[test] +fn fetch_published_ahead_is_delivered_when_awaited() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::NeedsIO(vec![request(0, 0)]), + PlannerStep::Continue, + PlannerStep::Await(vec![IoRequestId(0)]), + PlannerStep::Done, + ], + &log, + ); + driver(&source).run(root)?; + assert_eq!(source.performed(), vec![IoTarget::range(0, 4)]); + assert_eq!( + log.events(), + vec![ + "compute root", + "compute root", + "deliver root 0 bytes", + "compute root", + ] + ); + Ok(()) +} + +#[test] +fn continue_requeues() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Continue, + PlannerStep::Continue, + PlannerStep::Done, + ], + &log, + ); + driver(&source()).run(root)?; + assert_eq!( + log.events() + .iter() + .filter(|event| event.as_str() == "compute root") + .count(), + 3 + ); + Ok(()) +} + +#[test] +fn io_failure_propagates_after_two_performs() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + source.fail(IoTarget::range(4, 4), "disk on fire"); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Io(vec![request(0, 0), request(1, 4)]), + PlannerStep::Done, + ], + &log, + ); + let err = driver(&source).run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref().is_some_and(|m| m.contains("disk on fire")), + "{err:?}" + ); + assert_eq!(source.performed().len(), 2); + Ok(()) +} + +#[test] +fn completions_out_of_order_within_one_batch() -> VortexResult<()> { + let log = Log::default(); + let source = lifo_source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Io(vec![request(0, 0), request(1, 4)]), + PlannerStep::Done, + ], + &log, + ); + driver(&source).run(root)?; + assert_eq!( + source.performed(), + vec![IoTarget::range(0, 4), IoTarget::range(4, 4)], + "submitted in request order" + ); + assert_eq!( + log.events(), + vec![ + "publish root", + "deliver root 1 bytes", + "deliver root 0 bytes", + "compute root", + ] + ); + Ok(()) +} + +#[test] +fn completions_out_of_order_across_parked_items() -> VortexResult<()> { + let log = Log::default(); + let source = lifo_source(); + let child = |id: u32, offset: u64| { + PlannerStep::Planner( + scope(0..1), + vec![ + PlannerStep::Io(vec![request(id, offset)]), + PlannerStep::Done, + ], + ) + }; + let root = ScriptedPlanner::boxed( + "root", + vec![child(5, 0), child(6, 4), PlannerStep::Done], + &log, + ); + driver(&source).run(root)?; + let deliveries: Vec<_> = log + .events() + .into_iter() + .filter(|event| event.starts_with("deliver")) + .collect(); + assert_eq!( + deliveries, + vec!["deliver child 6 bytes", "deliver child 5 bytes"], + "both children were parked before either completed" + ); + Ok(()) +} + +#[test] +fn lying_source_fails_before_delivery() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + source.answer_with_size(IoTarget::range(0, 4), 4); + let root = ScriptedPlanner::boxed( + "root", + vec![PlannerStep::Io(vec![request(0, 0)]), PlannerStep::Done], + &log, + ); + let err = driver(&source).run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref().is_some_and(|m| m.contains("answered")), + "{err:?}" + ); + assert_eq!(log.events(), vec!["publish root"]); + Ok(()) +} + +#[test] +fn morsel_needs_io_midway() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Morsel( + scope(0..4), + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(1))), + MorselStep::Io(vec![request(0, 0)]), + MorselStep::Compute(MorselOutput::Batch(array_of(2))), + MorselStep::Compute(MorselOutput::Done), + ], + ), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source()).run(root)?; + assert_eq!(batches.len(), 2); + assert_arrays_eq!(batches[0].array, array_of(1), &mut ctx()); + assert_arrays_eq!(batches[1].array, array_of(2), &mut ctx()); + Ok(()) +} + +#[rstest] +#[case::empty_batch( + vec![PlannerStep::Morsel(scope(0..1), vec![MorselStep::Compute(MorselOutput::Batch(array_of(0)))])], + "empty batch" +)] +#[case::empty_io(vec![PlannerStep::Io(vec![])], "empty batch")] +#[case::republished_fetch( + vec![ + PlannerStep::NeedsIO(vec![request(0, 0)]), + PlannerStep::NeedsIO(vec![request(0, 0)]), + ], + "published while it is outstanding" +)] +#[case::waiting_for_nothing(vec![PlannerStep::Wait], "waits with no outstanding fetch")] +#[case::still_waiting_after_the_last_delivery( + vec![PlannerStep::Io(vec![request(0, 0)]), PlannerStep::Wait], + "waits with no outstanding fetch" +)] +fn protocol_errors(#[case] mut steps: Vec, #[case] message: &str) -> VortexResult<()> { + let log = Log::default(); + steps.push(PlannerStep::Done); + let root = ScriptedPlanner::boxed("root", steps, &log); + let err = driver(&source()).run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref().is_some_and(|m| m.contains(message)), + "{err:?}" + ); + Ok(()) +} + +#[test] +fn step_limit_is_enforced() -> VortexResult<()> { + let log = Log::default(); + let mut steps: Vec<_> = (0..100).map(|_| PlannerStep::Continue).collect(); + steps.push(PlannerStep::Done); + let root = ScriptedPlanner::boxed("root", steps, &log); + let driver = Driver::new(source() as Arc).with_step_limit(10); + let err = driver.run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref() + .is_some_and(|m| m.contains("step limit of 10")), + "{err:?}" + ); + Ok(()) +} + +#[rstest] +#[case(IoIntent::Announce)] +#[case(IoIntent::Prefetch)] +fn optional_publication_does_not_park_or_receive_bytes( + #[case] intent: IoIntent, +) -> VortexResult<()> { + let log = Log::default(); + let source = source(); + let mut optional = request(0, 0); + optional.intent = intent; + let root = ScriptedPlanner::boxed( + "warm", + vec![ + PlannerStep::NeedsIO(vec![optional]), + one_batch_morsel(scope(0..1), 1), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source).run(root)?; + assert_eq!(batches.len(), 1); + let submitted = source.submissions(); + assert_eq!(submitted.len(), 1); + assert_eq!(submitted[0].1.intent, intent); + assert!(source.performed().is_empty()); + assert!( + !log.events() + .iter() + .any(|event| event.starts_with("deliver")) + ); + Ok(()) +} + +#[test] +fn optional_requests_cannot_be_wait_dependencies() -> VortexResult<()> { + let source = source(); + let mut optional = request(0, 0); + optional.intent = IoIntent::Announce; + let root = ScriptedPlanner::boxed( + "root", + vec![PlannerStep::Io(vec![optional])], + &Log::default(), + ); + let error = driver(&source) + .run(root) + .err() + .map(|error| error.to_string()); + assert!(error.is_some_and(|error| error.contains("waits with no outstanding fetch"))); + assert!(source.performed().is_empty()); + Ok(()) +} + +#[test] +fn optional_registration_can_be_promoted_to_fetch() -> VortexResult<()> { + let source = source(); + let mut optional = request(0, 0); + optional.intent = IoIntent::Announce; + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::NeedsIO(vec![optional]), + PlannerStep::Io(vec![request(0, 0)]), + PlannerStep::Done, + ], + &Log::default(), + ); + driver(&source).run(root)?; + assert_eq!(source.submissions().len(), 2); + assert_eq!(source.performed().len(), 1); + Ok(()) +} + +const _: () = { + const fn assert_send() {} + assert_send::(); +}; + +fn ready_root(name: &'static str, log: &Log, rows: Range) -> Box { + ScriptedPlanner::boxed( + name, + vec![one_batch_morsel(scope(rows), 1), PlannerStep::Done], + log, + ) +} + +fn reading_root(name: &'static str, log: &Log, rows: Range) -> Box { + ScriptedPlanner::boxed( + name, + vec![ + PlannerStep::Io(vec![request(0, 0)]), + one_batch_morsel(scope(rows), 2), + PlannerStep::Done, + ], + log, + ) +} + +/// Drains `run` until it waits or idles, describing what it handed back. +fn drain(run: &mut Run) -> VortexResult> { + let mut events = Vec::new(); + loop { + events.push(match run.advance()? { + Progress::Batch(batch) => format!("batch {:?} {:?}", batch.root, batch.scope.rows), + Progress::RootDone(root) => format!("done {root:?}"), + Progress::Waiting => { + events.push("waiting".to_string()); + return Ok(events); + } + Progress::Idle => { + events.push("idle".to_string()); + return Ok(events); + } + }); + } +} + +/// 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] +fn roots_finish_independently() -> VortexResult<()> { + let log = Log::default(); + let (reading_io, ready_io) = (source(), source()); + let mut run = Run::new().with_step_limit(1_000); + let reading = run.admit( + reading_root("reading", &log, 0..2), + scope(0..2), + Arc::clone(&reading_io) as Arc, + ); + let ready = run.admit( + ready_root("ready", &log, 2..3), + scope(2..3), + Arc::clone(&ready_io) as Arc, + ); + assert_eq!((reading, ready), (RootId(0), RootId(1))); + assert_eq!(run.live_roots(), 2); + + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(1) 2..3", "done RootId(1)", "waiting"] + ); + assert_eq!(run.live_roots(), 1); + assert_eq!((reading_io.clears(), ready_io.clears()), (0, 1)); + assert!(ready_io.submissions().is_empty()); + + run.complete(reading_io.wait()?)?; + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(0) 0..2", "done RootId(0)", "idle"] + ); + assert_eq!(run.live_roots(), 0); + assert_eq!(reading_io.clears(), 1); + Ok(()) +} + +#[test] +fn poll_completion_delivers_from_every_waiting_root() -> VortexResult<()> { + let log = Log::default(); + let sources = [source(), source()]; + let mut run = Run::new().with_step_limit(1_000); + for (index, io) in sources.iter().enumerate() { + let rows = index as u64..index as u64 + 1; + run.admit( + reading_root("root", &log, rows.clone()), + scope(rows), + Arc::clone(io) as Arc, + ); + } + assert_eq!(drain(&mut run)?, vec!["waiting"]); + let mut cx = Context::from_waker(Waker::noop()); + assert!(matches!(run.poll_completion(&mut cx), Poll::Ready(Ok(())))); + let mut events = drain(&mut run)?; + events.sort(); + assert_eq!( + events, + vec![ + "batch RootId(0) 0..1", + "batch RootId(1) 1..2", + "done RootId(0)", + "done RootId(1)", + "idle", + ] + ); + Ok(()) +} + +#[test] +fn cancel_drops_a_roots_work_and_clears_its_session() -> VortexResult<()> { + let log = Log::default(); + let (cancelled_io, kept_io) = (source(), source()); + let mut run = Run::new().with_step_limit(1_000); + let cancelled = run.admit( + reading_root("cancelled", &log, 0..2), + scope(0..2), + Arc::clone(&cancelled_io) as Arc, + ); + run.admit( + reading_root("kept", &log, 2..4), + scope(2..4), + Arc::clone(&kept_io) as Arc, + ); + assert_eq!(drain(&mut run)?, vec!["waiting"]); + + run.cancel(cancelled); + assert_eq!(run.live_roots(), 1); + assert_eq!(cancelled_io.clears(), 1); + run.complete(kept_io.wait()?)?; + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(1) 2..4", "done RootId(1)", "idle"] + ); + assert!( + !log.events() + .contains(&"deliver cancelled 0 bytes".to_string()) + ); + Ok(()) +} + +#[test] +fn dropping_a_run_clears_every_live_session() -> VortexResult<()> { + let log = Log::default(); + let io = source(); + let mut run = Run::new(); + run.admit( + reading_root("root", &log, 0..2), + scope(0..2), + Arc::clone(&io) as Arc, + ); + assert_eq!(drain(&mut run)?, vec!["waiting"]); + drop(run); + assert_eq!(io.clears(), 1); + Ok(()) +} + +/// Waiting for M fetches costs the driver one delivery each, not a visit that re-lists the rest. +#[test] +fn a_waiting_item_is_not_revisited_per_delivery() -> VortexResult<()> { + const FETCHES: u32 = 64; + let log = Log::default(); + let io = RecordingIoSource::default(); + let batch: Vec = (0..FETCHES) + .map(|id| { + io.canned_bytes(u64::from(id) * 4, buffer![0u8; 4].into_byte_buffer()); + request(id, u64::from(id) * 4) + }) + .collect(); + let io = Arc::new(io); + let root = ScriptedPlanner::boxed( + "root", + vec![PlannerStep::Io(batch), PlannerStep::Done], + &log, + ); + // Publish and the final compute: the deliveries in between cost no visit. + Driver::new(Arc::clone(&io) as Arc) + .with_step_limit(2) + .run(root)?; + assert_eq!(io.performed().len(), FETCHES as usize); + Ok(()) +} + +#[test] +fn waking_subtree_precedes_an_already_discovered_sibling() -> VortexResult<()> { + let log = Log::default(); + let io = source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Planner( + scope(0..1), + vec![ + PlannerStep::Io(vec![request(0, 0)]), + one_batch_morsel(scope(0..1), 1), + PlannerStep::Done, + ], + ), + PlannerStep::Morsel( + scope(1..2), + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(1))), + MorselStep::Compute(MorselOutput::Batch(array_of(1))), + MorselStep::Compute(MorselOutput::Done), + ], + ), + PlannerStep::Done, + ], + &log, + ); + let mut run = Run::new().with_step_limit(1_000); + run.admit(root, scope(0..2), Arc::clone(&io) as Arc); + let Progress::Batch(batch) = run.advance()? else { + panic!("the later sibling should run while the first child waits"); + }; + assert_eq!(batch.scope, scope(1..2)); + + run.complete(io.wait()?)?; + // The waking child creates a grandchild after the sibling already exists. Its subtree + // must still win, including over the sibling's continuation after its first batch. + let Progress::Batch(batch) = run.advance()? else { + panic!("the waking subtree should produce a batch first"); + }; + assert_eq!(batch.scope, scope(0..1)); + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(0) 1..2", "done RootId(0)", "idle"] + ); + Ok(()) +} + +#[test] +fn runnable_roots_follow_admission_order() -> VortexResult<()> { + let log = Log::default(); + let mut run = Run::new().with_step_limit(1_000); + for rows in [0..1, 1..2] { + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Continue, + one_batch_morsel(scope(rows.clone()), 1), + PlannerStep::Done, + ], + &log, + ); + run.admit(root, scope(rows), source()); + } + assert_eq!( + drain(&mut run)?, + vec![ + "batch RootId(0) 0..1", + "done RootId(0)", + "batch RootId(1) 1..2", + "done RootId(1)", + "idle", + ] + ); + Ok(()) +} + +/// Makes the recording source's completions available to nonblocking polling. +struct PollingSource { + inner: Arc, + outstanding: AtomicUsize, +} + +impl IoSource for PollingSource { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + self.outstanding.fetch_add( + batch + .iter() + .filter(|request| request.intent == IoIntent::Fetch) + .count(), + Ordering::Relaxed, + ); + self.inner.submit(owner, batch) + } + + fn poll(&self) -> VortexResult> { + if self.outstanding.load(Ordering::Relaxed) == 0 { + return Ok(None); + } + self.wait().map(Some) + } + + fn poll_completion(&self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(self.wait()) + } + + fn wait(&self) -> VortexResult { + let completion = self.inner.wait()?; + self.outstanding.fetch_sub(1, Ordering::Relaxed); + Ok(completion) + } + + fn release(&self, owner: IoOwnerId) { + self.inner.release(owner); + } + + fn clear(&self) { + self.outstanding.store(0, Ordering::Relaxed); + self.inner.clear(); + } +} + +#[test] +fn polls_io_while_a_later_branch_keeps_computing() -> VortexResult<()> { + let log = Log::default(); + let io = Arc::new(PollingSource { + inner: source(), + outstanding: AtomicUsize::new(0), + }); + let mut run = Run::new().with_step_limit(1_000); + run.admit(reading_root("earlier", &log, 0..1), scope(0..1), io); + let mut script: Vec<_> = (0..IO_POLL_INTERVAL * 2) + .map(|_| PlannerStep::Continue) + .collect(); + script.push(one_batch_morsel(scope(1..2), 1)); + script.push(PlannerStep::Done); + run.admit( + ScriptedPlanner::boxed("later", script, &log), + scope(1..2), + source(), + ); + let Progress::Batch(batch) = run.advance()? else { + panic!("ready IO should wake the earlier root before the later root finishes"); + }; + assert_eq!(batch.root, RootId(0)); + let later_visits = log + .events() + .iter() + .filter(|event| event.as_str() == "compute later") + .count(); + assert!(later_visits > 0 && later_visits <= IO_POLL_INTERVAL); + Ok(()) +} diff --git a/vortex-scan/src/planning/mod.rs b/vortex-scan/src/planning/mod.rs index 5e293dae237..3095712f161 100644 --- a/vortex-scan/src/planning/mod.rs +++ b/vortex-scan/src/planning/mod.rs @@ -6,6 +6,10 @@ //! Planners and morsels are `Send` and move between threads with the run that owns them; //! stage-specific phases remain private to those implementations. +pub mod driver; pub mod morsel; pub mod next; pub mod planner; + +#[cfg(test)] +mod tests; diff --git a/vortex-scan/src/planning/tests/mod.rs b/vortex-scan/src/planning/tests/mod.rs new file mode 100644 index 00000000000..59229044cbc --- /dev/null +++ b/vortex-scan/src/planning/tests/mod.rs @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scripted work and IO sources for the planning driver tests. + +pub mod scripted; diff --git a/vortex-scan/src/planning/tests/scripted.rs b/vortex-scan/src/planning/tests/scripted.rs new file mode 100644 index 00000000000..4035bb8ac61 --- /dev/null +++ b/vortex-scan/src/planning/tests/scripted.rs @@ -0,0 +1,418 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scripted planners, morsels, and a recording IO source for driver tests. +//! +//! A scripted object follows its script exactly and panics when the driver calls it against the +//! protocol, so a driver bug becomes a test failure that names the script position. + +use std::ops::Range; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::planning::morsel::Morsel; +use crate::planning::morsel::MorselOutput; +use crate::planning::planner::Planner; +use crate::planning::planner::PlannerOutput; +use crate::planning::planner::State; +use crate::planning::planner::WorkScope; + +/// Shared record of the calls the driver made, as readable one-line events. +#[derive(Clone, Default)] +pub struct Log(Arc>>); + +impl Log { + fn push(&self, event: String) { + self.0.lock().push(event); + } + + pub fn events(&self) -> Vec { + self.0.lock().clone() + } +} + +/// One scripted planner step. Every step but `Await` and `Wait` is the output of one `compute()`. +pub enum PlannerStep { + Done, + Continue, + /// Publishes the batch and goes straight on to the next step. + NeedsIO(IoBatch), + Planner(WorkScope, Vec), + Morsel(WorkScope, Vec), + /// Publishes the batch, then waits until every request in it is delivered. + Io(IoBatch), + /// Waits for requests an earlier `NeedsIO` step published. + Await(Vec), + /// Reports `Waiting` without having published anything. + Wait, +} + +/// One scripted morsel step. +pub enum MorselStep { + Compute(MorselOutput), + /// Publishes the batch, then waits until every request in it is delivered. + Io(IoBatch), +} + +/// The ids a waiting step still expects. +struct Pending(Vec); + +impl Pending { + fn of(batch: &IoBatch) -> Self { + Self(batch.iter().map(|request| request.request).collect()) + } + + /// Removes a delivered id and returns whether the step has everything it waited for. + fn deliver(&mut self, name: &str, position: usize, request: IoRequestId) -> bool { + let index = self + .0 + .iter() + .position(|pending| *pending == request) + .unwrap_or_else(|| { + panic!("{name}: delivery of unknown request {request:?} at step {position}") + }); + self.0.remove(index); + self.0.is_empty() + } +} + +pub struct ScriptedPlanner { + name: &'static str, + script: Vec, + position: usize, + /// What the step at `position` waits for, once it has published or received something. + pending: Option, + log: Log, +} + +impl ScriptedPlanner { + pub fn new(name: &'static str, script: Vec, log: &Log) -> Self { + Self { + name, + script, + position: 0, + pending: None, + log: log.clone(), + } + } + + /// A fresh scripted planner, boxed as the driver takes it. + pub fn boxed(name: &'static str, script: Vec, log: &Log) -> Box { + Box::new(Self::new(name, script, log)) + } +} + +impl IoConsumer for ScriptedPlanner { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + self.log.push(format!( + "deliver {} {} {}", + self.name, + request.0, + result.kind() + )); + let position = self.position; + if let Some(PlannerStep::Await(ids)) = self.script.get(position) + && self.pending.is_none() + { + self.pending = Some(Pending(ids.clone())); + } + let Some(pending) = self.pending.as_mut() else { + panic!( + "{}: unexpected delivery at script step {position}", + self.name + ); + }; + if pending.deliver(self.name, position, request) { + self.pending = None; + self.position += 1; + } + } +} + +impl Planner for ScriptedPlanner { + fn state(&self) -> State { + match self.script.get(self.position) { + None => State::Done, + Some(PlannerStep::Io(_)) if self.pending.is_none() => State::NeedsCompute, + Some(PlannerStep::Io(_) | PlannerStep::Await(_) | PlannerStep::Wait) => State::Waiting, + Some(_) => State::NeedsCompute, + } + } + + fn compute(&mut self) -> VortexResult { + let position = self.position; + if let Some(PlannerStep::Io(batch)) = self.script.get(position) { + assert!( + self.pending.is_none(), + "{}: compute called while script step {position} waits for IO", + self.name + ); + self.log.push(format!("publish {}", self.name)); + // A batch with nothing to wait for is left for the driver to reject. + self.pending = Some(Pending::of(batch)); + return Ok(PlannerOutput::NeedsIO(batch.clone())); + } + self.log.push(format!("compute {}", self.name)); + self.position += 1; + let Some(slot) = self.script.get_mut(position) else { + panic!("{}: compute past the end of the script", self.name); + }; + let step = std::mem::replace(slot, PlannerStep::Done); + Ok(match step { + PlannerStep::Done => PlannerOutput::Done, + PlannerStep::Continue => PlannerOutput::Continue, + PlannerStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + PlannerStep::Planner(scope, script) => { + PlannerOutput::Planner(scope, Self::boxed("child", script, &self.log)) + } + PlannerStep::Morsel(scope, script) => PlannerOutput::Morsel( + scope, + Box::new(ScriptedMorsel::new("morsel", script, &self.log)), + ), + PlannerStep::Io(_) | PlannerStep::Await(_) | PlannerStep::Wait => panic!( + "{}: compute called while script step {position} waits for IO", + self.name + ), + }) + } +} + +pub struct ScriptedMorsel { + name: &'static str, + script: Vec, + position: usize, + pending: Option, + log: Log, +} + +impl ScriptedMorsel { + pub fn new(name: &'static str, script: Vec, log: &Log) -> Self { + Self { + name, + script, + position: 0, + pending: None, + log: log.clone(), + } + } +} + +impl IoConsumer for ScriptedMorsel { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + self.log.push(format!( + "deliver {} {} {}", + self.name, + request.0, + result.kind() + )); + let position = self.position; + let Some(pending) = self.pending.as_mut() else { + panic!( + "{}: unexpected delivery at script step {position}", + self.name + ); + }; + if pending.deliver(self.name, position, request) { + self.pending = None; + self.position += 1; + } + } +} + +impl Morsel for ScriptedMorsel { + fn state(&self) -> State { + match self.script.get(self.position) { + None => State::Done, + Some(MorselStep::Io(_)) if self.pending.is_some() => State::Waiting, + Some(_) => State::NeedsCompute, + } + } + + fn compute(&mut self) -> VortexResult { + let position = self.position; + match self.script.get_mut(position) { + Some(MorselStep::Io(batch)) => { + assert!( + self.pending.is_none(), + "{}: compute called while script step {position} waits for IO", + self.name + ); + let batch = batch.clone(); + self.log.push(format!("publish {}", self.name)); + self.pending = Some(Pending::of(&batch)); + Ok(MorselOutput::NeedsIO(batch)) + } + Some(MorselStep::Compute(output)) => { + self.log.push(format!("compute {}", self.name)); + self.position += 1; + Ok(std::mem::replace(output, MorselOutput::Done)) + } + None => panic!("{}: compute past the end of the script", self.name), + } + } +} + +enum Canned { + Bytes(ByteBuffer), + Size(u64), + Fail(String), +} + +/// Which submitted request `wait` completes first. +#[derive(Clone, Copy, Default)] +pub enum Order { + /// Oldest submission first. + #[default] + Fifo, + /// Newest submission first, so a batch and concurrently parked items complete in reverse. + Lifo, +} + +/// An [`IoSource`] answering from canned buffers. Every submission is resolved at once and +/// held until `wait`, which hands completions back in the configured [`Order`]. It records the +/// targets in submission order, and how often it was cleared. +#[derive(Default)] +pub struct RecordingIoSource { + canned: Mutex>, + performed: Mutex>, + ready: Mutex>, + order: Order, + submissions: Mutex>, + clears: Mutex, +} + +impl RecordingIoSource { + pub fn lifo() -> Self { + Self { + order: Order::Lifo, + ..Default::default() + } + } + + /// Answers a range at `offset` with `bytes`. + pub fn canned_bytes(&self, offset: u64, bytes: ByteBuffer) { + self.canned + .lock() + .insert(IoTarget::range(offset, bytes.len()), Canned::Bytes(bytes)); + } + + /// Answers `target` with a size regardless of what it asked for, to test the driver's check. + pub fn answer_with_size(&self, target: IoTarget, size: u64) { + self.canned.lock().insert(target, Canned::Size(size)); + } + + pub fn fail(&self, target: IoTarget, message: &str) { + self.canned + .lock() + .insert(target, Canned::Fail(message.to_string())); + } + + pub fn submissions(&self) -> Vec<(IoOwnerId, IoRequest)> { + self.submissions.lock().clone() + } + + pub fn performed(&self) -> Vec { + self.performed.lock().clone() + } + + /// How many times the source was cleared. + pub fn clears(&self) -> usize { + *self.clears.lock() + } +} + +impl IoSource for RecordingIoSource { + fn submit(&self, work: IoOwnerId, batch: Vec) -> VortexResult<()> { + for request in batch { + self.submissions.lock().push((work, request)); + if request.intent != IoIntent::Fetch { + continue; + } + self.performed.lock().push(request.target); + let result = match self.canned.lock().get(&request.target) { + Some(Canned::Bytes(bytes)) => { + Ok(IoResult::Bytes(BufferHandle::new_host(bytes.clone()))) + } + Some(Canned::Size(size)) => Ok(IoResult::Size(*size)), + Some(Canned::Fail(message)) => Err(vortex_err!("{message}")), + None => Err(vortex_err!("no canned answer for {:?}", request.target)), + }; + self.ready.lock().push(Completion { + owner: work, + request: request.request, + result, + }); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + Ok(None) + } + + fn poll_completion(&self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(self.wait()) + } + + fn release(&self, work: IoOwnerId) { + self.ready + .lock() + .retain(|completion| completion.owner != work); + } + + fn clear(&self) { + *self.clears.lock() += 1; + self.ready.lock().clear(); + } + + fn wait(&self) -> VortexResult { + let mut ready = self.ready.lock(); + let completion = match self.order { + Order::Fifo if !ready.is_empty() => Some(ready.remove(0)), + Order::Fifo => None, + Order::Lifo => ready.pop(), + }; + completion.ok_or_else(|| vortex_err!("RecordingIoSource: wait with nothing outstanding")) + } +} + +/// A small non-nullable `u32` array of `n` rows. +pub fn array_of(n: u64) -> ArrayRef { + PrimitiveArray::from_iter(0..u32::try_from(n).expect("small test array")).into_array() +} + +/// An execution context over a bare array session, for `assert_arrays_eq!`. +pub fn ctx() -> ExecutionCtx { + array_session().create_execution_ctx() +} + +pub fn scope(rows: Range) -> WorkScope { + WorkScope { + file_ordinal: 0, + rows, + } +} From 7c1cb9be29969cdd2cff4d2deeaba1b832c3ed86 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 19:09:05 +0000 Subject: [PATCH 03/14] chore(scan): clippy fixes and record the single-fork design in the handoff doc Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017cUKVPYqfN7SLcXetK59i6 --- vortex-layout/src/scan/v2/mod.rs | 2 +- vortex-layout/src/scan/v2/tests.rs | 3 +- vortex-scan/design/planner-decode-slice.md | 53 ++++++++++------------ 3 files changed, 26 insertions(+), 32 deletions(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 16003898523..3beea785157 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -162,7 +162,7 @@ impl PreparedScan { let root = lower(&file.layout)?; let projection = optimize(plan_row_idx_expression(builder.projection, root.clone())?)?; let projection_starts: Arc<[u64]> = chunk_starts([&projection])?.into(); - let rows = builder.row_range.unwrap_or(0..root.row_count()); + let rows = builder.row_range.unwrap_or_else(|| 0..root.row_count()); let splits = projection_splits(&projection_starts, rows); let io = Arc::new(SegmentScanIo::new( file.segments, diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 9c71dc0f990..2c12e5b8131 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -64,8 +64,7 @@ fn table_chunk(start: usize) -> VortexResult { .map(|row| i32::try_from(row).unwrap_or(i32::MAX)), ); let strings = VarBinArray::from_iter( - rows.clone() - .map(|row| (row % 3 != 0).then(|| format!("row {row}"))), + rows.map(|row| (row % 3 != 0).then(|| format!("row {row}"))), DType::Utf8(true.into()), ); let elements = Buffer::from_iter((0..2 * CHUNK_ROWS).map(|i| i64::try_from(i).unwrap_or(0))); diff --git a/vortex-scan/design/planner-decode-slice.md b/vortex-scan/design/planner-decode-slice.md index 6c3cb157a92..6c3d4fa926d 100644 --- a/vortex-scan/design/planner-decode-slice.md +++ b/vortex-scan/design/planner-decode-slice.md @@ -119,35 +119,30 @@ Depends on A. New `vortex-layout/src/scan/planning/` and `vortex-layout/src/scan `vortex-layout/src/plan/physical/pipeline/tests.rs` (`lower(&layout)`), driven by `Driver::run` from PR A; compare with `assert_arrays_eq!` against the default scan. -### PR C: opt-in hook (~550 + tests) - -- `scan/v2/mod.rs` (113): `ENV_VAR = "VORTEX_SCAN_V2"`, `enabled()`, `supports()` (no list - columns), `ScanFile { layout, locations, segments }`. Drop `io` and `plans` fields and the - `conjuncts`, `file`, `prefetch`, `pruning`, `repeated_scan`, `share` modules. Gate - `supports_scan` on `builder.filter.is_none()` (scans with a filter use the default executor). -- `scan/v2/scan_builder.rs` (333): `ScanBuilder::from_default` and `into_stream`; replace - `prepare()` / `RepeatedScanV2` with a small local prepare: lower the projection with - `plan::physical::lower`, `optimize`, compute splits via `chunk_starts` + `projection_splits`, - build one `SplitPlan` per split. -- `scan/v2/stream.rs` (141): `into_stream` and `LazyScanStream`; drop the - `pruning::file_pruning_enabled` import and branch. -- `vortex-layout/src/scan/scan_builder.rs`: add `into_parts()` and `ScanParts` (+44 in the - source commit; drop `has_filter_and_limit`, which is subsumed by the no-filter gate). -- `vortex-file/src/planning/mod.rs`: only `segment_locations` and `scan_file(&VortexFile)` - (source L51-78, ~30 lines). It must be in `vortex-file` because - `Footer::segment_specs_with_metadata` (`vortex-file/src/footer/mod.rs` L210) is `pub(crate)`. - Use `VortexFile::segment_source()` (`vortex-file/src/file.rs` L151); drop `file.scan_io()` and - `scan_file_cache` (both from #10463 / plan caching). -- Tests: `v2/tests.rs` cases without filters, plus `vortex-file` round trips like the - `v2::into_stream(file.scan()?, scan_file(&file))` test in `vortex-file/src/open.rs` (~L571-608 - in the source commit). - -### PR D (optional): first caller - -- `vortex-datafusion/src/persistent/opener.rs` on develop, L454-468: keep the builder chain, - then `if scan::v2::enabled() && scan::v2::supports_scan(&b, &file) { v2::into_stream(...) } - else { b.into_stream() }` (~+15). Skip `diagnostics.rs` and `benchmark_segment_cache`. -- Non-engine alternative: next to `VortexFile::scan()` (`vortex-file/src/file.rs` L217). +### PR C: opt-in executor behind a single fork + +The source commit forked at every caller (`if v2::enabled() && v2::supports_scan(..) { +v2::into_stream(..) } else { b.into_stream() }` in the DataFusion opener, `vortex-file`, DuckDB). +This slice forks in exactly one place instead, and nothing V2-specific is added to the V1 path: + +- `ScanBuilder`'s lazy prepare step calls `prepare_stream`, which is the only fork: + `match v2::select(self)? { Executor::V1(b) => b.prepare()?.execute_stream(None), + Executor::V2(scan) => scan.execute_stream() }`. `into_stream`, `into_iter` and the Arrow + record-batch stream all go through it, so no caller changes. +- `v2::select` owns the whole gate: `VORTEX_SCAN_V2=1`, the builder carries a `ScanFile`, no + filter, no limit, and every layout is Flat/Chunked/Struct/Dict/Zoned/LegacyStats (no List + layout; lists written as flat leaves, the default, are supported). It logs the decision at + `debug` on target `vortex_layout::scan::v2`. Each later PR widens this gate in one place. +- `VortexFile::scan()` attaches `ScanFile` via `ScanBuilder::with_scan_file`, built from + `Footer::segment_specs_with_metadata`. Callers that build a `ScanBuilder` from a bare layout + reader (DataFusion opener, DuckDB) stay on V1 until they attach one with the same setter: wiring, + not a fork. +- `ScanBuilder::prepare()` / `build()` keep returning the V1 `RepeatedScan`. + +### PR D (optional): more callers + +- `vortex-datafusion/src/persistent/opener.rs`: add `.with_scan_file(vxf.scan_file())` to the + builder chain. No `if`. ### Later From e1e3f4ef979844001d6551f176abbaebaf706008 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 19:14:55 +0000 Subject: [PATCH 04/14] fix(scan): split V2 scans at chunk starts under zoned plans Develop's lower keeps Zoned and LegacyStats nodes in the plan, so the boundary walk stopped at them and every default-written file ran as a single split. Follow their data child. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017cUKVPYqfN7SLcXetK59i6 --- vortex-layout/src/scan/v2/splits.rs | 10 +++++++++- vortex-layout/src/scan/v2/tests.rs | 18 ++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/vortex-layout/src/scan/v2/splits.rs b/vortex-layout/src/scan/v2/splits.rs index 7537dfc2534..3d84f19b995 100644 --- a/vortex-layout/src/scan/v2/splits.rs +++ b/vortex-layout/src/scan/v2/splits.rs @@ -15,15 +15,17 @@ use vortex_error::VortexResult; use crate::plan::physical::Concat; use crate::plan::physical::Eval; use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; use crate::plan::physical::Pack; use crate::plan::physical::PlanRef; use crate::plan::physical::Take; +use crate::plan::physical::Zoned; /// The row positions where the chunks of `plans` start, over their shared row domain, sorted and /// without duplicates. /// /// This follows plans whose children cover the same rows, or known parts of them, and stops at any -/// other plan: a take's values cover other rows. +/// other plan: a take's values cover other rows, and zone statistics cover zones. pub(super) fn chunk_starts<'a>( plans: impl IntoIterator, ) -> VortexResult> { @@ -42,6 +44,12 @@ fn collect_starts(plan: &PlanRef, offset: u64, starts: &mut Vec) -> VortexR starts.push(offset + child_offset); collect_starts(&plan.child_required(index)?, offset + child_offset, starts)?; } + } else if let Some(zoned) = plan.as_opt::() { + if let Some(data) = zoned.data_plan()? { + collect_starts(&data, offset, starts)?; + } + } else if let Some(stats) = plan.as_opt::() { + collect_starts(&stats.data_plan()?, offset, starts)?; } else if plan.is::() { collect_starts(&plan.child_required(0)?, offset, starts)?; } else if plan.is::() || plan.is::() || plan.is::() { diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 2c12e5b8131..9eb421cc677 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -336,3 +336,21 @@ async fn primitive_column() -> VortexResult<()> { ); Ok(()) } + +/// Every chunk of the default zoned table is its own split, so the scan runs them in parallel. +#[rstest] +#[tokio::test] +async fn splits_at_chunk_starts( + #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, strategy.get().as_ref()).await?; + let Executor::V2(scan) = select_supported(builder(&session, &segments, &layout)?)? else { + panic!("the scan should run on the planner-driven executor"); + }; + let expected = (0..CHUNKS as u64) + .map(|chunk| chunk * CHUNK_ROWS as u64..(chunk + 1) * CHUNK_ROWS as u64) + .collect::>(); + assert_eq!(scan.splits, expected); + Ok(()) +} From 4dd799e302b4723c483229675a40fbdbe2c8905d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 19:33:07 +0000 Subject: [PATCH 05/14] perf(scan): share decoded dictionary values across a V2 scan's splits Each split had its own ValueCache, so a dictionary shared by many chunks was decoded once per split (HashTags 2.2x, CMSprovider 1.7x slower than V1). Share one cache per scan, as V1's reader does. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017cUKVPYqfN7SLcXetK59i6 --- vortex-layout/src/plan/physical/pipeline/values.rs | 8 ++++---- vortex-layout/src/scan/planning/mod.rs | 13 ++++--------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/vortex-layout/src/plan/physical/pipeline/values.rs b/vortex-layout/src/plan/physical/pipeline/values.rs index b4dd7623336..ad37fcb28b5 100644 --- a/vortex-layout/src/plan/physical/pipeline/values.rs +++ b/vortex-layout/src/plan/physical/pipeline/values.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Whole values, such as dictionary values, kept between the scans of one split. +//! Whole values, such as dictionary values, kept between the scans that share them. use std::sync::Arc; @@ -15,9 +15,9 @@ use crate::plan::physical::PlanRef; /// Values a plan produces over its whole domain whatever rows it is compiled for, kept for the /// later scans that share this cache. /// -/// A dictionary's values are read whole by every reader of its codes, so the filter and the -/// projection of one split decode them once. The cache lives as long as the split that owns it -/// and holds only the values that split read: nothing is kept for the rest of a scan. +/// A dictionary's values are read whole by every reader of its codes, so the scans sharing a +/// cache decode them once. The cache lives as long as its owner, such as a split or a whole +/// scan, and holds only the values its scans read. #[derive(Clone, Default)] pub struct ValueCache(Arc>>); diff --git a/vortex-layout/src/scan/planning/mod.rs b/vortex-layout/src/scan/planning/mod.rs index b632e55e139..da8209f873e 100644 --- a/vortex-layout/src/scan/planning/mod.rs +++ b/vortex-layout/src/scan/planning/mod.rs @@ -61,8 +61,9 @@ pub struct ScanPlans { pub projection_starts: Arc<[u64]>, /// The global row index of the root row domain's first row. pub row_offset: u64, - /// Whole values, such as dictionary values, kept between the scans of one split. Each split - /// starts with its own, dropped when the split finishes. + /// Whole values, such as dictionary values, shared by every split of the scan, so a + /// dictionary read by many chunks is decoded once, as the default executor's reader caches + /// it. Dropped with the scan. pub values: ValueCache, } @@ -76,11 +77,5 @@ pub struct SelectedRows { /// The work for one split: project the rows of `scope` selected by `mask`. pub fn plan_projection(plans: ScanPlans, scope: WorkScope, mask: Mask) -> Box { - Box::new(ProjectionPlanner::new( - ScanPlans { - values: ValueCache::default(), - ..plans - }, - SelectedRows { scope, mask }, - )) + Box::new(ProjectionPlanner::new(plans, SelectedRows { scope, mask })) } From 6432c580db405621340a41ad985f1b10b4537f3d Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 15 Oct 2026 22:55:29 +0100 Subject: [PATCH 06/14] perf(scan): reduce V2 decompression planning and memory overhead Reuse V1 split policies, avoid redundant identity-projection planning, and release decoded dictionary values after their consuming splits finish. Reduce pipeline allocation, plan retention, and ready-I/O scheduling overhead. Validated with repeated V1/V2/Parquet decompression comparisons, output parity across 14 datasets, and 175 focused pipeline and scan tests. Some workloads still retain modest RSS and parallel duplicate-read overhead versus V1. Signed-off-by: Joe Isaacs --- .../src/plan/physical/pipeline/compile.rs | 49 ++++---- .../src/plan/physical/pipeline/port.rs | 83 ++++++++++--- .../src/plan/physical/pipeline/scan.rs | 15 ++- .../src/plan/physical/pipeline/values.rs | 82 ++++++++++++- vortex-layout/src/plan/physical/plans/pack.rs | 6 +- vortex-layout/src/scan/planning/mod.rs | 5 +- vortex-layout/src/scan/planning/projection.rs | 48 +++++--- vortex-layout/src/scan/scan_builder.rs | 4 +- vortex-layout/src/scan/split_by.rs | 21 +++- vortex-layout/src/scan/v2/io.rs | 76 +++++++++++- vortex-layout/src/scan/v2/mod.rs | 73 ++++++++++-- vortex-layout/src/scan/v2/splits.rs | 99 ++++++++++++++-- vortex-layout/src/scan/v2/tests.rs | 112 ++++++++++++++++++ 13 files changed, 568 insertions(+), 105 deletions(-) diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index 992830a4eaf..74f5296cd19 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -192,7 +192,7 @@ impl Compiler<'_> { } return port; } - let ranges = std::mem::take(shares.reaches.get_mut(*segment as usize)?); + let ranges = shares.shared.remove(&segment)?; let mut claimants: SmallVec<[usize; 4]> = SmallVec::new(); for rows in &ranges { claimants.extend( @@ -240,10 +240,11 @@ pub(crate) struct Shares { ordered: bool, /// Splits that have finished, and so will claim nothing. finished: Vec, - /// By segment id, which a file numbers densely: for each segment more than one reader will - /// read and whose decoding pipeline is not built yet, the range of plan rows each of its - /// readers reads it for. Every split overlapping a range reads it once for that range. - reaches: Vec; 1]>>, + /// Segment readers collected before the scan starts. A split may read only a few of a + /// file's segments, so recording a reader must not allocate by the file-wide segment id. + reaches: Vec<(SegmentId, Range)>, + /// The row ranges of segments read more than once, until their decoding pipeline is built. + shared: FxHashMap; 1]>>, /// By split: the unclaimed ports of each shared segment the split reads. ports: Vec>>, } @@ -257,6 +258,7 @@ impl Shares { splits, ordered, reaches: Vec::new(), + shared: FxHashMap::default(), } } @@ -290,7 +292,7 @@ impl Shares { pub(crate) fn add(&mut self, plan: &PlanRef, rows: Range) -> VortexResult<()> { let reaches = &mut self.reaches; plan.reach(rows, &Reach::Offset(0), &mut |segment, rows| { - record(reaches, segment, rows) + reaches.push((segment, rows)) }) } @@ -298,23 +300,25 @@ impl Shares { /// pipeline that needs them. pub(crate) fn retain_shared(&mut self) { let mut reaches = std::mem::take(&mut self.reaches); - for ranges in &mut reaches { - let shared = match ranges.as_slice() { - [] => false, - // One reader per split overlapping the range: shared when the range crosses a - // split boundary. - [rows] => self + reaches.sort_unstable_by_key(|(segment, _)| *segment); + for readers in reaches.chunk_by(|(left, _), (right, _)| left == right) { + let (segment, rows) = &readers[0]; + // One reader per split overlapping the range: shared when the range crosses a + // split boundary, or when multiple plan nodes read the segment. + if readers.len() == 1 + && self .overlapping(rows) .filter(|&split| self.overlaps(split, rows)) .nth(1) - .is_some(), - _ => true, - }; - if !shared { - *ranges = SmallVec::new(); + .is_none() + { + continue; } + self.shared.insert( + *segment, + readers.iter().map(|(_, rows)| rows.clone()).collect(), + ); } - self.reaches = reaches; } /// Every port made for a reader that has not claimed it yet. @@ -371,15 +375,6 @@ impl Reach { } } -/// Adds a reader of `segment` reading it for `rows`. -fn record(reaches: &mut Vec; 1]>>, segment: SegmentId, rows: Range) { - let index = *segment as usize; - if index >= reaches.len() { - reaches.resize_with(index + 1, SmallVec::new); - } - reaches[index].push(rows); -} - /// The chunks of `concat` overlapping `rows`: each chunk's index, where it starts, and the rows /// of `rows` it holds, in the concatenation's domain. Found by binary search, so a split pays /// for the chunks it overlaps, not for the chunks of the file. diff --git a/vortex-layout/src/plan/physical/pipeline/port.rs b/vortex-layout/src/plan/physical/pipeline/port.rs index a5c4c38a6e4..e089a9f0f76 100644 --- a/vortex-layout/src/plan/physical/pipeline/port.rs +++ b/vortex-layout/src/plan/physical/pipeline/port.rs @@ -34,9 +34,10 @@ pub(crate) enum SourceEnd { /// A queue of batches with one sink and one source. /// -/// Its buffer is allocated once, for the capacity. Only an unbounded port, a share's waiting for a -/// source compiled later, grows past the default. +/// The first batch is inline: a segment scan commonly emits only one batch. Additional batches +/// use a deque, whose buffer is retained when the port is reused. pub(crate) struct Queue { + front: Option, batches: VecDeque, /// Batches the queue may hold before its sink is blocked. capacity: usize, @@ -57,7 +58,8 @@ pub(crate) struct Queue { impl Queue { fn new(capacity: usize, sink: Option, source: SourceEnd) -> Self { Self { - batches: VecDeque::with_capacity(capacity.clamp(1, DEFAULT_CAPACITY)), + front: None, + batches: VecDeque::new(), capacity: capacity.max(1), closed: false, changed: false, @@ -69,7 +71,7 @@ impl Queue { } pub(crate) fn has_room(&self) -> bool { - self.batches.len() < self.capacity + self.len() < self.capacity } pub(crate) fn closed(&self) -> bool { @@ -83,7 +85,7 @@ impl Queue { } pub(crate) fn len(&self) -> usize { - self.batches.len() + usize::from(self.front.is_some()) + self.batches.len() } #[cfg(debug_assertions)] @@ -96,7 +98,7 @@ impl Queue { } pub(crate) fn pop(&mut self) -> Option { - self.batches.pop_front() + std::mem::replace(&mut self.front, self.batches.pop_front()) } } @@ -108,12 +110,12 @@ pub struct Inlet<'a> { impl Inlet<'_> { /// The oldest batch, in place, so the source can slice it down without removing it. pub fn peek_mut(&mut self) -> Option<&mut ArrayRef> { - self.queue.batches.front_mut() + self.queue.front.as_mut() } /// The oldest batch, removed. Removing a batch is what frees its slot. pub fn take(&mut self) -> Option { - self.queue.batches.pop_front() + self.queue.pop() } /// Whether the sink has finished. Closed and empty is the end of the inlet. @@ -123,23 +125,28 @@ impl Inlet<'_> { /// Whether no batch is queued. pub fn is_empty(&self) -> bool { - self.queue.batches.is_empty() + self.queue.front.is_none() } /// Batches queued. pub fn len(&self) -> usize { - self.queue.batches.len() + self.queue.len() } /// Whether the sink can push no more until the source takes a batch. An unbounded inlet /// counts as full at the default capacity, so a source waiting for it to fill still runs. pub fn full(&self) -> bool { - self.queue.batches.len() >= self.queue.capacity.min(DEFAULT_CAPACITY) + self.queue.len() >= self.queue.capacity.min(DEFAULT_CAPACITY) } /// Rows queued, over every batch. pub fn rows(&self) -> usize { - self.queue.batches.iter().map(|batch| batch.len()).sum() + self.queue + .front + .iter() + .chain(self.queue.batches.iter()) + .map(|batch| batch.len()) + .sum() } } @@ -214,10 +221,9 @@ impl Arena { ) -> PortId { match self.free.pop() { Some(id) => { - // A freed slot keeps its buffer, so reusing it allocates only for a larger port. + // A freed slot keeps any overflow buffer for reuse. let queue = &mut self.queues[id]; let capacity = capacity.max(1); - queue.batches.reserve_exact(capacity.min(DEFAULT_CAPACITY)); queue.capacity = capacity; queue.closed = false; queue.changed = false; @@ -257,7 +263,11 @@ impl Arena { "push into port {id} with no consumer" ); if !batch.is_empty() { - queue.batches.push_back(batch); + if queue.front.is_none() { + queue.front = Some(batch); + } else { + queue.batches.push_back(batch); + } queue.changed = true; } } @@ -278,6 +288,7 @@ impl Arena { let queue = &mut self.queues[id]; queue.source_gone = true; queue.source = SourceEnd::Unclaimed; + queue.front = None; queue.batches.clear(); self.maybe_free(id); } @@ -292,6 +303,7 @@ impl Arena { ); queue.closed = true; queue.source_gone = true; + queue.front = None; queue.batches.clear(); self.maybe_free(id); } @@ -300,6 +312,7 @@ impl Arena { let queue = &mut self.queues[id]; if queue.closed && queue.source_gone && !queue.free { queue.free = true; + queue.front = None; queue.batches.clear(); self.free.push(id); } @@ -317,3 +330,43 @@ impl Arena { self.queues.len() - self.free.len() } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_buffer::buffer; + + use super::Arena; + use super::SourceEnd; + + #[rstest] + #[case::single(1)] + #[case::bounded(4)] + #[case::unbounded(usize::MAX)] + fn queued_batches_keep_order_and_reused_ports_are_empty(#[case] capacity: usize) { + let mut arena = Arena::default(); + let port = arena.create(capacity, None, SourceEnd::Scan(0)); + let arrays: Vec<_> = (0..capacity.min(4)) + .map(|i| buffer![i as u64].into_array()) + .collect(); + for array in &arrays { + assert!(arena.has_room(port)); + arena.push(port, array.clone()); + } + assert_eq!(arena.inlet(port).len(), arrays.len()); + assert_eq!(arena.inlet(port).rows(), arrays.len()); + arena.close(port); + for expected in arrays { + let actual = arena.inlet(port).take().expect("queued batch"); + assert!(ArrayRef::ptr_eq(&actual, &expected)); + } + assert!(arena.inlet(port).is_empty()); + arena.drop_source(port); + let reused = arena.create(1, None, SourceEnd::Scan(0)); + assert_eq!(reused, port); + assert!(arena.inlet(reused).is_empty()); + assert!(arena.has_room(reused)); + } +} diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index fc444340e83..9d650ebbd26 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -626,7 +626,7 @@ struct Active { /// reads of later splits are in flight while earlier ones compute. pub struct Scan { core: Core, - root: Root, + root: Option, splits: Vec, next: usize, active: Vec>, @@ -685,7 +685,7 @@ impl Scan { core.dirty_flags = vec![false; 1]; Ok(Self { core, - root, + root: Some(root), splits, next: 0, active: vec![None], @@ -766,12 +766,21 @@ impl Scan { let index = self.next; self.next += 1; let split = self.splits[index].clone(); - let output = match &self.root { + let output = match self + .root + .as_ref() + .ok_or_else(|| vortex_err!("Scan has no plan for its next split"))? + { Root::Plan(plan) => { self.core .compile_stage(plan, split.rows.clone(), &split.mask, (slot, index))? } }; + if self.next == self.splits.len() { + // Compiled operators own what they need. The last split no longer needs the + // lowering tree or the source layouts retained by its lazy children. + self.root = None; + } if output.is_none() { self.core.finish_split(index); } else { diff --git a/vortex-layout/src/plan/physical/pipeline/values.rs b/vortex-layout/src/plan/physical/pipeline/values.rs index ad37fcb28b5..9b40314f279 100644 --- a/vortex-layout/src/plan/physical/pipeline/values.rs +++ b/vortex-layout/src/plan/physical/pipeline/values.rs @@ -12,19 +12,24 @@ use vortex_array::ArrayRef; use crate::plan::physical::PlanByPtr; use crate::plan::physical::PlanRef; +type CachedValue = Arc>>; + /// Values a plan produces over its whole domain whatever rows it is compiled for, kept for the /// later scans that share this cache. /// -/// A dictionary's values are read whole by every reader of its codes, so the scans sharing a -/// cache decode them once. The cache lives as long as its owner, such as a split or a whole -/// scan, and holds only the values its scans read. +/// A dictionary's values are read whole by every reader of its codes, so readers sharing this +/// cache can reuse them across filters, projections, and splits. The caller controls its lifetime; +/// only whole values that a reader requested are retained. #[derive(Clone, Default)] -pub struct ValueCache(Arc>>); +pub struct ValueCache(Arc>>); impl ValueCache { /// The values `plan` produced earlier. pub(crate) fn get(&self, plan: &PlanRef) -> Option { - self.0.lock().get(&PlanByPtr(plan.clone())).cloned() + self.0 + .lock() + .get(&PlanByPtr(plan.clone())) + .and_then(|value| value.lock().clone()) } /// Records the values of `plan`, keeping the first when two race. @@ -32,7 +37,72 @@ impl ValueCache { self.0 .lock() .entry(PlanByPtr(plan.clone())) - .or_insert(values) + .or_default() + .lock() + .get_or_insert(values) .clone() } + + /// Shares only these entries. Once the parent cache is dropped, a value lives until the + /// last cache using its plan is dropped, rather than until every split finishes. + pub(crate) fn for_plans<'a>(&self, plans: impl Iterator) -> Self { + let mut entries = self.0.lock(); + let scoped = plans + .map(|plan| { + let key = PlanByPtr(plan.clone()); + let value = entries.entry(key.clone()).or_default(); + (key, Arc::clone(value)) + }) + .collect(); + Self(Arc::new(Mutex::new(scoped))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + + use super::ValueCache; + use crate::plan::physical::SegmentScanPlan; + use crate::segments::SegmentId; + + #[test] + fn scoped_values_live_until_the_last_consumer() -> VortexResult<()> { + let array = buffer![1_i32, 2, 3].into_array(); + let plans = [0_u32, 1].map(|id| { + SegmentScanPlan::new( + array.dtype().clone(), + array.len() as u64, + SegmentId::from(id), + ReadContext::new([]), + None, + ) + .into_plan() + }); + let parent = ValueCache::default(); + let first = parent.for_plans(std::iter::once(&plans[0])); + let last = parent.for_plans(std::iter::once(&plans[0])); + let unrelated = parent.for_plans(std::iter::once(&plans[1])); + drop(parent); + + first.insert(&plans[0], array.clone()); + unrelated.insert(&plans[1], array.clone()); + let entry = Arc::downgrade(first.0.lock().values().next().expect("one entry")); + drop(first); + assert!(ArrayRef::ptr_eq( + &last.get(&plans[0]).expect("shared value"), + &array, + )); + assert!(entry.upgrade().is_some()); + drop(last); + assert!(entry.upgrade().is_none()); + assert!(unrelated.get(&plans[1]).is_some()); + Ok(()) + } } diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index 825897e33ee..25f12380493 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -235,10 +235,10 @@ impl PlanVTable for Pack { return Ok(None); } let mut chains = Vec::with_capacity(count); - for child in plan.children().iter() { + for child in plan.children().iter_refs() { chains.push( compiler - .compile(&child?, rows.clone(), mask)? + .compile(child?, rows.clone(), mask)? .ok_or_else(|| vortex_err!("A Pack field produced no rows"))?, ); } @@ -258,7 +258,7 @@ impl PlanVTable for Pack { at: &Reach, visit: &mut dyn FnMut(SegmentId, Range), ) -> VortexResult<()> { - for child in plan.children().iter() { + for child in plan.children().iter_refs() { child?.reach(rows.clone(), at, visit)?; } Ok(()) diff --git a/vortex-layout/src/scan/planning/mod.rs b/vortex-layout/src/scan/planning/mod.rs index da8209f873e..c110846a9b5 100644 --- a/vortex-layout/src/scan/planning/mod.rs +++ b/vortex-layout/src/scan/planning/mod.rs @@ -61,9 +61,8 @@ pub struct ScanPlans { pub projection_starts: Arc<[u64]>, /// The global row index of the root row domain's first row. pub row_offset: u64, - /// Whole values, such as dictionary values, shared by every split of the scan, so a - /// dictionary read by many chunks is decoded once, as the default executor's reader caches - /// it. Dropped with the scan. + /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans + /// always read the same complete row domain, independently of a split's selection. pub values: ValueCache, } diff --git a/vortex-layout/src/scan/planning/projection.rs b/vortex-layout/src/scan/planning/projection.rs index e230a3789a6..22930a4e996 100644 --- a/vortex-layout/src/scan/planning/projection.rs +++ b/vortex-layout/src/scan/planning/projection.rs @@ -1,10 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -32,7 +31,7 @@ use crate::scan::v2::splits::projection_splits; /// with a selected row becomes one [`ProjectionMorsel`], so a morsel reads at most one chunk of /// every column. The planner finishes once it has handed out every morsel. pub struct ProjectionPlanner { - plans: ScanPlans, + plans: Option, /// The split's selected rows, until the first compute cuts them. selected: Option, /// The projection splits not yet handed out, last first. @@ -43,7 +42,7 @@ impl ProjectionPlanner { /// Creates a planner for `selected`. pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { Self { - plans, + plans: Some(plans), selected: Some(selected), pending: Vec::new(), } @@ -54,7 +53,8 @@ impl ProjectionPlanner { let SelectedRows { scope, mask } = selected; let start = scope.rows.start; let index = |row: u64| usize::try_from(row - start).vortex_expect("split row fits usize"); - let mut splits = projection_splits(&self.plans.projection_starts, scope.rows.clone()) + let plans = self.plans.as_ref().vortex_expect("uncut planner has plans"); + let mut splits = projection_splits(&plans.projection_starts, scope.rows.clone()) .into_iter() .filter_map(|rows| { let mask = mask.slice(index(rows.start)..index(rows.end)); @@ -94,9 +94,19 @@ impl Planner for ProjectionPlanner { return Ok(PlannerOutput::Done); }; let scope = selected.scope.clone(); + // The driver's continuation follows its children. Give the last morsel ownership so + // the completed planner does not keep the lowering tree alive while that morsel runs. + let plans = if self.pending.is_empty() { + self.plans.take().vortex_expect("planner has plans") + } else { + self.plans + .as_ref() + .vortex_expect("planner has plans") + .clone() + }; Ok(PlannerOutput::Morsel( scope, - Box::new(ProjectionMorsel::new(self.plans.clone(), selected)), + Box::new(ProjectionMorsel::new(plans, selected)), )) } } @@ -104,7 +114,8 @@ impl Planner for ProjectionPlanner { /// Runs the projection plan over a split's selected rows and returns them as one batch, in row /// order. A selection that projects to no rows finishes without a batch. pub struct ProjectionMorsel { - plans: ScanPlans, + plans: Option, + dtype: DType, selected: SelectedRows, graph: Option, pieces: Vec, @@ -115,7 +126,8 @@ impl ProjectionMorsel { /// Creates a morsel projecting `selected`. pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { Self { - plans, + dtype: plans.projection.dtype().clone(), + plans: Some(plans), selected, graph: None, pieces: Vec::new(), @@ -134,7 +146,7 @@ impl ProjectionMorsel { let array = match arrays.len() { 0 => return Ok(MorselOutput::Done), 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.plans.projection.dtype().clone())?.into_array(), + _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), }; Ok(MorselOutput::Batch(array)) } @@ -165,19 +177,19 @@ impl Morsel for ProjectionMorsel { vortex_bail!("ProjectionMorsel: compute called after Done"); } let Some(graph) = self.graph.as_mut() else { + let plans = self + .plans + .take() + .vortex_expect("an uncompiled projection still has its plans"); let graph = ScanGraph::try_new( - self.plans.session.clone(), - &self.plans.projection, + plans.session, + &plans.projection, self.selected.scope.rows.clone(), self.selected.mask.clone(), - self.plans.row_offset, - self.plans.values.clone(), + plans.row_offset, + plans.values, )?; - self.graph = Some(ProtocolGraph::new( - graph, - Arc::clone(&self.plans.locations), - 0, - )); + self.graph = Some(ProtocolGraph::new(graph, plans.locations, 0)); return Ok(MorselOutput::Continue); }; if graph.state() == State::Done { diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 47b4d80a80d..590d2b644fe 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -69,10 +69,10 @@ pub struct ScanBuilder { // TODO(joe): replace this is usage of row_id selection, see pub(super) selection: Selection, /// How to split the file for concurrent processing. - split_by: SplitBy, + pub(super) split_by: SplitBy, /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare) /// uses them instead of walking the layout. - natural_splits: Option>, + pub(super) natural_splits: Option>, /// The number of splits to make progress on concurrently **per-thread**. pub(super) concurrency: usize, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index b995e46222f..761959132a6 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -56,18 +56,29 @@ impl SplitBy { layout_reader: &dyn LayoutReader, row_range: &Range, field_mask: &[FieldMask], + ) -> VortexResult> { + let boundaries = if matches!(self, Self::RowCount(_)) { + Vec::new() + } else { + layout_boundaries(layout_reader, row_range, field_mask)? + }; + self.split_boundaries(boundaries, row_range) + } + + /// Applies this split policy to natural boundaries from a layout or physical plan. + pub(super) fn split_boundaries( + &self, + boundaries: Vec, + row_range: &Range, ) -> VortexResult> { Ok(match *self { - SplitBy::Layout => layout_boundaries(layout_reader, row_range, field_mask)?, + SplitBy::Layout => boundaries, SplitBy::LayoutSubSplitting { max_rows } => { vortex_ensure!( max_rows > 0, "SplitBy::LayoutSubSplitting requires a non-zero max_rows" ); - subdivide_large_spans( - layout_boundaries(layout_reader, row_range, field_mask)?, - max_rows, - ) + subdivide_large_spans(boundaries, max_rows) } SplitBy::RowCount(n) => row_range .clone() diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 66c4ee6d4a7..888e53506a7 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -79,6 +79,7 @@ pub(super) struct SegmentIoSource { segments: Arc, ranges: SegmentRanges, reads: Mutex>>, + ready: Mutex>, /// Reads started by prefetches, consumed by the next fetch of their segment. prefetched: Mutex>, /// Announced segments, requested and never polled. @@ -94,6 +95,7 @@ impl SegmentIoSource { segments, ranges, reads: Mutex::new(FuturesUnordered::new()), + ready: Mutex::default(), prefetched: Mutex::default(), announced: Mutex::default(), } @@ -103,6 +105,7 @@ impl SegmentIoSource { impl IoSource for SegmentIoSource { fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { let reads = self.reads.lock(); + let mut ready = self.ready.lock(); let mut prefetched = self.prefetched.lock(); for request in batch { let IoTarget::Range { offset, len, .. } = request.target else { @@ -143,7 +146,7 @@ impl IoSource for SegmentIoSource { }); continue; } - let bytes = match prefetched.remove(&segment) { + let mut bytes = match prefetched.remove(&segment) { Some(read) => read .map_err(move |err| { vortex_err!("prefetched read of segment {segment} failed: {err}") @@ -157,6 +160,16 @@ impl IoSource for SegmentIoSource { .unwrap_or_else(|| self.segments.request(segment)), }; let id = request.request; + // In-memory and cached segments need no allocated asynchronous task. A pending + // read is polled again by FuturesUnordered before the driver can go to sleep. + if let Poll::Ready(result) = bytes.poll_unpin(&mut Context::from_waker(Waker::noop())) { + ready.push_back(Completion { + owner, + request: id, + result: result.map(IoResult::Bytes), + }); + continue; + } reads.push( async move { Completion { @@ -172,6 +185,9 @@ impl IoSource for SegmentIoSource { } fn poll(&self) -> VortexResult> { + if let Some(completion) = self.ready.lock().pop_front() { + return Ok(Some(completion)); + } // The split's future polls again with its own waker before it waits, so a read that // finishes after this check is not missed. let mut cx = Context::from_waker(Waker::noop()); @@ -182,6 +198,9 @@ impl IoSource for SegmentIoSource { } fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + if let Some(completion) = self.ready.lock().pop_front() { + return Poll::Ready(Ok(completion)); + } match self.reads.lock().poll_next_unpin(cx) { Poll::Ready(Some(completion)) => Poll::Ready(Ok(completion)), Poll::Ready(None) => Poll::Ready(Err(vortex_err!( @@ -199,6 +218,7 @@ impl IoSource for SegmentIoSource { fn clear(&self) { self.reads.lock().clear(); + self.ready.lock().clear(); self.prefetched.lock().clear(); self.announced.lock().clear(); } @@ -211,11 +231,14 @@ mod tests { use std::sync::atomic::Ordering; use futures::FutureExt; + use futures::channel::oneshot; use futures::future; + use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_io::request::IoIntent; use vortex_io::request::IoOwnerId; use vortex_io::request::IoRequest; @@ -250,6 +273,57 @@ mod tests { } } + struct MixedSegments(Mutex>>); + + impl SegmentSource for MixedSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + if *id == 0 { + return future::ready(Ok(BufferHandle::new_host(ByteBuffer::from(vec![0; 4])))) + .boxed(); + } + let receiver = self.0.lock().take().expect("one pending read"); + async move { + receiver + .await + .map(BufferHandle::new_host) + .map_err(|err| vortex_err!("{err}")) + } + .boxed() + } + } + + #[tokio::test] + async fn pending_read_wakes_after_an_immediate_completion() -> VortexResult<()> { + let (sender, receiver) = oneshot::channel(); + let segments = Arc::new(MixedSegments(Mutex::new(Some(receiver)))); + let locations = [0, 4].map(|offset| SegmentLocation { + offset, + length: 4, + alignment: Alignment::none(), + }); + let io = SegmentIoSource::new(segments, segment_ranges(&locations)); + let mut pending = request(IoIntent::Fetch, 1); + pending.target = IoTarget::range(4, 4); + io.submit(IoOwnerId(0), vec![pending, request(IoIntent::Fetch, 0)])?; + let immediate = io.poll()?.expect("the in-memory read is already ready"); + assert_eq!(immediate.request, IoRequestId(0)); + assert!(immediate.result.is_ok()); + assert!(io.poll()?.is_none()); + + let delivery = tokio::spawn(async move { + tokio::task::yield_now().await; + sender + .send(ByteBuffer::from(vec![1; 4])) + .expect("read is live"); + }); + let completed = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completed.request, IoRequestId(1)); + assert!(completed.result.is_ok()); + delivery.await.expect("delivery task succeeds"); + assert!(io.poll()?.is_none()); + 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] diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 3beea785157..6d0a87e949b 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -21,6 +21,7 @@ pub(crate) mod splits; mod tests; use std::env; +use std::iter; use std::ops::Range; use std::sync::Arc; use std::sync::LazyLock; @@ -46,6 +47,7 @@ use crate::layouts::flat::Flat; use crate::layouts::struct_::Struct; use crate::layouts::zoned::LegacyStats; use crate::layouts::zoned::Zoned; +use crate::plan::physical::PlanRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; @@ -105,7 +107,9 @@ fn select_supported(builder: ScanBuilder) -> VortexResult< } else { match &builder.scan_file { None => Some("no scan file"), - Some(file) if !supports_layout(&file.layout)? => Some("layout"), + Some(file) if !supports_layout(&file.layout, builder.projection.is_root())? => { + Some("layout") + } Some(_) => None, } }; @@ -117,10 +121,10 @@ fn select_supported(builder: ScanBuilder) -> VortexResult< Ok(Executor::V2(PreparedScan::try_new(builder)?)) } -/// Whether every layout under `layout` lowers to plans with a pipeline implementation. +/// 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) -> VortexResult { +fn supports_layout(layout: &LayoutRef, identity: bool) -> VortexResult { if !(layout.is::() || layout.is::() || layout.is::() @@ -130,8 +134,15 @@ fn supports_layout(layout: &LayoutRef) -> VortexResult { { 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), + }; + } for child in layout.children()? { - if !supports_layout(&child)? { + if !supports_layout(&child, identity)? { return Ok(false); } } @@ -142,6 +153,8 @@ fn supports_layout(layout: &LayoutRef) -> VortexResult { pub(crate) struct PreparedScan { session: VortexSession, plans: ScanPlans, + /// Whole dictionary values and the root rows that use them, for an identity projection. + value_scopes: Option, PlanRef)>>, /// Serves the splits' reads. io: Arc, /// The scan's splits, over the root row domain, in row order. @@ -160,10 +173,38 @@ impl PreparedScan { .scan_file .ok_or_else(|| vortex_err!("a V2 scan needs a scan file"))?; let root = lower(&file.layout)?; - let projection = optimize(plan_row_idx_expression(builder.projection, root.clone())?)?; - let projection_starts: Arc<[u64]> = chunk_starts([&projection])?.into(); - let rows = builder.row_range.unwrap_or_else(|| 0..root.row_count()); - let splits = projection_splits(&projection_starts, rows); + 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 projection = if identity { + root + } else { + optimize(plan_row_idx_expression(builder.projection, root)?)? + }; + let mut value_scopes = identity.then(Vec::new); + let projection_starts: Arc<[u64]> = + chunk_starts([&projection], value_scopes.as_mut())?.into(); + let rows = builder.row_range.unwrap_or(0..row_count); + // Whole-column dictionaries already need the scan's full lifetime. Only dictionaries + // covering part of the scan need separate cache owners for their consuming splits. + let value_scopes = value_scopes.filter(|values| { + values + .iter() + .any(|(used, _)| used.start > rows.start || used.end < rows.end) + }); + let splits = if let Some(boundaries) = builder.natural_splits { + projection_splits(&boundaries, rows) + } else { + let natural = projection_splits(&projection_starts, rows.clone()); + let boundaries = natural + .iter() + .map(|split| split.start) + .chain(iter::once(rows.end)) + .collect(); + let boundaries = builder.split_by.split_boundaries(boundaries, &rows)?; + projection_splits(&boundaries, rows) + }; let io = Arc::new(SegmentScanIo::new( file.segments, segment_ranges(&file.locations), @@ -178,6 +219,7 @@ impl PreparedScan { row_offset: builder.row_offset, values: Default::default(), }, + value_scopes, io, splits, selection: builder.selection, @@ -200,12 +242,17 @@ impl PreparedScan { file_ordinal: 0, rows: row_mask.row_range(), }; + let mut plans = self.plans.clone(); + if let Some(values) = &self.value_scopes { + plans.values = self.plans.values.for_plans( + values + .iter() + .filter(|(used, _)| used.start < rows.end && rows.start < used.end) + .map(|(_, plan)| plan), + ); + } Some(SplitPlan { - root: plan_projection( - self.plans.clone(), - scope.clone(), - row_mask.mask().clone(), - ), + root: plan_projection(plans, scope.clone(), row_mask.mask().clone()), scope, }) }) diff --git a/vortex-layout/src/scan/v2/splits.rs b/vortex-layout/src/scan/v2/splits.rs index 3d84f19b995..b1c5cf6c604 100644 --- a/vortex-layout/src/scan/v2/splits.rs +++ b/vortex-layout/src/scan/v2/splits.rs @@ -26,35 +26,78 @@ use crate::plan::physical::Zoned; /// /// This follows plans whose children cover the same rows, or known parts of them, and stops at any /// other plan: a take's values cover other rows, and zone statistics cover zones. +/// For a lowered identity projection, `values` also records the root rows that need each whole +/// dictionary, so its cached value can be released after those splits finish. pub(super) fn chunk_starts<'a>( plans: impl IntoIterator, + mut values: Option<&mut Vec<(Range, PlanRef)>>, ) -> VortexResult> { let mut starts = Vec::new(); for plan in plans { - collect_starts(plan, 0, &mut starts)?; + collect_starts(plan, 0, &mut starts, values.as_deref_mut())?; } starts.sort_unstable(); starts.dedup(); Ok(starts) } -fn collect_starts(plan: &PlanRef, offset: u64, starts: &mut Vec) -> VortexResult<()> { +fn collect_starts( + plan: &PlanRef, + offset: u64, + starts: &mut Vec, + mut values: Option<&mut Vec<(Range, PlanRef)>>, +) -> VortexResult<()> { if let Some(concat) = plan.as_opt::() { for (index, &child_offset) in concat.row_offsets().iter().enumerate() { starts.push(offset + child_offset); - collect_starts(&plan.child_required(index)?, offset + child_offset, starts)?; + collect_starts( + &plan.child_required(index)?, + offset + child_offset, + starts, + values.as_deref_mut(), + )?; } } else if let Some(zoned) = plan.as_opt::() { if let Some(data) = zoned.data_plan()? { - collect_starts(&data, offset, starts)?; + collect_starts(&data, offset, starts, values)?; } - } else if let Some(stats) = plan.as_opt::() { - collect_starts(&stats.data_plan()?, offset, starts)?; + } else if let Some(legacy) = plan.as_opt::() { + collect_starts(&legacy.data_plan()?, offset, starts, values)?; } else if plan.is::() { - collect_starts(&plan.child_required(0)?, offset, starts)?; + if let Some(values) = values.as_deref_mut() { + let whole = plan.child_required(1)?; + let rows = offset..offset + plan.row_count(); + values.push((rows.clone(), whole.clone())); + collect_nested_values(&whole, &rows, values)?; + } + collect_starts(&plan.child_required(0)?, offset, starts, values)?; } else if plan.is::() || plan.is::() || plan.is::() { - for child in plan.children().iter() { - collect_starts(&child?, offset, starts)?; + for child in plan.children().iter_refs() { + collect_starts(child?, offset, starts, values.as_deref_mut())?; + } + } + Ok(()) +} + +/// A dictionary's whole values can contain further dictionaries. They are needed over the +/// outer dictionary's row domain, regardless of their own row counts. +fn collect_nested_values( + plan: &PlanRef, + rows: &Range, + values: &mut Vec<(Range, PlanRef)>, +) -> VortexResult<()> { + if let Some(zoned) = plan.as_opt::() { + if let Some(data) = zoned.data_plan()? { + collect_nested_values(&data, rows, values)?; + } + } else if let Some(legacy) = plan.as_opt::() { + collect_nested_values(&legacy.data_plan()?, rows, values)?; + } else { + if plan.is::() { + values.push((rows.clone(), plan.child_required(1)?)); + } + for child in plan.children().iter_refs() { + collect_nested_values(child?, rows, values)?; } } Ok(()) @@ -62,6 +105,9 @@ fn collect_starts(plan: &PlanRef, offset: u64, starts: &mut Vec) -> VortexR /// The projection splits of `rows`: cut at every one of `starts` inside it. pub(crate) fn projection_splits(starts: &[u64], rows: Range) -> Vec> { + if rows.is_empty() { + return Vec::new(); + } let first = starts.partition_point(|&start| start <= rows.start); let last = starts.partition_point(|&start| start < rows.end); iter::once(rows.start) @@ -76,14 +122,22 @@ pub(crate) fn projection_splits(starts: &[u64], rows: Range) -> Vec, @@ -91,4 +145,31 @@ mod tests { ) { assert_eq!(projection_splits(starts, rows), expected); } + + #[test] + fn dictionary_values_use_the_outer_codes_row_domain() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, false.into()); + let segment = |id, rows| { + SegmentScanPlan::new( + dtype.clone(), + rows, + SegmentId::from(id), + ReadContext::new([]), + None, + ) + .into_plan() + }; + let inner_values = segment(0_u32, 3); + let outer_values = TakePlan::new(segment(1, 2), inner_values.clone()).into_plan(); + let dict = TakePlan::new(segment(2, 4), outer_values.clone()).into_plan(); + let plan = ConcatPlan::try_new(dtype.clone(), vec![segment(3, 5), dict])?.into_plan(); + let mut values = Vec::new(); + assert_eq!(chunk_starts([&plan], Some(&mut values))?, vec![0, 5]); + assert_eq!(values.len(), 2); + for ((rows, actual), expected) in values.iter().zip([outer_values, inner_values]) { + assert_eq!(rows, &(5..9)); + assert!(PlanRef::ptr_eq(actual, &expected)); + } + Ok(()) + } } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 9eb421cc677..6fbdca35343 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -4,6 +4,8 @@ use std::num::NonZeroUsize; use std::ops::Range; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use futures::TryStreamExt; use futures::stream; @@ -28,6 +30,7 @@ use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; @@ -38,6 +41,7 @@ use super::select_supported; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; +use crate::layouts::dict::DictLayout; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::row_idx::row_idx; use crate::layouts::table::TableStrategy; @@ -45,6 +49,9 @@ use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; use crate::scan::planning::SegmentLocation; use crate::scan::scan_builder::ScanBuilder; +use crate::scan::split_by::SplitBy; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::segments::TestSegments; use crate::sequence::SequenceId; @@ -256,6 +263,111 @@ async fn projection_matches_v1( .await } +#[tokio::test] +async fn zoned_projection_keeps_data_splits() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, zoned_table().as_ref()).await?; + let projection = get_item("ints", root()).bind(layout.dtype())?; + let Executor::V2(scan) = select_supported( + builder(&session, &segments, &layout)? + .with_projection(projection) + .with_row_range(250..3100), + )? + else { + panic!("the scan should run on the planner-driven executor"); + }; + assert_eq!( + scan.splits, + vec![250..1000, 1000..2000, 2000..3000, 3000..3100] + ); + Ok(()) +} + +#[rstest] +#[case::layout(SplitBy::Layout, vec![0..1000, 1000..2000, 2000..3000, 3000..4000])] +#[case::bounded( + SplitBy::LayoutSubSplitting { max_rows: 600 }, + vec![0..500, 500..1000, 1000..1500, 1500..2000, 2000..2500, 2500..3000, 3000..3500, 3500..4000] +)] +#[case::rows(SplitBy::RowCount(1500), vec![0..1500, 1500..3000, 3000..4000])] +#[tokio::test] +async fn split_policy_is_preserved( + #[case] policy: SplitBy, + #[case] expected: Vec>, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let Executor::V2(scan) = + select_supported(builder(&session, &segments, &layout)?.with_split_by(policy))? + else { + panic!("the scan should run on the planner-driven executor"); + }; + assert_eq!(scan.splits, expected); + Ok(()) +} + +#[tokio::test] +async fn supplied_splits_take_precedence_and_are_clamped() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let Executor::V2(scan) = select_supported( + builder(&session, &segments, &layout)? + .with_split_by(SplitBy::RowCount(1000)) + .with_natural_splits(Arc::from([0, 2500, 4000])) + .with_row_range(500..3000), + )? + else { + panic!("the scan should run on the planner-driven executor"); + }; + assert_eq!(scan.splits, vec![500..2500, 2500..3000]); + Ok(()) +} + +struct CountingSegments { + inner: Arc, + reads: AtomicUsize, +} + +impl SegmentSource for CountingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.reads.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } +} + +#[tokio::test] +async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let ints = layout.slot(1)?.ok_or_else(|| vortex_err!("missing ints"))?; + // The codes 0..4000 index the same values, so each split returns its own row numbers. + let layout = DictLayout::new(Arc::clone(&ints), ints).into_layout(); + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments: Arc = counted.clone(); + let Executor::V2(scan) = select_supported(builder(&session, &segments, &layout)?)? else { + panic!("the scan should run on the planner-driven executor"); + }; + let splits = scan.split_plans(); + assert_eq!(splits.len(), CHUNKS); + for (index, split) in splits.into_iter().enumerate() { + let before = counted.reads.load(Ordering::Relaxed); + let arrays = split.run(Arc::clone(&scan.io)).await?; + let reads = counted.reads.load(Ordering::Relaxed) - before; + assert_eq!(reads, if index == 0 { CHUNKS } else { 1 }); + let start = i32::try_from(index * CHUNK_ROWS)?; + let expected = PrimitiveArray::from_iter(start..start + i32::try_from(CHUNK_ROWS)?); + assert_arrays_eq!( + ChunkedArray::try_new(arrays, layout.dtype().clone())?.into_array(), + expected, + &mut session.create_execution_ctx() + ); + } + Ok(()) +} + #[tokio::test] async fn filters_and_limits_run_on_v1() -> VortexResult<()> { let session = new_session().with_tokio(); From 5076f3123380945a16418d6b51161ad85df4cbc5 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 16 Oct 2026 00:08:40 +0100 Subject: [PATCH 07/14] bench(scan): cover pipeline overhead and fix scan lints Measure preparation, compilation and in-memory segment delivery for narrow and wide pipeline scans in the existing CodSpeed layout shard. Document the stable identity of scoped dictionary-cache keys and use an explicit Arc clone in the V2 test so the scan code passes the CI Clippy configurations. Signed-off-by: Joe Isaacs --- vortex-layout/Cargo.toml | 4 + vortex-layout/benches/pipeline_scan.rs | 132 ++++++++++++++++++ .../src/plan/physical/pipeline/values.rs | 4 + vortex-layout/src/scan/v2/tests.rs | 2 +- 4 files changed, 141 insertions(+), 1 deletion(-) create mode 100644 vortex-layout/benches/pipeline_scan.rs diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 1b286be4671..22768d7e4f1 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -92,6 +92,10 @@ harness = false name = "aggregate_partial_merge" harness = false +[[bench]] +name = "pipeline_scan" +harness = false + [lints] workspace = true diff --git a/vortex-layout/benches/pipeline_scan.rs b/vortex-layout/benches/pipeline_scan.rs new file mode 100644 index 00000000000..db4beeaa043 --- /dev/null +++ b/vortex-layout/benches/pipeline_scan.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Isolates scan preparation, compilation, and segment delivery from Arrow conversion and IO. + +#![expect(clippy::unwrap_used)] + +use std::hint::black_box; + +use divan::Bencher; +use mimalloc::MiMalloc; +use vortex_array::ArrayContext; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::serde::SerializeOptions; +use vortex_buffer::Alignment; +use vortex_buffer::ByteBufferMut; +use vortex_layout::plan::physical::PackPlan; +use vortex_layout::plan::physical::PlanRef; +use vortex_layout::plan::physical::SegmentScanPlan; +use vortex_layout::plan::physical::pipeline::Scan; +use vortex_layout::plan::physical::pipeline::Split; +use vortex_layout::plan::physical::pipeline::Turn; +use vortex_layout::segments::SegmentId; +use vortex_layout::session::LayoutSession; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +const ROWS: u64 = 1024; + +fn main() { + divan::main(); +} + +struct Fixture { + session: VortexSession, + plan: PlanRef, + segment: BufferHandle, +} + +impl Fixture { + fn new(columns: u32) -> Self { + let session = vortex_array::array_session().with::(); + let array = PrimitiveArray::from_iter(0_i32..1024).into_array(); + let context = ArrayContext::empty(); + let buffers = array + .serialize( + &context, + &session, + &SerializeOptions { + offset: 0, + include_padding: true, + }, + ) + .unwrap(); + let mut bytes = ByteBufferMut::empty_aligned(Alignment::new(64)); + for buffer in buffers { + bytes.extend_from_slice(buffer.as_ref()); + } + let fields = StructFields::from_iter( + (0..columns).map(|index| (format!("c{index}"), array.dtype().clone())), + ); + let children = (0..columns) + .map(|index| { + SegmentScanPlan::new( + array.dtype().clone(), + ROWS, + SegmentId::from(index), + ReadContext::new(context.to_ids()), + None, + ) + .into_plan() + }) + .collect(); + Self { + session, + plan: PackPlan::try_new(fields, Nullability::NonNullable, ROWS, children, None) + .unwrap() + .into_plan(), + segment: BufferHandle::new_host(bytes.freeze()), + } + } + + fn scan(&self) -> Scan { + Scan::try_new( + self.session.clone(), + self.plan.clone(), + vec![Split::all(0..ROWS)], + ) + .unwrap() + } +} + +#[divan::bench(args = [1, 100, 1000])] +fn prepare(bencher: Bencher, columns: u32) { + let fixture = Fixture::new(columns); + bencher.bench_local(|| black_box(fixture.scan())); +} + +#[divan::bench(args = [1, 100, 1000])] +fn prepare_and_compile(bencher: Bencher, columns: u32) { + let fixture = Fixture::new(columns); + bencher.bench_local(|| { + let mut scan = fixture.scan(); + assert!(matches!(scan.step().unwrap(), Turn::Read(_))); + black_box(scan); + }); +} + +#[divan::bench(args = [1, 100, 1000])] +fn scan(bencher: Bencher, columns: u32) { + let fixture = Fixture::new(columns); + bencher.bench_local(|| { + let mut scan = fixture.scan(); + loop { + match scan.step().unwrap() { + Turn::Read(read) => scan.deliver(read.id, fixture.segment.clone()).unwrap(), + Turn::Output(_, array) => { + black_box(array); + } + Turn::Waiting => unreachable!("every read is delivered immediately"), + Turn::Done => break, + } + } + }); +} diff --git a/vortex-layout/src/plan/physical/pipeline/values.rs b/vortex-layout/src/plan/physical/pipeline/values.rs index 9b40314f279..99c226da7d2 100644 --- a/vortex-layout/src/plan/physical/pipeline/values.rs +++ b/vortex-layout/src/plan/physical/pipeline/values.rs @@ -47,6 +47,10 @@ impl ValueCache { /// last cache using its plan is dropped, rather than until every split finishes. pub(crate) fn for_plans<'a>(&self, plans: impl Iterator) -> Self { let mut entries = self.0.lock(); + #[allow( + clippy::mutable_key_type, + reason = "PlanByPtr hashes and compares only the stable Arc address" + )] let scoped = plans .map(|plan| { let key = PlanByPtr(plan.clone()); diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 6fbdca35343..f946dc49a51 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -346,7 +346,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { inner: segments, reads: AtomicUsize::new(0), }); - let segments: Arc = counted.clone(); + let segments = Arc::clone(&counted) as Arc; let Executor::V2(scan) = select_supported(builder(&session, &segments, &layout)?)? else { panic!("the scan should run on the planner-driven executor"); }; From 7b15cfc70a1e66953c1434d404a1d03e42306ec8 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 16 Oct 2026 00:22:43 +0100 Subject: [PATCH 08/14] docs(scan): fix planner documentation links Keep private implementation names as code and remove the redundant v2 link target so the planner documentation builds with warnings denied. Signed-off-by: Joe Isaacs --- vortex-layout/src/scan/planning/mod.rs | 2 +- vortex-layout/src/scan/scan_builder.rs | 2 +- vortex-layout/src/scan/v2/mod.rs | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vortex-layout/src/scan/planning/mod.rs b/vortex-layout/src/scan/planning/mod.rs index c110846a9b5..dae73e2780b 100644 --- a/vortex-layout/src/scan/planning/mod.rs +++ b/vortex-layout/src/scan/planning/mod.rs @@ -3,7 +3,7 @@ //! Layout-level planners and morsels for the planning protocol. //! -//! A split is evaluated over physical plans: a [`ProjectionPlanner`] emits a [`ProjectionMorsel`] +//! A split is evaluated over physical plans: a [`ProjectionPlanner`] emits a `ProjectionMorsel` //! per projection split, which runs the projection plan over the split's selected rows through //! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published //! as a protocol request, so the planning driver decides when IO happens. diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 590d2b644fe..8134e86d948 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -274,7 +274,7 @@ impl ScanBuilder { /// Attach the segment map of the file the layout reader was opened over. /// - /// Scans with a segment map may run on the planner-driven executor, see [`v2`](crate::scan::v2). + /// Scans with a segment map may run on the planner-driven executor, see [`v2`]. pub fn with_scan_file(mut self, file: ScanFile) -> Self { self.scan_file = Some(file); self diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 6d0a87e949b..be7ccdd53e6 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -8,7 +8,7 @@ //! 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. //! -//! [`select`] is the one place a streaming [`ScanBuilder`] picks between this executor and the +//! `select` is the one place a streaming [`ScanBuilder`] picks between this executor and the //! default one. It picks this one when `VORTEX_SCAN_V2=1`, the builder carries a [`ScanFile`], //! and the scan only uses what this executor implements so far: no filter, no limit, and a layout //! that lowers to plans with a pipeline implementation. Every other scan runs on the default From 30b65de6a8ddfd21488c7bc3565ca3c9ed59afd3 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 11:17:05 +0000 Subject: [PATCH 09/14] bench(compress): decompress Vortex on the V2 scan executor Adds ScanBuilder::with_scan_v2 to choose the planner-driven executor for one scan, overriding VORTEX_SCAN_V2, and has compress-bench's Vortex decompression call it, so CI measures V2 decompression against develop's V1 baseline without relying on an environment variable. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VtNyxw9KerLDyt5eXnpnRE Signed-off-by: Claude --- benchmarks/compress-bench/src/vortex.rs | 7 ++++++- vortex-layout/src/scan/scan_builder.rs | 11 +++++++++++ vortex-layout/src/scan/v2/mod.rs | 5 +++-- 3 files changed, 20 insertions(+), 3 deletions(-) diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index a28a06df381..0ef9fa3778e 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -68,7 +68,12 @@ impl Compressor for VortexCompressor { }; let start = Instant::now(); - let mut scan = SESSION.open_options().open_buffer(data.clone())?.scan()?; + // Decode on the planner-driven scan executor. + let mut scan = SESSION + .open_options() + .open_buffer(data.clone())? + .scan()? + .with_scan_v2(true); let source_dtype = scan.dtype()?; let root_columns = source_dtype .as_struct_fields_opt() diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 8134e86d948..b457a1eae42 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -85,6 +85,8 @@ pub struct ScanBuilder { pub(super) row_offset: u64, /// Where the scanned file's segments live, which lets the planner-driven executor read it. pub(super) scan_file: Option, + /// Overrides [`v2::enabled`] for this scan when set. + pub(super) scan_v2: Option, } impl ScanBuilder { @@ -109,6 +111,7 @@ impl ScanBuilder { limit: None, row_offset: 0, scan_file: None, + scan_v2: None, } } @@ -280,6 +283,13 @@ impl ScanBuilder { self } + /// Whether a supported scan runs on the planner-driven executor, overriding + /// [`v2::ENV_VAR`] for this scan. + pub fn with_scan_v2(mut self, enabled: bool) -> Self { + self.scan_v2 = Some(enabled); + self + } + /// The [`DType`] returned by the scan, after applying the projection. pub fn dtype(&self) -> VortexResult { Ok(self.projection.dtype().clone()) @@ -311,6 +321,7 @@ impl ScanBuilder { limit: self.limit, row_offset: self.row_offset, scan_file: self.scan_file, + scan_v2: self.scan_v2, map_fn: Arc::new(move |a| old_map_fn(a).and_then(&map_fn)), } } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index be7ccdd53e6..9fd6f01284a 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -9,7 +9,8 @@ //! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. //! //! `select` is the one place a streaming [`ScanBuilder`] picks between this executor and the -//! default one. It picks this one when `VORTEX_SCAN_V2=1`, the builder carries a [`ScanFile`], +//! default one. It picks this one when `VORTEX_SCAN_V2=1` (or +//! [`ScanBuilder::with_scan_v2`] asks for it), the builder carries a [`ScanFile`], //! and the scan only uses what this executor implements so far: no filter, no limit, and a layout //! that lowers to plans with a pipeline implementation. Every other scan runs on the default //! executor, unchanged. @@ -92,7 +93,7 @@ pub(crate) enum Executor { /// Picks the executor for `builder` and prepares the scan when it is this one. pub(crate) fn select(builder: ScanBuilder) -> VortexResult> { - if !enabled() { + if !builder.scan_v2.unwrap_or_else(enabled) { return Ok(Executor::V1(builder)); } select_supported(builder) From 3b36128a7658b4e86f7f97ce9c992733cff33dc8 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 13:00:56 +0000 Subject: [PATCH 10/14] docs(scan): remove the planner decode-slice handoff doc Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VtNyxw9KerLDyt5eXnpnRE Signed-off-by: Claude --- vortex-scan/design/planner-decode-slice.md | 189 --------------------- 1 file changed, 189 deletions(-) delete mode 100644 vortex-scan/design/planner-decode-slice.md diff --git a/vortex-scan/design/planner-decode-slice.md b/vortex-scan/design/planner-decode-slice.md deleted file mode 100644 index 6c3d4fa926d..00000000000 --- a/vortex-scan/design/planner-decode-slice.md +++ /dev/null @@ -1,189 +0,0 @@ - - - -# Handoff: minimal planner-driven decode slice - -Status: handoff for a new implementation session. Verified against `origin/develop` at -`18aebeaac8` (which includes #10462) and source commit `42b29cc459`. - -## Goal - -Make the `vortex-scan` planning driver (`Run` / `Planner` / `Morsel`) execute develop's sans-IO -`pipeline::Scan` decode/read path, with all IO going through the planning protocol -(`IoRequest` / `Completion`), as the first mergeable step towards planner-driven V2 scans. - -Source material: - -- `42b29cc459`, PR [#10464](https://github.com/vortex-data/vortex/pull/10464) "feat(scan): - planner-driven V2 scans on the physical plan pipelines" (+11367 lines). Not an ancestor of - develop; its merge base is `1a9ed685d1`. -- Full-change reference: PR [#10466](https://github.com/vortex-data/vortex/pull/10466), branch - `origin/ji/vibrant-sagan-2w7pov`. -- PR [#10462](https://github.com/vortex-data/vortex/pull/10462) (Zoned pruning, `Share`, - `ValueCache`) is **already merged into develop** as `18aebeaac8`. -- PR [#10463](https://github.com/vortex-data/vortex/pull/10463) (`3c37da1a2a`, coalescing IO - service, `IoIntent::Forget`, `vortex_io::request::trace`) is **not** on develop and not needed. - -## Scope - -In scope: driver, a protocol adapter from `pipeline::Scan` to `IoRequest`s, a projection-only -planner, segment IO via the existing `SegmentSource`, and an opt-in hook behind -`VORTEX_SCAN_V2=1` for scans with no filter and no list columns. - -Non-goals (later PRs): filter evaluation (`FilterPlanner`, conjunct ordering), zone pruning in -the planner, IO announcements and prefetch, the coalescing IO service, plan caching, repeated -scans, DuckDB integration, DataFusion diagnostics, and any performance work. - -## What develop already has - -| Piece | Location on develop | -| --- | --- | -| `WorkScope`, `State`, `PlannerOutput`, `Planner` | `vortex-scan/src/planning/planner.rs` (81 lines) | -| `Morsel`, `MorselOutput` | `vortex-scan/src/planning/morsel.rs` (54) | -| `Next`, `next_fn` | `vortex-scan/src/planning/next.rs` (74); module root `planning/mod.rs` (11) | -| `IoRequest`, `IoBatch`, `IoIntent::{Fetch,Prefetch,Announce}`, `IoTarget`, `IoResult`, `Completion`, `IoSource`, `IoService`, `IoConsumer` | `vortex-io/src/request/mod.rs` (223); no `Forget`, no `trace` module | -| `pipeline::Scan` with `try_new`, `with_value_cache` (L712), `step() -> Turn::{Read,Output,Waiting,Done}` (L718), `deliver(ReadId, BufferHandle)` (L755) | `vortex-layout/src/plan/physical/pipeline/scan.rs`; `Turn` at L82 | -| `ValueCache` | `vortex-layout/src/plan/physical/pipeline/values.rs` L22 | -| `lower(&LayoutRef) -> PlanRef`, public, handles Flat/Chunked/Struct/Dict/List/Zoned/LegacyStats | `vortex-layout/src/plan/physical/lower.rs` L50 | -| `optimize(PlanRef)` | `vortex-layout/src/plan/physical/optimize.rs` L28 | -| `plan_row_idx_expression` | `vortex-layout/src/plan/physical/plans/row_idx.rs` L135 | -| `attempt_split_ranges` | `vortex-layout/src/scan/splits.rs` L33 | - -## What lives only in `42b29cc459` - -Driver (`vortex-scan`): - -- `src/planning/driver/mod.rs` (644): `RootId` L56, `Batch` L59, `Progress` L69, `Run` L194 with - `new` L222, `admit` L241, `cancel` L262, `advance` L282, `complete` L331, `poll_completion` - L372; blocking `Driver` L600-644 (`run` L624). -- `src/planning/driver/trace.rs` (75), `driver/tests.rs` (814), `planning/tests/mod.rs` (6), - `planning/tests/scripted.rs` (418). -- `planner.rs` +12 / `morsel.rs` +6: only `trace_name()` / `trace_selection()` default methods, - used solely by tracing. -- `Cargo.toml` +6: `vortex-utils` dependency; dev-deps `parking_lot`, `rstest`, - `vortex-array` (`_test-harness`). - -Hidden dependency on #10463: `driver/mod.rs` L41 imports `vortex_io::request::trace::timestamp_ns` -and L254 calls `IoSource::trace_id()`; `trace.rs` imports `trace::{next_id, timestamp_ns}`. - -## PR plan - -Extract each file with `git show 42b29cc459: > ` and then trim as described. - -### PR A: driver in `vortex-scan` (~600 + ~1.2k tests) - -Standalone on develop. - -1. Copy `vortex-scan/src/planning/driver/{mod.rs,tests.rs}`, - `vortex-scan/src/planning/tests/{mod.rs,scripted.rs}`; add `pub mod driver;` and - `#[cfg(test)] mod tests;` to `planning/mod.rs`. -2. Strip tracing: drop `trace.rs`, the `trace` field on `Run` (L195) and its setup in `new` - (L224), every `tracing::debug!(ts_ns = timestamp_ns(), ...)` block (admit, cancel, completion, - selection ~L499-530), and the `timestamp_ns` import. About 60 lines. Do not add - `trace_name`/`trace_selection` to `Planner`/`Morsel`. - Alternative: cherry-pick the ~40 lines of `vortex-io` from `3c37da1a2a` - (`request/trace.rs` +24, `request/mod.rs` `trace_id()`); not recommended for this slice. -3. Add the `Cargo.toml` deps listed above. -4. Tests: `driver/tests.rs` and `scripted.rs` exercise the driver with scripted planners and a - fake `IoSource`; drop any test that asserts trace output. - -### PR B: protocol adapter in `vortex-layout` (~700 + ~300 tests) - -Depends on A. New `vortex-layout/src/scan/planning/` and `vortex-layout/src/scan/v2/` modules. - -- `scan/planning/graph.rs` (257): `ScanGraph` wraps `pipeline::Scan`; `ProtocolGraph` maps - `Turn::Read` to a `Fetch` `IoRequest` (L207-216) and delivers `IoResult`s back. Keep - `.with_value_cache(values)` (L73) and the `ValueCache` import (L27): they compile on develop - now that #10462 landed. -- `scan/planning/projection.rs` (568): keep `ProjectionPlanner` with `new` and `cut` (L104) and - `ProjectionMorsel` (L251-~362). Trim `compute` (L202) to: take `selected`, `cut`, pop a split, - emit `PlannerOutput::Morsel`. Drop `prefetch` (L70), `forget_pruned` (L131, needs - `IoIntent::Forget` from #10463), the `VORTEX_SCAN_IO_FORGET_PRUNED_PROJECTION` switch, the - `plan_segments` import from `v2::prefetch`, and the tests at L382+ that cover them. -- `scan/planning/mod.rs` (171): keep `SegmentLocation` (L44) and `ScanPlans` (L66, **including** - `values: ValueCache`). Move `SelectedRows` here from `filter.rs` L62. Drop `announce`, - `filter`, `plan_split`/`plan_sparse_projection` (they wrap `AnnouncePlanner`); replace with a - `plan_projection(plans, scope, mask) -> Box` returning `ProjectionPlanner`. -- `scan/v2/io.rs` (321): keep `segment_ranges`, `SegmentScanIo`, `SegmentIoSource` (L36-212), - adapting `SegmentSource` to `IoService`. Delete the `IoIntent::Forget` branch (L111). Keep the - two tests (L213+), rewriting any that build `Forget` requests. -- `scan/v2/split.rs` (90): keep `SplitPlan::run`, the `admit` / `advance` / - `poll_completion` loop (~L76-85). Drop `EARLY_ANNOUNCE` (L44-60). -- `scan/v2/splits.rs` (343): keep `projection_splits` (L252), `chunk_starts` (L135) / - `collect_starts`, `layout_chunk_starts` (L54) / `collect_layout_starts`, `cut_at` (~110 lines). - Drop `max_split_rows`, `MAX_SPARSE_ROWS`, the filter/sparse boundary functions. -- `scan/v2/lower.rs` (207): **drop entirely**. Develop's `plan::physical::lower` now handles - Zoned/LegacyStats, lowering them to their data when unoptimized. -- Add `pub(crate) mod planning; pub mod v2;` to `scan/mod.rs`. -- Tests: an in-memory `SegmentSource` plus layouts built like - `vortex-layout/src/plan/physical/pipeline/tests.rs` (`lower(&layout)`), driven by - `Driver::run` from PR A; compare with `assert_arrays_eq!` against the default scan. - -### PR C: opt-in executor behind a single fork - -The source commit forked at every caller (`if v2::enabled() && v2::supports_scan(..) { -v2::into_stream(..) } else { b.into_stream() }` in the DataFusion opener, `vortex-file`, DuckDB). -This slice forks in exactly one place instead, and nothing V2-specific is added to the V1 path: - -- `ScanBuilder`'s lazy prepare step calls `prepare_stream`, which is the only fork: - `match v2::select(self)? { Executor::V1(b) => b.prepare()?.execute_stream(None), - Executor::V2(scan) => scan.execute_stream() }`. `into_stream`, `into_iter` and the Arrow - record-batch stream all go through it, so no caller changes. -- `v2::select` owns the whole gate: `VORTEX_SCAN_V2=1`, the builder carries a `ScanFile`, no - filter, no limit, and every layout is Flat/Chunked/Struct/Dict/Zoned/LegacyStats (no List - layout; lists written as flat leaves, the default, are supported). It logs the decision at - `debug` on target `vortex_layout::scan::v2`. Each later PR widens this gate in one place. -- `VortexFile::scan()` attaches `ScanFile` via `ScanBuilder::with_scan_file`, built from - `Footer::segment_specs_with_metadata`. Callers that build a `ScanBuilder` from a bare layout - reader (DataFusion opener, DuckDB) stay on V1 until they attach one with the same setter: wiring, - not a fork. -- `ScanBuilder::prepare()` / `build()` keep returning the V1 `RepeatedScan`. - -### PR D (optional): more callers - -- `vortex-datafusion/src/persistent/opener.rs`: add `.with_scan_file(vxf.scan_file())` to the - builder chain. No `if`. - -### Later - -1. `FilterPlanner` (`scan/planning/filter.rs` 511 + `filter/tests.rs` 190) and - `scan/v2/conjuncts.rs` (226); lift the no-filter gate. -2. Zone pruning in the planner (`scan/v2/pruning.rs` 275; `ZonedPlan` is already on develop). -3. #10463 coalescing service, `FileScanIo` (`vortex-file/src/segments/{io_service,scan_io}.rs`), - `IoIntent::Forget`, tracing. -4. `announce.rs` (409), `prefetch.rs` (145), `share.rs` (147), split heuristics (sparse, - `max_split_rows`), `file.rs` plan caching (102), `repeated_scan.rs` (606). -5. `vortex-file/src/planning/{footer_open,footer_prune,tests}` (~1.8k incl. fixtures). -6. DuckDB (`file_prefetch.rs`, `ffi.rs`, C++ reader), DataFusion diagnostics. - -## Size - -Slice: ~3.3k lines including tests (~1.85k non-test). Full stack: ~16.3k (#10462 +1622, -#10463 +3301, #10464 +11367); #10462 is already merged. - -## Verification - -Run only when requested, scoped per crate (`vortex-scan`, `vortex-layout`, `vortex-file`, -`vortex-datafusion`): - -```bash -cargo clippy -p --all-targets --all-features -- -D warnings -cargo +nightly-2026-09-10 fmt -p # NIGHTLY_TOOLCHAIN in .github/workflows/ci.yml -cargo nextest run -p -``` - -Then a parity check with `VORTEX_SCAN_V2=1` on an unfiltered DataFusion query. - -## Open risks - -- **CI regression, cause unknown.** On CI's 96-core ARM runners the full V2 stack regressed - (Vortex queries up to 12x, FineWeb vortex-compact 1.25-1.8x, TPC-H SF10 q06 +28%) but not on - 4-core local runs. Per-split IO locks and the file read driver were ruled out. This slice is a - good base to bisect it on CI. -- **Local q01 +10%:** extra decimal to Arrow conversion and bit-packed unpack versus V1's - filter-on-compressed path. Expect the slice to be slower than V1 until filters land. -- `SegmentScanIo` issues one `SegmentSource` request per segment with no coalescing. -- Develop's `plan/physical` changed since the source commit's merge base (`1a9ed685d1`, ~2.2k - lines via #10462). `Scan::try_new(session, plan, splits)` is unchanged, but recheck imports - such as `pipeline::Split` and `ValueCache` when porting `graph.rs`. -- List columns cannot compile before their offsets are read; such files stay on V1. From 972aecbde188f3c998a1b84eef36c2a22b139d09 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 16 Oct 2026 14:16:54 +0100 Subject: [PATCH 11/14] refactor(scan): select V2 through a separate builder Use an explicit V2 ScanBuilder import for benchmark decompression and restore the V1 builder. Remove runtime executor selection and fallback, restrict planning and graph helper visibility, and cover the direct stream and Arrow paths. Signed-off-by: Joe Isaacs --- benchmarks/compress-bench/src/vortex.rs | 13 +- vortex-file/src/file.rs | 10 +- vortex-layout/src/scan/arrow.rs | 2 +- vortex-layout/src/scan/mod.rs | 2 +- vortex-layout/src/scan/planning/graph.rs | 22 +- vortex-layout/src/scan/planning/mod.rs | 30 +-- vortex-layout/src/scan/planning/projection.rs | 10 +- vortex-layout/src/scan/scan_builder.rs | 62 +---- vortex-layout/src/scan/v2/mod.rs | 224 +++++++++++++----- vortex-layout/src/scan/v2/splits.rs | 2 +- vortex-layout/src/scan/v2/tests.rs | 156 ++++++------ 11 files changed, 303 insertions(+), 230 deletions(-) diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index 0ef9fa3778e..392a34e4755 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -19,6 +19,7 @@ use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; +use vortex::layout::scan::v2::ScanBuilder; use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; @@ -68,13 +69,9 @@ impl Compressor for VortexCompressor { }; let start = Instant::now(); - // Decode on the planner-driven scan executor. - let mut scan = SESSION - .open_options() - .open_buffer(data.clone())? - .scan()? - .with_scan_v2(true); - let source_dtype = scan.dtype()?; + let file = SESSION.open_options().open_buffer(data.clone())?; + let mut scan = ScanBuilder::new(SESSION.clone(), file.scan_file()); + let source_dtype = scan.dtype().clone(); let root_columns = source_dtype .as_struct_fields_opt() .map_or(0, |fields| fields.nfields()); @@ -86,7 +83,7 @@ impl Compressor for VortexCompressor { .optimize_recursive()?; scan = scan.with_projection(projection); } - let schema = Arc::new(SESSION.arrow().to_arrow_schema(&scan.dtype()?)?); + let schema = Arc::new(SESSION.arrow().to_arrow_schema(scan.dtype())?); let stream = scan.into_record_batch_stream(schema)?; pin_mut!(stream); diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 50aa797c4a5..9444d8cb1a5 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -217,14 +217,16 @@ impl VortexFile { /// Initiate a scan of the file, returning a builder for projection, filtering, selection, and /// execution options. pub fn scan(&self) -> VortexResult> { - Ok( - ScanBuilder::new(self.session.clone(), self.layout_reader()?) - .with_scan_file(self.scan_file()), - ) + Ok(ScanBuilder::new( + self.session.clone(), + self.layout_reader()?, + )) } /// The parts of this file the planner-driven scan executor reads through: its root layout, /// where each of its segments lives, and its segment source. + /// + /// Pass this to [`vortex_layout::scan::v2::ScanBuilder::new`] to build a V2 scan explicitly. pub fn scan_file(&self) -> ScanFile { let locations = self .footer diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 224035a45fe..b128b92b452 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -64,7 +64,7 @@ impl ScanBuilder { } } -fn to_record_batch( +pub(super) fn to_record_batch( chunk: ArrayRef, data_type: &Field, ctx: &mut ExecutionCtx, diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index c8a02ed74d9..fcf15137e06 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -6,7 +6,7 @@ mod filter; pub mod layout; mod limit; pub mod multi; -pub(crate) mod planning; +mod planning; pub mod repeated_scan; pub mod scan_builder; pub mod split_by; diff --git a/vortex-layout/src/scan/planning/graph.rs b/vortex-layout/src/scan/planning/graph.rs index 679fc6c6e44..e5fa1a3664c 100644 --- a/vortex-layout/src/scan/planning/graph.rs +++ b/vortex-layout/src/scan/planning/graph.rs @@ -30,17 +30,17 @@ use crate::segments::SegmentId; /// Rows a running plan produced, in row order. #[derive(Clone)] -pub(crate) struct Piece { +pub(super) struct Piece { /// Orders the piece among the pieces of one graph. - pub(crate) rows: Range, + pub(super) rows: Range, /// The selected rows the piece holds. - pub(crate) array: ArrayRef, + pub(super) array: ArrayRef, } /// A plan running over one row range on the pipeline [`Scan`]. /// /// The scan produces the selected rows in row order and leaves its reads to its owner. -pub(crate) struct ScanGraph { +pub(super) struct ScanGraph { scan: Scan, /// A step the scan already took that the owner has not seen. pending: Option, @@ -60,7 +60,7 @@ impl ScanGraph { /// Compiles nothing yet: the scan compiles `plan` over `rows`, restricted to `mask`, on its /// first step. `row_offset` is the global row index of the plan's first row, and `values` the /// whole values shared with the split's other scans. - pub(crate) fn try_new( + pub(super) fn try_new( session: VortexSession, plan: &PlanRef, rows: Range, @@ -119,7 +119,7 @@ impl ScanGraph { } /// What one [`ProtocolGraph::compute`] produced. -pub(crate) enum GraphStep { +pub(super) enum GraphStep { /// Control returns to the owner with CPU work still ready. Yield, /// Requests to publish, each returned once. @@ -133,7 +133,7 @@ pub(crate) enum GraphStep { /// The scan names segments by id; the protocol names byte ranges. Each scan read becomes a /// `Fetch` of the segment's range at the segment's alignment, and each delivery is handed back to /// the scan under its own id. -pub(crate) struct ProtocolGraph { +pub(super) struct ProtocolGraph { graph: ScanGraph, locations: Arc<[SegmentLocation]>, /// Published and undelivered requests, with the scan read and segment each one answers. @@ -147,7 +147,7 @@ pub(crate) struct ProtocolGraph { impl ProtocolGraph { /// Wraps `graph`, numbering its requests from `first_id`. Request ids must be unique over an /// owner's lifetime. - pub(crate) fn new(graph: ScanGraph, locations: Arc<[SegmentLocation]>, first_id: u32) -> Self { + pub(super) fn new(graph: ScanGraph, locations: Arc<[SegmentLocation]>, first_id: u32) -> Self { Self { graph, locations, @@ -158,7 +158,7 @@ impl ProtocolGraph { } /// The owner's protocol state while the graph is running. - pub(crate) fn state(&self) -> State { + pub(super) fn state(&self) -> State { if self.failed.is_some() { return State::NeedsCompute; } @@ -169,7 +169,7 @@ impl ProtocolGraph { } } - pub(crate) fn compute(&mut self) -> VortexResult { + pub(super) fn compute(&mut self) -> VortexResult { if let Some(error) = self.failed.take() { return Err(error); } @@ -214,7 +214,7 @@ impl ProtocolGraph { /// Hands a delivered read to the scan. The driver only delivers requests this graph /// published, so anything else is a driver bug. Bytes the scan cannot take fail the owner's /// next compute. - pub(crate) fn set_io_result(&mut self, id: request::IoRequestId, result: IoResult) { + pub(super) fn set_io_result(&mut self, id: request::IoRequestId, result: IoResult) { let Some((read, segment_id)) = self.outstanding.remove(&id.0) else { vortex_panic!("ProtocolGraph: delivery of {id:?}, which is not outstanding"); }; diff --git a/vortex-layout/src/scan/planning/mod.rs b/vortex-layout/src/scan/planning/mod.rs index dae73e2780b..e95180ab1da 100644 --- a/vortex-layout/src/scan/planning/mod.rs +++ b/vortex-layout/src/scan/planning/mod.rs @@ -3,17 +3,16 @@ //! Layout-level planners and morsels for the planning protocol. //! -//! A split is evaluated over physical plans: a [`ProjectionPlanner`] emits a `ProjectionMorsel` +//! A split is evaluated over physical plans: a `ProjectionPlanner` emits a `ProjectionMorsel` //! per projection split, which runs the projection plan over the split's selected rows through //! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published //! as a protocol request, so the planning driver decides when IO happens. -pub(crate) mod graph; +mod graph; mod projection; use std::sync::Arc; -pub use projection::ProjectionPlanner; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; use vortex_mask::Mask; @@ -21,6 +20,7 @@ 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::ValueCache; @@ -37,7 +37,7 @@ pub struct SegmentLocation { impl SegmentLocation { /// The byte range a read of the segment asks for. - pub fn target(&self) -> IoTarget { + fn target(&self) -> IoTarget { IoTarget::Range { offset: self.offset, len: self.length as usize, @@ -48,33 +48,33 @@ impl SegmentLocation { /// What every split of a scan shares: the projection plan and where its segments live. #[derive(Clone)] -pub struct ScanPlans { +pub(super) struct ScanPlans { /// The session used for decoding and expression evaluation. - pub session: VortexSession, + pub(super) session: VortexSession, /// The location of every segment, indexed by segment id. - pub locations: Arc<[SegmentLocation]>, + pub(super) locations: Arc<[SegmentLocation]>, /// The optimized projection plan, over the scan's root row domain. - pub projection: PlanRef, + pub(super) projection: PlanRef, /// Where the chunks the projection reads start, over the root row domain, sorted. A split's /// projection is cut at each of these, so every projection split reads at most one chunk of /// every column. - pub projection_starts: Arc<[u64]>, + pub(super) projection_starts: Arc<[u64]>, /// The global row index of the root row domain's first row. - pub row_offset: u64, + pub(super) row_offset: u64, /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans /// always read the same complete row domain, independently of a split's selection. - pub values: ValueCache, + pub(super) values: ValueCache, } /// The rows of a split selected for projection. -pub struct SelectedRows { +struct SelectedRows { /// The split the rows belong to. - pub scope: WorkScope, + scope: WorkScope, /// Which rows of `scope.rows` are selected. - pub mask: Mask, + mask: Mask, } /// The work for one split: project the rows of `scope` selected by `mask`. -pub fn plan_projection(plans: ScanPlans, scope: WorkScope, mask: Mask) -> Box { +pub(super) fn plan_projection(plans: ScanPlans, scope: WorkScope, mask: Mask) -> Box { Box::new(ProjectionPlanner::new(plans, SelectedRows { scope, mask })) } diff --git a/vortex-layout/src/scan/planning/projection.rs b/vortex-layout/src/scan/planning/projection.rs index 22930a4e996..70cd2e59ca0 100644 --- a/vortex-layout/src/scan/planning/projection.rs +++ b/vortex-layout/src/scan/planning/projection.rs @@ -28,9 +28,9 @@ use crate::scan::v2::splits::projection_splits; /// Turns a split's selected rows into the morsels that project them. /// /// The rows are cut into projection splits where the projection's chunks start, and each split -/// with a selected row becomes one [`ProjectionMorsel`], so a morsel reads at most one chunk of +/// with a selected row becomes one `ProjectionMorsel`, so a morsel reads at most one chunk of /// every column. The planner finishes once it has handed out every morsel. -pub struct ProjectionPlanner { +pub(super) struct ProjectionPlanner { plans: Option, /// The split's selected rows, until the first compute cuts them. selected: Option, @@ -40,7 +40,7 @@ pub struct ProjectionPlanner { impl ProjectionPlanner { /// Creates a planner for `selected`. - pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { Self { plans: Some(plans), selected: Some(selected), @@ -113,7 +113,7 @@ impl Planner for ProjectionPlanner { /// Runs the projection plan over a split's selected rows and returns them as one batch, in row /// order. A selection that projects to no rows finishes without a batch. -pub struct ProjectionMorsel { +struct ProjectionMorsel { plans: Option, dtype: DType, selected: SelectedRows, @@ -124,7 +124,7 @@ pub struct ProjectionMorsel { impl ProjectionMorsel { /// Creates a morsel projecting `selected`. - pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + fn new(plans: ScanPlans, selected: SelectedRows) -> Self { Self { dtype: plans.projection.dtype().clone(), plans: Some(plans), diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index b457a1eae42..814c0759453 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -42,9 +42,6 @@ use crate::scan::repeated_scan::RepeatedScan; use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; use crate::scan::splits::attempt_split_ranges; -use crate::scan::v2; -use crate::scan::v2::Executor; -use crate::scan::v2::ScanFile; /// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs. /// @@ -57,36 +54,32 @@ use crate::scan::v2::ScanFile; /// Projection and filter expressions must be bound against the reader dtype. Work is divided by /// the configured [`SplitBy`] strategy or by explicit selection ranges. pub struct ScanBuilder { - pub(super) session: VortexSession, + session: VortexSession, layout_reader: LayoutReaderRef, - pub(super) projection: BoundExpression, - pub(super) filter: Option, + projection: BoundExpression, + filter: Option, /// Whether the scan needs to return splits in the order they appear in the file. - pub(super) ordered: bool, + ordered: bool, /// Optionally read a subset of the rows in the file. - pub(super) row_range: Option>, + row_range: Option>, /// The selection mask to apply to the selected row range. // TODO(joe): replace this is usage of row_id selection, see - pub(super) selection: Selection, + selection: Selection, /// How to split the file for concurrent processing. - pub(super) split_by: SplitBy, + split_by: SplitBy, /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare) /// uses them instead of walking the layout. - pub(super) natural_splits: Option>, + natural_splits: Option>, /// The number of splits to make progress on concurrently **per-thread**. - pub(super) concurrency: usize, + concurrency: usize, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. - pub(super) map_fn: Arc VortexResult + Send + Sync>, + map_fn: Arc VortexResult + Send + Sync>, metrics_registry: Option>, /// Maximal number of rows to read after filtering. - pub(super) limit: Option, + limit: Option, /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression, /// but not by the scan [`Selection`] which remains relative. - pub(super) row_offset: u64, - /// Where the scanned file's segments live, which lets the planner-driven executor read it. - pub(super) scan_file: Option, - /// Overrides [`v2::enabled`] for this scan when set. - pub(super) scan_v2: Option, + row_offset: u64, } impl ScanBuilder { @@ -110,8 +103,6 @@ impl ScanBuilder { metrics_registry: None, limit: None, row_offset: 0, - scan_file: None, - scan_v2: None, } } @@ -275,21 +266,6 @@ impl ScanBuilder { self } - /// Attach the segment map of the file the layout reader was opened over. - /// - /// Scans with a segment map may run on the planner-driven executor, see [`v2`]. - pub fn with_scan_file(mut self, file: ScanFile) -> Self { - self.scan_file = Some(file); - self - } - - /// Whether a supported scan runs on the planner-driven executor, overriding - /// [`v2::ENV_VAR`] for this scan. - pub fn with_scan_v2(mut self, enabled: bool) -> Self { - self.scan_v2 = Some(enabled); - self - } - /// The [`DType`] returned by the scan, after applying the projection. pub fn dtype(&self) -> VortexResult { Ok(self.projection.dtype().clone()) @@ -320,8 +296,6 @@ impl ScanBuilder { metrics_registry: self.metrics_registry, limit: self.limit, row_offset: self.row_offset, - scan_file: self.scan_file, - scan_v2: self.scan_v2, map_fn: Arc::new(move |a| old_map_fn(a).and_then(&map_fn)), } } @@ -389,16 +363,6 @@ impl ScanBuilder { )) } - /// Prepares the scan and builds its stream on the executor [`v2::select`] picks. - /// - /// This is the only place a streaming scan chooses between executors. - fn prepare_stream(self) -> VortexResult>> { - match v2::select(self)? { - Executor::V1(builder) => builder.prepare()?.execute_stream(None), - Executor::V2(scan) => scan.execute_stream(), - } - } - /// Constructs a task per row split of the scan, returned as a vector of futures. /// /// Scans with row limits must use [`Self::into_stream`] or [`Self::into_iter`] so the scan @@ -462,7 +426,7 @@ impl Stream for LazyScanStream { // task: poll_next must only wait for and poll an already-constructed stream. // This also keeps construction errors on the Preparing -> Error path rather // than running construction on the caller's executor. - let task = handle.spawn_cpu(move || builder.prepare_stream()); + let task = handle.spawn_cpu(move || builder.prepare()?.execute_stream(None)); self.state = LazyScanState::Preparing(task); } LazyScanState::Preparing(task) => match ready!(Pin::new(task).poll(cx)) { diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 9fd6f01284a..4b2ee26175e 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -8,36 +8,40 @@ //! 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. //! -//! `select` is the one place a streaming [`ScanBuilder`] picks between this executor and the -//! default one. It picks this one when `VORTEX_SCAN_V2=1` (or -//! [`ScanBuilder::with_scan_v2`] asks for it), the builder carries a [`ScanFile`], -//! and the scan only uses what this executor implements so far: no filter, no limit, and a layout -//! that lowers to plans with a pipeline implementation. Every other scan runs on the default -//! executor, unchanged. +//! 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. mod io; mod split; -pub(crate) mod splits; +pub(super) mod splits; #[cfg(test)] mod tests; -use std::env; use std::iter; use std::ops::Range; use std::sync::Arc; -use std::sync::LazyLock; +use arrow_array::RecordBatch; +use arrow_schema::ArrowError; +use arrow_schema::Field; +use arrow_schema::SchemaRef; +use futures::Stream; use futures::StreamExt; use futures::TryStreamExt; 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::VortexResult; -use vortex_error::vortex_err; +use vortex_error::vortex_ensure; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; 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; @@ -52,10 +56,11 @@ use crate::plan::physical::PlanRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; +use crate::scan::arrow::to_record_batch; use crate::scan::planning::ScanPlans; pub use crate::scan::planning::SegmentLocation; use crate::scan::planning::plan_projection; -use crate::scan::scan_builder::ScanBuilder; +use crate::scan::split_by::SplitBy; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; @@ -63,19 +68,10 @@ use crate::scan::v2::splits::chunk_starts; use crate::scan::v2::splits::projection_splits; use crate::segments::SegmentSource; -/// Set to `1` to run supported scans on the planner-driven executor. -pub const ENV_VAR: &str = "VORTEX_SCAN_V2"; - -/// Whether supported scans run on the planner-driven executor. -pub fn enabled() -> bool { - static ENABLED: LazyLock = LazyLock::new(|| env::var(ENV_VAR).is_ok_and(|v| v == "1")); - *ENABLED -} - /// The parts of a file the planner-driven executor reads through. #[derive(Clone)] pub struct ScanFile { - /// The root layout the scan builder's reader was opened over. + /// The root layout to scan. pub layout: LayoutRef, /// The location of every segment, indexed by segment id. pub locations: Arc<[SegmentLocation]>, @@ -83,43 +79,157 @@ pub struct ScanFile { pub segments: Arc, } -/// The executor a streaming scan runs on. -pub(crate) enum Executor { - /// The default executor, over the layout reader. - V1(ScanBuilder), - /// The planner-driven executor, prepared. - V2(PreparedScan), +/// 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. +pub struct ScanBuilder { + session: VortexSession, + file: ScanFile, + projection: BoundExpression, + ordered: bool, + row_range: Option>, + selection: Selection, + split_by: SplitBy, + natural_splits: Option>, + concurrency: usize, + map_fn: Arc VortexResult + Send + Sync>, + row_offset: u64, } -/// Picks the executor for `builder` and prepares the scan when it is this one. -pub(crate) fn select(builder: ScanBuilder) -> VortexResult> { - if !builder.scan_v2.unwrap_or_else(enabled) { - return Ok(Executor::V1(builder)); +impl ScanBuilder { + /// Creates a V2 scan over the file's layout and segment source. + pub fn new(session: VortexSession, file: ScanFile) -> Self { + Self { + projection: BoundExpression::new_root(file.layout.dtype().clone()), + session, + file, + ordered: true, + row_range: None, + selection: Selection::default(), + split_by: SplitBy::default(), + natural_splits: None, + 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)))) } - select_supported(builder) } -/// Prepares `builder` on this executor when it supports the scan. -fn select_supported(builder: ScanBuilder) -> VortexResult> { - let declined = if builder.filter.is_some() { - Some("filter") - } else if builder.limit.is_some() { - Some("limit") - } else { - match &builder.scan_file { - None => Some("no scan file"), - Some(file) if !supports_layout(&file.layout, builder.projection.is_root())? => { - Some("layout") - } - Some(_) => None, +impl ScanBuilder { + /// The output dtype after projection, before any output mapping. + pub fn dtype(&self) -> &DType { + self.projection.dtype() + } + + /// Sets a projection bound against the file's dtype. + pub fn with_projection(mut self, projection: BoundExpression) -> Self { + self.projection = projection; + 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); + self + } + + /// Applies a row selection within the scan's row range. + pub fn with_selection(mut self, selection: Selection) -> Self { + self.selection = selection; + self + } + + /// Selects rows by strictly sorted file-local indices. + pub fn with_row_indices(self, row_indices: StrictSortedBuffer) -> Self { + self.with_selection(Selection::IncludeByIndex(row_indices)) + } + + /// Sets the root row offset used by row-index expressions. + pub fn with_row_offset(mut self, row_offset: u64) -> Self { + self.row_offset = row_offset; + self + } + + /// Configures how natural scan work is split for concurrency. + pub fn with_split_by(mut self, split_by: SplitBy) -> Self { + self.split_by = split_by; + self + } + + /// Supplies strictly increasing full-file boundaries, overriding the split policy. + /// Boundaries are clamped to the scan's row range. + pub fn with_natural_splits(mut self, boundaries: Arc<[u64]>) -> Self { + debug_assert!( + boundaries.windows(2).all(|w| w[0] < w[1]), + "natural split boundaries must be strictly increasing" + ); + self.natural_splits = Some(boundaries); + self + } + + /// Sets the number of active splits 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. + 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, + ordered: self.ordered, + row_range: self.row_range, + selection: self.selection, + split_by: self.split_by, + natural_splits: self.natural_splits, + concurrency: self.concurrency, + row_offset: self.row_offset, + map_fn: Arc::new(move |array| old_map_fn(array).and_then(&map_fn)), } - }; - if let Some(reason) = declined { - tracing::debug!(target: "vortex_layout::scan::v2", reason, "scan runs on v1"); - return Ok(Executor::V1(builder)); } - tracing::debug!(target: "vortex_layout::scan::v2", "scan runs on v2"); - Ok(Executor::V2(PreparedScan::try_new(builder)?)) + + /// 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 || PreparedScan::try_new(self)?.execute_stream()) + .await + }) + .try_flatten() + .boxed()) + } } /// Whether every required layout lowers to plans with a pipeline implementation. @@ -151,7 +261,7 @@ fn supports_layout(layout: &LayoutRef, identity: bool) -> VortexResult { } /// A scan planned over its file, ready to run its splits. -pub(crate) struct PreparedScan { +struct PreparedScan { session: VortexSession, plans: ScanPlans, /// Whole dictionary values and the root rows that use them, for an identity projection. @@ -170,9 +280,11 @@ impl PreparedScan { /// Lowers the file's layout, plans the projection over it, and cuts the scan's row range /// where the projection's chunks start. fn try_new(builder: ScanBuilder) -> VortexResult { - let file = builder - .scan_file - .ok_or_else(|| vortex_err!("a V2 scan needs a scan file"))?; + let file = builder.file; + vortex_ensure!( + supports_layout(&file.layout, builder.projection.is_root())?, + "V2 scans do not support this file's layout" + ); let root = lower(&file.layout)?; let row_count = root.row_count(); // An identity projection needs no rewrites. Keep the lowered children lazy, including @@ -264,7 +376,7 @@ impl PreparedScan { /// /// 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. - pub(crate) fn execute_stream(self) -> VortexResult>> { + 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(); diff --git a/vortex-layout/src/scan/v2/splits.rs b/vortex-layout/src/scan/v2/splits.rs index b1c5cf6c604..6e11e223d15 100644 --- a/vortex-layout/src/scan/v2/splits.rs +++ b/vortex-layout/src/scan/v2/splits.rs @@ -104,7 +104,7 @@ fn collect_nested_values( } /// The projection splits of `rows`: cut at every one of `starts` inside it. -pub(crate) fn projection_splits(starts: &[u64], rows: Range) -> Vec> { +pub(in crate::scan) fn projection_splits(starts: &[u64], rows: Range) -> Vec> { if rows.is_empty() { return Vec::new(); } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index f946dc49a51..cb3ef8f9de3 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -7,6 +7,8 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use arrow_array::cast::AsArray; +use arrow_array::types::Int32Type; use futures::TryStreamExt; use futures::stream; use rstest::rstest; @@ -27,6 +29,7 @@ use vortex_array::expr::lit; use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::validity::Validity; +use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -35,9 +38,9 @@ use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; -use super::Executor; +use super::PreparedScan; +use super::ScanBuilder; use super::ScanFile; -use super::select_supported; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; @@ -48,7 +51,7 @@ use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; use crate::scan::planning::SegmentLocation; -use crate::scan::scan_builder::ScanBuilder; +use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; use crate::scan::split_by::SplitBy; use crate::segments::SegmentFuture; use crate::segments::SegmentId; @@ -162,13 +165,10 @@ fn builder( segments: &Arc, layout: &LayoutRef, ) -> VortexResult> { - let reader = layout.new_reader( - "".into(), - Arc::clone(segments), - session, - &Default::default(), - )?; - Ok(ScanBuilder::new(session.clone(), reader).with_scan_file(scan_file(segments, layout)?)) + Ok(ScanBuilder::new( + session.clone(), + scan_file(segments, layout)?, + )) } async fn collect( @@ -184,12 +184,29 @@ async fn assert_same_as_v1( session: &VortexSession, make: impl Fn() -> VortexResult>, ) -> VortexResult<()> { - let dtype = make()?.dtype()?; - let expected = collect(dtype.clone(), make()?.prepare()?.execute_stream(None)?).await?; - let Executor::V2(scan) = select_supported(make()?)? else { - panic!("the scan should run on the planner-driven executor"); - }; - let actual = collect(dtype, scan.execute_stream()?).await?; + let builder = make()?; + let dtype = builder.dtype().clone(); + let reader = builder.file.layout.new_reader( + "".into(), + Arc::clone(&builder.file.segments), + session, + &Default::default(), + )?; + let mut v1 = V1ScanBuilder::new(session.clone(), reader) + .with_projection(builder.projection.clone()) + .with_selection(builder.selection.clone()) + .with_row_offset(builder.row_offset) + .with_ordered(builder.ordered) + .with_split_by(builder.split_by) + .with_concurrency(builder.concurrency); + if let Some(rows) = builder.row_range.clone() { + v1 = v1.with_row_range(rows); + } + if let Some(boundaries) = builder.natural_splits.clone() { + v1 = v1.with_natural_splits(boundaries); + } + let expected = collect(dtype.clone(), v1.into_stream()?).await?; + let actual = collect(dtype, builder.into_stream()?).await?; assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); Ok(()) } @@ -263,19 +280,42 @@ async fn projection_matches_v1( .await } +#[tokio::test] +async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let projection = + pack([("value", get_item("ints", root()))], false.into()).bind(layout.dtype())?; + 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 values: Vec<_> = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, (990_i32..1010).collect::>()); + Ok(()) +} + #[tokio::test] async fn zoned_projection_keeps_data_splits() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, zoned_table().as_ref()).await?; let projection = get_item("ints", root()).bind(layout.dtype())?; - let Executor::V2(scan) = select_supported( + let scan = PreparedScan::try_new( builder(&session, &segments, &layout)? .with_projection(projection) .with_row_range(250..3100), - )? - else { - panic!("the scan should run on the planner-driven executor"); - }; + )?; assert_eq!( scan.splits, vec![250..1000, 1000..2000, 2000..3000, 3000..3100] @@ -297,11 +337,7 @@ async fn split_policy_is_preserved( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let Executor::V2(scan) = - select_supported(builder(&session, &segments, &layout)?.with_split_by(policy))? - else { - panic!("the scan should run on the planner-driven executor"); - }; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?.with_split_by(policy))?; assert_eq!(scan.splits, expected); Ok(()) } @@ -310,15 +346,12 @@ async fn split_policy_is_preserved( async fn supplied_splits_take_precedence_and_are_clamped() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let Executor::V2(scan) = select_supported( + let scan = PreparedScan::try_new( builder(&session, &segments, &layout)? .with_split_by(SplitBy::RowCount(1000)) .with_natural_splits(Arc::from([0, 2500, 4000])) .with_row_range(500..3000), - )? - else { - panic!("the scan should run on the planner-driven executor"); - }; + )?; assert_eq!(scan.splits, vec![500..2500, 2500..3000]); Ok(()) } @@ -347,9 +380,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let Executor::V2(scan) = select_supported(builder(&session, &segments, &layout)?)? else { - panic!("the scan should run on the planner-driven executor"); - }; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; let splits = scan.split_plans(); assert_eq!(splits.len(), CHUNKS); for (index, split) in splits.into_iter().enumerate() { @@ -368,43 +399,19 @@ 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. #[tokio::test] -async fn filters_and_limits_run_on_v1() -> VortexResult<()> { - let session = new_session().with_tokio(); - let (segments, layout) = write(&session, table().as_ref()).await?; - let dtype = layout.dtype().clone(); - let filter = gt(get_item("ints", root()), lit(10i32)).bind(&dtype)?; - let filtered = builder(&session, &segments, &layout)?.with_filter(filter); - assert!(matches!(select_supported(filtered)?, Executor::V1(_))); - let limited = builder(&session, &segments, &layout)?.with_limit(10); - assert!(matches!(select_supported(limited)?, Executor::V1(_))); - Ok(()) -} - -#[tokio::test] -async fn scans_without_a_scan_file_run_on_v1() -> VortexResult<()> { - let session = new_session().with_tokio(); - let (segments, layout) = write(&session, table().as_ref()).await?; - let reader = layout.new_reader( - "".into(), - Arc::clone(&segments), - &session, - &Default::default(), - )?; - let builder = ScanBuilder::new(session.clone(), reader); - assert!(matches!(select_supported(builder)?, Executor::V1(_))); - Ok(()) -} - -/// List layouts are scanned by the default executor: their element ranges are only known once -/// their offsets are read, after a pipeline would have been compiled. -#[tokio::test] -async fn list_layouts_run_on_v1() -> VortexResult<()> { +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)?; - assert!(matches!(select_supported(builder)?, Executor::V1(_))); + let error = builder + .into_stream()? + .try_next() + .await + .expect_err("V2 must reject an unsupported layout"); + assert!(error.to_string().contains("V2 scans do not support")); Ok(()) } @@ -412,11 +419,7 @@ async fn list_layouts_run_on_v1() -> VortexResult<()> { async fn unordered_returns_every_row() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let Executor::V2(scan) = - select_supported(builder(&session, &segments, &layout)?.with_ordered(false))? - else { - panic!("the scan should run on the planner-driven executor"); - }; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; let chunks: Vec = scan.execute_stream()?.try_collect().await?; let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); assert_eq!(rows, CHUNKS * CHUNK_ROWS); @@ -428,14 +431,11 @@ 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 Executor::V2(scan) = select_supported( + let scan = PreparedScan::try_new( builder(&session, &segments, &layout)? .with_projection(get_item("ints", root()).bind(&dtype)?) .with_row_range(990..1010), - )? - else { - panic!("the scan should run on the planner-driven executor"); - }; + )?; let actual = collect( DType::Primitive(vortex_array::dtype::PType::I32, false.into()), scan.execute_stream()?, @@ -457,9 +457,7 @@ async fn splits_at_chunk_starts( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, strategy.get().as_ref()).await?; - let Executor::V2(scan) = select_supported(builder(&session, &segments, &layout)?)? else { - panic!("the scan should run on the planner-driven executor"); - }; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; let expected = (0..CHUNKS as u64) .map(|chunk| chunk * CHUNK_ROWS as u64..(chunk + 1) * CHUNK_ROWS as u64) .collect::>(); From e4cfe34a30b8d3e38f5bb9b76215d3123f4f8152 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 16 Oct 2026 14:53:33 +0100 Subject: [PATCH 12/14] refactor(scan): use fixed V2 chunks and an execution adaptor Signed-off-by: Joe Isaacs --- .../src/plan/physical/pipeline/values.rs | 80 +------ .../graph.rs => execution/adaptor.rs} | 187 ++++++---------- .../src/scan/{planning => execution}/mod.rs | 14 +- .../src/scan/execution/projection.rs | 163 ++++++++++++++ vortex-layout/src/scan/mod.rs | 2 +- vortex-layout/src/scan/planning/projection.rs | 207 ------------------ vortex-layout/src/scan/split_by.rs | 21 +- vortex-layout/src/scan/v2/io.rs | 4 +- vortex-layout/src/scan/v2/mod.rs | 102 +++------ vortex-layout/src/scan/v2/splits.rs | 175 --------------- vortex-layout/src/scan/v2/tests.rs | 125 +++++------ 11 files changed, 334 insertions(+), 746 deletions(-) rename vortex-layout/src/scan/{planning/graph.rs => execution/adaptor.rs} (62%) rename vortex-layout/src/scan/{planning => execution}/mod.rs (80%) create mode 100644 vortex-layout/src/scan/execution/projection.rs delete mode 100644 vortex-layout/src/scan/planning/projection.rs delete mode 100644 vortex-layout/src/scan/v2/splits.rs diff --git a/vortex-layout/src/plan/physical/pipeline/values.rs b/vortex-layout/src/plan/physical/pipeline/values.rs index 99c226da7d2..18c2116b637 100644 --- a/vortex-layout/src/plan/physical/pipeline/values.rs +++ b/vortex-layout/src/plan/physical/pipeline/values.rs @@ -12,8 +12,6 @@ use vortex_array::ArrayRef; use crate::plan::physical::PlanByPtr; use crate::plan::physical::PlanRef; -type CachedValue = Arc>>; - /// Values a plan produces over its whole domain whatever rows it is compiled for, kept for the /// later scans that share this cache. /// @@ -21,15 +19,12 @@ type CachedValue = Arc>>; /// cache can reuse them across filters, projections, and splits. The caller controls its lifetime; /// only whole values that a reader requested are retained. #[derive(Clone, Default)] -pub struct ValueCache(Arc>>); +pub struct ValueCache(Arc>>); impl ValueCache { /// The values `plan` produced earlier. pub(crate) fn get(&self, plan: &PlanRef) -> Option { - self.0 - .lock() - .get(&PlanByPtr(plan.clone())) - .and_then(|value| value.lock().clone()) + self.0.lock().get(&PlanByPtr(plan.clone())).cloned() } /// Records the values of `plan`, keeping the first when two race. @@ -37,76 +32,7 @@ impl ValueCache { self.0 .lock() .entry(PlanByPtr(plan.clone())) - .or_default() - .lock() - .get_or_insert(values) + .or_insert(values) .clone() } - - /// Shares only these entries. Once the parent cache is dropped, a value lives until the - /// last cache using its plan is dropped, rather than until every split finishes. - pub(crate) fn for_plans<'a>(&self, plans: impl Iterator) -> Self { - let mut entries = self.0.lock(); - #[allow( - clippy::mutable_key_type, - reason = "PlanByPtr hashes and compares only the stable Arc address" - )] - let scoped = plans - .map(|plan| { - let key = PlanByPtr(plan.clone()); - let value = entries.entry(key.clone()).or_default(); - (key, Arc::clone(value)) - }) - .collect(); - Self(Arc::new(Mutex::new(scoped))) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::registry::ReadContext; - - use super::ValueCache; - use crate::plan::physical::SegmentScanPlan; - use crate::segments::SegmentId; - - #[test] - fn scoped_values_live_until_the_last_consumer() -> VortexResult<()> { - let array = buffer![1_i32, 2, 3].into_array(); - let plans = [0_u32, 1].map(|id| { - SegmentScanPlan::new( - array.dtype().clone(), - array.len() as u64, - SegmentId::from(id), - ReadContext::new([]), - None, - ) - .into_plan() - }); - let parent = ValueCache::default(); - let first = parent.for_plans(std::iter::once(&plans[0])); - let last = parent.for_plans(std::iter::once(&plans[0])); - let unrelated = parent.for_plans(std::iter::once(&plans[1])); - drop(parent); - - first.insert(&plans[0], array.clone()); - unrelated.insert(&plans[1], array.clone()); - let entry = Arc::downgrade(first.0.lock().values().next().expect("one entry")); - drop(first); - assert!(ArrayRef::ptr_eq( - &last.get(&plans[0]).expect("shared value"), - &array, - )); - assert!(entry.upgrade().is_some()); - drop(last); - assert!(entry.upgrade().is_none()); - assert!(unrelated.get(&plans[1]).is_some()); - Ok(()) - } } diff --git a/vortex-layout/src/scan/planning/graph.rs b/vortex-layout/src/scan/execution/adaptor.rs similarity index 62% rename from vortex-layout/src/scan/planning/graph.rs rename to vortex-layout/src/scan/execution/adaptor.rs index e5fa1a3664c..fc6053d55ca 100644 --- a/vortex-layout/src/scan/planning/graph.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Adapts pipeline execution to the planner I/O and compute protocol. + use std::ops::Range; use std::sync::Arc; @@ -15,111 +17,25 @@ use vortex_io::request; use vortex_io::request::IoBatch; use vortex_io::request::IoIntent; use vortex_io::request::IoResult; -use vortex_mask::Mask; use vortex_scan::planning::planner::State; -use vortex_session::VortexSession; -use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ReadId; use crate::plan::physical::pipeline::Scan; -use crate::plan::physical::pipeline::Split; use crate::plan::physical::pipeline::Turn; -use crate::plan::physical::pipeline::ValueCache; -use crate::scan::planning::SegmentLocation; +use crate::scan::execution::SegmentLocation; use crate::segments::SegmentId; /// Rows a running plan produced, in row order. #[derive(Clone)] pub(super) struct Piece { - /// Orders the piece among the pieces of one graph. + /// Orders the piece among the outputs of one scan. pub(super) rows: Range, /// The selected rows the piece holds. pub(super) array: ArrayRef, } -/// A plan running over one row range on the pipeline [`Scan`]. -/// -/// The scan produces the selected rows in row order and leaves its reads to its owner. -pub(super) struct ScanGraph { - scan: Scan, - /// A step the scan already took that the owner has not seen. - pending: Option, - done: bool, - /// Orders the pieces handed out. - next_piece: u64, -} - -/// What a graph wants from its owner. -enum GraphState { - NeedsCompute, - Waiting, - Done, -} - -impl ScanGraph { - /// Compiles nothing yet: the scan compiles `plan` over `rows`, restricted to `mask`, on its - /// first step. `row_offset` is the global row index of the plan's first row, and `values` the - /// whole values shared with the split's other scans. - pub(super) fn try_new( - session: VortexSession, - plan: &PlanRef, - rows: Range, - mask: Mask, - row_offset: u64, - values: ValueCache, - ) -> VortexResult { - let scan = Scan::try_new(session, plan.clone(), vec![Split { rows, mask }])? - .with_row_offset(row_offset) - .with_value_cache(values); - Ok(Self { - scan, - pending: None, - done: false, - next_piece: 0, - }) - } - - fn state(&self) -> GraphState { - match &self.pending { - _ if self.done => GraphState::Done, - Some(Turn::Waiting) => GraphState::Waiting, - _ => GraphState::NeedsCompute, - } - } - - fn step(&mut self) -> VortexResult { - let turn = match self.pending.take() { - Some(Turn::Waiting) | None => self.scan.step()?, - Some(turn) => turn, - }; - match &turn { - Turn::Done => self.done = true, - Turn::Waiting => self.pending = Some(Turn::Waiting), - Turn::Read(_) | Turn::Output(..) => {} - } - Ok(turn) - } - - fn piece(&mut self, array: ArrayRef) -> Piece { - let start = self.next_piece; - self.next_piece += 1; - Piece { - rows: start..start + 1, - array, - } - } - - fn deliver(&mut self, read: ReadId, bytes: BufferHandle) -> VortexResult<()> { - self.scan.deliver(read, bytes)?; - if matches!(self.pending, Some(Turn::Waiting)) { - self.pending = None; - } - Ok(()) - } -} - -/// What one [`ProtocolGraph::compute`] produced. -pub(super) enum GraphStep { +/// What one [`ScanAdaptor::compute`] produced. +pub(super) enum AdaptorStep { /// Control returns to the owner with CPU work still ready. Yield, /// Requests to publish, each returned once. @@ -128,13 +44,18 @@ pub(super) enum GraphStep { Piece(Piece), } -/// A [`ScanGraph`] whose segment reads are published as protocol requests. +/// Adapts a pipeline [`Scan`] to the planner I/O and compute protocol. /// /// The scan names segments by id; the protocol names byte ranges. Each scan read becomes a /// `Fetch` of the segment's range at the segment's alignment, and each delivery is handed back to /// the scan under its own id. -pub(super) struct ProtocolGraph { - graph: ScanGraph, +pub(super) struct ScanAdaptor { + scan: Scan, + /// A turn already taken while batching reads, retained for the next compute. + pending: Option, + done: bool, + /// Orders the pieces handed out. + next_piece: u64, locations: Arc<[SegmentLocation]>, /// Published and undelivered requests, with the scan read and segment each one answers. outstanding: FxHashMap, @@ -144,65 +65,87 @@ pub(super) struct ProtocolGraph { next_id: u32, } -impl ProtocolGraph { - /// Wraps `graph`, numbering its requests from `first_id`. Request ids must be unique over an - /// owner's lifetime. - pub(super) fn new(graph: ScanGraph, locations: Arc<[SegmentLocation]>, first_id: u32) -> Self { +impl ScanAdaptor { + /// Wraps a scan, numbering requests uniquely within this adaptor's lifetime. + pub(super) fn new(scan: Scan, locations: Arc<[SegmentLocation]>) -> Self { Self { - graph, + scan, + pending: None, + done: false, + next_piece: 0, locations, outstanding: FxHashMap::default(), failed: None, - next_id: first_id, + next_id: 0, } } - /// The owner's protocol state while the graph is running. + /// The owner's protocol state while the scan is running. pub(super) fn state(&self) -> State { - if self.failed.is_some() { - return State::NeedsCompute; - } - match self.graph.state() { - GraphState::Done => State::Done, - GraphState::NeedsCompute => State::NeedsCompute, - GraphState::Waiting => State::Waiting, + match &self.pending { + _ if self.failed.is_some() => State::NeedsCompute, + _ if self.done => State::Done, + Some(Turn::Waiting) => State::Waiting, + _ => State::NeedsCompute, } } - pub(super) fn compute(&mut self) -> VortexResult { + pub(super) fn compute(&mut self) -> VortexResult { if let Some(error) = self.failed.take() { return Err(error); } - Ok(match self.graph.step()? { - Turn::Waiting | Turn::Done => GraphStep::Yield, - Turn::Output(_, array) => GraphStep::Piece(self.graph.piece(array)), + Ok(match self.step()? { + Turn::Waiting | Turn::Done => AdaptorStep::Yield, + Turn::Output(_, array) => AdaptorStep::Piece(self.piece(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)?]; loop { - match self.graph.step()? { + match self.step()? { Turn::Read(read) => batch.push(self.request(read.id, read.segment_id)?), turn => { if !matches!(turn, Turn::Waiting) { - self.graph.pending = Some(turn); - self.graph.done = false; + self.pending = Some(turn); + self.done = false; } break; } } } - GraphStep::NeedsIO(batch) + AdaptorStep::NeedsIO(batch) } }) } + fn step(&mut self) -> VortexResult { + let turn = match self.pending.take() { + Some(Turn::Waiting) | None => self.scan.step()?, + Some(turn) => turn, + }; + match &turn { + Turn::Done => self.done = true, + Turn::Waiting => self.pending = Some(Turn::Waiting), + Turn::Read(_) | Turn::Output(..) => {} + } + Ok(turn) + } + + fn piece(&mut self, array: ArrayRef) -> Piece { + let start = self.next_piece; + self.next_piece += 1; + Piece { + rows: start..start + 1, + array, + } + } + fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { let location = self.location(segment_id)?; let id = self.next_id; self.next_id = self .next_id .checked_add(1) - .ok_or_else(|| vortex_err!("ProtocolGraph ran out of request ids"))?; + .ok_or_else(|| vortex_err!("ScanAdaptor ran out of request ids"))?; self.outstanding.insert(id, (read, segment_id)); Ok(request::IoRequest { intent: IoIntent::Fetch, @@ -211,15 +154,15 @@ impl ProtocolGraph { }) } - /// Hands a delivered read to the scan. The driver only delivers requests this graph + /// Hands a delivered read to the scan. The driver only delivers requests this adaptor /// published, so anything else is a driver bug. Bytes the scan cannot take fail the owner's /// next compute. pub(super) fn set_io_result(&mut self, id: request::IoRequestId, result: IoResult) { let Some((read, segment_id)) = self.outstanding.remove(&id.0) else { - vortex_panic!("ProtocolGraph: delivery of {id:?}, which is not outstanding"); + vortex_panic!("ScanAdaptor: delivery of {id:?}, which is not outstanding"); }; let IoResult::Bytes(bytes) = result else { - vortex_panic!("ProtocolGraph: segment {segment_id} answered with a size"); + vortex_panic!("ScanAdaptor: segment {segment_id} answered with a size"); }; if let Err(error) = self.deliver(read, segment_id, bytes) && self.failed.is_none() @@ -237,7 +180,11 @@ impl ProtocolGraph { // A source that honours the target's alignment makes this a check, not a copy. let alignment = self.location(segment_id)?.alignment; let bytes = BufferHandle::new_host(bytes.try_into_host_sync()?.aligned(alignment)); - self.graph.deliver(read, bytes) + self.scan.deliver(read, bytes)?; + if matches!(self.pending, Some(Turn::Waiting)) { + self.pending = None; + } + Ok(()) } fn location(&self, id: SegmentId) -> VortexResult { diff --git a/vortex-layout/src/scan/planning/mod.rs b/vortex-layout/src/scan/execution/mod.rs similarity index 80% rename from vortex-layout/src/scan/planning/mod.rs rename to vortex-layout/src/scan/execution/mod.rs index e95180ab1da..9d8b8897ef3 100644 --- a/vortex-layout/src/scan/planning/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -1,14 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Layout-level planners and morsels for the planning protocol. +//! Scan execution over physical plans through the planner I/O and compute protocol. //! -//! A split is evaluated over physical plans: a `ProjectionPlanner` emits a `ProjectionMorsel` -//! per projection split, which runs the projection plan over the split's selected rows through +//! A `ProjectionPlanner` emits one `ProjectionMorsel` for each fixed-size work chunk. The morsel +//! runs the physical projection plan over the chunk's selected rows through //! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published -//! as a protocol request, so the planning driver decides when IO happens. +//! through a `ScanAdaptor`, so the planning driver can advance other work while reads wait. -mod graph; +mod adaptor; mod projection; use std::sync::Arc; @@ -55,10 +55,6 @@ pub(super) struct ScanPlans { pub(super) locations: Arc<[SegmentLocation]>, /// The optimized projection plan, over the scan's root row domain. pub(super) projection: PlanRef, - /// Where the chunks the projection reads start, over the root row domain, sorted. A split's - /// projection is cut at each of these, so every projection split reads at most one chunk of - /// every column. - pub(super) projection_starts: Arc<[u64]>, /// The global row index of the root row domain's first row. pub(super) row_offset: u64, /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs new file mode 100644 index 00000000000..b6180725c27 --- /dev/null +++ b/vortex-layout/src/scan/execution/projection.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::dtype::DType; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_scan::planning::morsel::Morsel; +use vortex_scan::planning::morsel::MorselOutput; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; + +use 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; +use crate::scan::execution::adaptor::Piece; +use crate::scan::execution::adaptor::ScanAdaptor; + +/// Hands one fixed-size work chunk to a projection morsel. +pub(super) struct ProjectionPlanner { + work: Option<(ScanPlans, SelectedRows)>, +} + +impl ProjectionPlanner { + /// Creates a planner for `selected`. + pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + Self { + work: Some((plans, selected)), + } + } +} + +impl IoConsumer for ProjectionPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for ProjectionPlanner { + fn state(&self) -> State { + if self.work.is_some() { + State::NeedsCompute + } else { + State::Done + } + } + + fn compute(&mut self) -> VortexResult { + let Some((plans, selected)) = self.work.take() else { + return Ok(PlannerOutput::Done); + }; + if selected.mask.all_false() { + return Ok(PlannerOutput::Done); + } + Ok(PlannerOutput::Morsel( + selected.scope.clone(), + Box::new(ProjectionMorsel::new(plans, selected)), + )) + } +} + +/// Runs the projection plan over a split's selected rows and returns them as one batch, in row +/// order. A selection that projects to no rows finishes without a batch. +struct ProjectionMorsel { + plans: Option, + dtype: DType, + selected: SelectedRows, + adaptor: Option, + pieces: Vec, + done: bool, +} + +impl ProjectionMorsel { + /// Creates a morsel projecting `selected`. + fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + Self { + dtype: plans.projection.dtype().clone(), + plans: Some(plans), + selected, + adaptor: None, + pieces: Vec::new(), + done: false, + } + } + + fn finish(&mut self) -> VortexResult { + self.done = true; + self.pieces.sort_by_key(|piece| piece.rows.start); + let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) + .into_iter() + .map(|piece| piece.array) + .filter(|array| !array.is_empty()) + .collect(); + let array = match arrays.len() { + 0 => return Ok(MorselOutput::Done), + 1 => arrays.remove(0), + _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), + }; + Ok(MorselOutput::Batch(array)) + } +} + +impl IoConsumer for ProjectionMorsel { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some(adaptor) = self.adaptor.as_mut() { + adaptor.set_io_result(request, result); + } + } +} + +impl Morsel for ProjectionMorsel { + fn state(&self) -> State { + match &self.adaptor { + _ if self.done => State::Done, + None => State::NeedsCompute, + Some(adaptor) => match adaptor.state() { + State::Done => State::NeedsCompute, + state => state, + }, + } + } + + fn compute(&mut self) -> VortexResult { + if self.done { + vortex_bail!("ProjectionMorsel: compute called after Done"); + } + let Some(adaptor) = self.adaptor.as_mut() else { + let plans = self + .plans + .take() + .vortex_expect("an uncompiled projection still has its plans"); + let scan = Scan::try_new( + plans.session, + 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); + self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); + return Ok(MorselOutput::Continue); + }; + if adaptor.state() == State::Done { + return self.finish(); + } + Ok(match adaptor.compute()? { + AdaptorStep::Yield => MorselOutput::Continue, + AdaptorStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), + AdaptorStep::Piece(piece) => { + self.pieces.push(piece); + MorselOutput::Continue + } + }) + } +} diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index fcf15137e06..2d3d3b70031 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -2,11 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub mod arrow; +mod execution; mod filter; pub mod layout; mod limit; pub mod multi; -mod planning; pub mod repeated_scan; pub mod scan_builder; pub mod split_by; diff --git a/vortex-layout/src/scan/planning/projection.rs b/vortex-layout/src/scan/planning/projection.rs deleted file mode 100644 index 70cd2e59ca0..00000000000 --- a/vortex-layout/src/scan/planning/projection.rs +++ /dev/null @@ -1,207 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_array::IntoArray; -use vortex_array::arrays::ChunkedArray; -use vortex_array::dtype::DType; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_io::request::IoConsumer; -use vortex_io::request::IoRequestId; -use vortex_io::request::IoResult; -use vortex_scan::planning::morsel::Morsel; -use vortex_scan::planning::morsel::MorselOutput; -use vortex_scan::planning::planner::Planner; -use vortex_scan::planning::planner::PlannerOutput; -use vortex_scan::planning::planner::State; -use vortex_scan::planning::planner::WorkScope; - -use crate::scan::planning::ScanPlans; -use crate::scan::planning::SelectedRows; -use crate::scan::planning::graph::GraphStep; -use crate::scan::planning::graph::Piece; -use crate::scan::planning::graph::ProtocolGraph; -use crate::scan::planning::graph::ScanGraph; -use crate::scan::v2::splits::projection_splits; - -/// Turns a split's selected rows into the morsels that project them. -/// -/// The rows are cut into projection splits where the projection's chunks start, and each split -/// with a selected row becomes one `ProjectionMorsel`, so a morsel reads at most one chunk of -/// every column. The planner finishes once it has handed out every morsel. -pub(super) struct ProjectionPlanner { - plans: Option, - /// The split's selected rows, until the first compute cuts them. - selected: Option, - /// The projection splits not yet handed out, last first. - pending: Vec, -} - -impl ProjectionPlanner { - /// Creates a planner for `selected`. - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { - Self { - plans: Some(plans), - selected: Some(selected), - pending: Vec::new(), - } - } - - /// The projection splits of `selected` that have a selected row, last first. - fn cut(&self, selected: SelectedRows) -> Vec { - let SelectedRows { scope, mask } = selected; - let start = scope.rows.start; - let index = |row: u64| usize::try_from(row - start).vortex_expect("split row fits usize"); - let plans = self.plans.as_ref().vortex_expect("uncut planner has plans"); - let mut splits = projection_splits(&plans.projection_starts, scope.rows.clone()) - .into_iter() - .filter_map(|rows| { - let mask = mask.slice(index(rows.start)..index(rows.end)); - (!mask.all_false()).then(|| SelectedRows { - scope: WorkScope { - file_ordinal: scope.file_ordinal, - rows, - }, - mask, - }) - }) - .collect::>(); - splits.reverse(); - splits - } -} - -impl IoConsumer for ProjectionPlanner { - fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} -} - -impl Planner for ProjectionPlanner { - fn state(&self) -> State { - if self.selected.is_none() && self.pending.is_empty() { - State::Done - } else { - State::NeedsCompute - } - } - - fn compute(&mut self) -> VortexResult { - if let Some(selected) = self.selected.take() { - self.pending = self.cut(selected); - } - // No projection split has a selected row. - let Some(selected) = self.pending.pop() else { - return Ok(PlannerOutput::Done); - }; - let scope = selected.scope.clone(); - // The driver's continuation follows its children. Give the last morsel ownership so - // the completed planner does not keep the lowering tree alive while that morsel runs. - let plans = if self.pending.is_empty() { - self.plans.take().vortex_expect("planner has plans") - } else { - self.plans - .as_ref() - .vortex_expect("planner has plans") - .clone() - }; - Ok(PlannerOutput::Morsel( - scope, - Box::new(ProjectionMorsel::new(plans, selected)), - )) - } -} - -/// Runs the projection plan over a split's selected rows and returns them as one batch, in row -/// order. A selection that projects to no rows finishes without a batch. -struct ProjectionMorsel { - plans: Option, - dtype: DType, - selected: SelectedRows, - graph: Option, - pieces: Vec, - done: bool, -} - -impl ProjectionMorsel { - /// Creates a morsel projecting `selected`. - fn new(plans: ScanPlans, selected: SelectedRows) -> Self { - Self { - dtype: plans.projection.dtype().clone(), - plans: Some(plans), - selected, - graph: None, - pieces: Vec::new(), - done: false, - } - } - - fn finish(&mut self) -> VortexResult { - self.done = true; - self.pieces.sort_by_key(|piece| piece.rows.start); - let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) - .into_iter() - .map(|piece| piece.array) - .filter(|array| !array.is_empty()) - .collect(); - let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), - 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), - }; - Ok(MorselOutput::Batch(array)) - } -} - -impl IoConsumer for ProjectionMorsel { - fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { - if let Some(graph) = self.graph.as_mut() { - graph.set_io_result(request, result); - } - } -} - -impl Morsel for ProjectionMorsel { - fn state(&self) -> State { - match &self.graph { - _ if self.done => State::Done, - None => State::NeedsCompute, - Some(graph) => match graph.state() { - State::Done => State::NeedsCompute, - state => state, - }, - } - } - - fn compute(&mut self) -> VortexResult { - if self.done { - vortex_bail!("ProjectionMorsel: compute called after Done"); - } - let Some(graph) = self.graph.as_mut() else { - let plans = self - .plans - .take() - .vortex_expect("an uncompiled projection still has its plans"); - let graph = ScanGraph::try_new( - plans.session, - &plans.projection, - self.selected.scope.rows.clone(), - self.selected.mask.clone(), - plans.row_offset, - plans.values, - )?; - self.graph = Some(ProtocolGraph::new(graph, plans.locations, 0)); - return Ok(MorselOutput::Continue); - }; - if graph.state() == State::Done { - return self.finish(); - } - Ok(match graph.compute()? { - GraphStep::Yield => MorselOutput::Continue, - GraphStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), - GraphStep::Piece(piece) => { - self.pieces.push(piece); - MorselOutput::Continue - } - }) - } -} diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 761959132a6..b995e46222f 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -56,29 +56,18 @@ impl SplitBy { layout_reader: &dyn LayoutReader, row_range: &Range, field_mask: &[FieldMask], - ) -> VortexResult> { - let boundaries = if matches!(self, Self::RowCount(_)) { - Vec::new() - } else { - layout_boundaries(layout_reader, row_range, field_mask)? - }; - self.split_boundaries(boundaries, row_range) - } - - /// Applies this split policy to natural boundaries from a layout or physical plan. - pub(super) fn split_boundaries( - &self, - boundaries: Vec, - row_range: &Range, ) -> VortexResult> { Ok(match *self { - SplitBy::Layout => boundaries, + SplitBy::Layout => layout_boundaries(layout_reader, row_range, field_mask)?, SplitBy::LayoutSubSplitting { max_rows } => { vortex_ensure!( max_rows > 0, "SplitBy::LayoutSubSplitting requires a non-zero max_rows" ); - subdivide_large_spans(boundaries, max_rows) + subdivide_large_spans( + layout_boundaries(layout_reader, row_range, field_mask)?, + max_rows, + ) } SplitBy::RowCount(n) => row_range .clone() diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 888e53506a7..ce565846d7e 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -29,7 +29,7 @@ use vortex_io::request::IoSource; use vortex_io::request::IoTarget; use vortex_utils::aliases::hash_map::HashMap; -use crate::scan::planning::SegmentLocation; +use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; @@ -248,7 +248,7 @@ mod tests { use super::SegmentIoSource; use super::segment_ranges; - use crate::scan::planning::SegmentLocation; + use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 4b2ee26175e..d500fb29403 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -10,15 +10,14 @@ //! //! 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. +//! error; filters and limits are not part of this builder's API. Work is partitioned into fixed +//! 65,536-row chunks, independently of the stored layout's chunk boundaries. mod io; mod split; -pub(super) mod splits; #[cfg(test)] mod tests; -use std::iter; use std::ops::Range; use std::sync::Arc; @@ -52,22 +51,35 @@ use crate::layouts::flat::Flat; use crate::layouts::struct_::Struct; use crate::layouts::zoned::LegacyStats; use crate::layouts::zoned::Zoned; -use crate::plan::physical::PlanRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; -use crate::scan::planning::ScanPlans; -pub use crate::scan::planning::SegmentLocation; -use crate::scan::planning::plan_projection; -use crate::scan::split_by::SplitBy; +use crate::scan::execution::ScanPlans; +pub use crate::scan::execution::SegmentLocation; +use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; -use crate::scan::v2::splits::chunk_starts; -use crate::scan::v2::splits::projection_splits; use crate::segments::SegmentSource; +/// Eight of the default writer's 8,192-row blocks per scan task. +const ROWS_PER_CHUNK: usize = 8 * 8192; + +/// Fixed boundaries in the file's row domain, clipped to the requested range. +fn row_chunks(rows: Range) -> Vec> { + if rows.is_empty() { + return Vec::new(); + } + let first = rows.start / ROWS_PER_CHUNK as u64 * ROWS_PER_CHUNK as u64; + (first..rows.end) + .step_by(ROWS_PER_CHUNK) + .map(|start| { + start.max(rows.start)..start.saturating_add(ROWS_PER_CHUNK as u64).min(rows.end) + }) + .collect() +} + /// The parts of a file the planner-driven executor reads through. #[derive(Clone)] pub struct ScanFile { @@ -90,8 +102,6 @@ pub struct ScanBuilder { ordered: bool, row_range: Option>, selection: Selection, - split_by: SplitBy, - natural_splits: Option>, concurrency: usize, map_fn: Arc VortexResult + Send + Sync>, row_offset: u64, @@ -107,8 +117,6 @@ impl ScanBuilder { ordered: true, row_range: None, selection: Selection::default(), - split_by: SplitBy::default(), - natural_splits: None, concurrency: 4, map_fn: Arc::new(Ok), row_offset: 0, @@ -173,23 +181,6 @@ impl ScanBuilder { self } - /// Configures how natural scan work is split for concurrency. - pub fn with_split_by(mut self, split_by: SplitBy) -> Self { - self.split_by = split_by; - self - } - - /// Supplies strictly increasing full-file boundaries, overriding the split policy. - /// Boundaries are clamped to the scan's row range. - pub fn with_natural_splits(mut self, boundaries: Arc<[u64]>) -> Self { - debug_assert!( - boundaries.windows(2).all(|w| w[0] < w[1]), - "natural split boundaries must be strictly increasing" - ); - self.natural_splits = Some(boundaries); - self - } - /// Sets the number of active splits per worker thread. Must be greater than zero. pub fn with_concurrency(mut self, concurrency: usize) -> Self { assert!(concurrency > 0); @@ -210,8 +201,6 @@ impl ScanBuilder { ordered: self.ordered, row_range: self.row_range, selection: self.selection, - split_by: self.split_by, - natural_splits: self.natural_splits, concurrency: self.concurrency, row_offset: self.row_offset, map_fn: Arc::new(move |array| old_map_fn(array).and_then(&map_fn)), @@ -264,8 +253,6 @@ fn supports_layout(layout: &LayoutRef, identity: bool) -> VortexResult { struct PreparedScan { session: VortexSession, plans: ScanPlans, - /// Whole dictionary values and the root rows that use them, for an identity projection. - value_scopes: Option, PlanRef)>>, /// Serves the splits' reads. io: Arc, /// The scan's splits, over the root row domain, in row order. @@ -277,8 +264,8 @@ struct PreparedScan { } impl PreparedScan { - /// Lowers the file's layout, plans the projection over it, and cuts the scan's row range - /// where the projection's chunks start. + /// 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!( @@ -295,29 +282,7 @@ impl PreparedScan { } else { optimize(plan_row_idx_expression(builder.projection, root)?)? }; - let mut value_scopes = identity.then(Vec::new); - let projection_starts: Arc<[u64]> = - chunk_starts([&projection], value_scopes.as_mut())?.into(); - let rows = builder.row_range.unwrap_or(0..row_count); - // Whole-column dictionaries already need the scan's full lifetime. Only dictionaries - // covering part of the scan need separate cache owners for their consuming splits. - let value_scopes = value_scopes.filter(|values| { - values - .iter() - .any(|(used, _)| used.start > rows.start || used.end < rows.end) - }); - let splits = if let Some(boundaries) = builder.natural_splits { - projection_splits(&boundaries, rows) - } else { - let natural = projection_splits(&projection_starts, rows.clone()); - let boundaries = natural - .iter() - .map(|split| split.start) - .chain(iter::once(rows.end)) - .collect(); - let boundaries = builder.split_by.split_boundaries(boundaries, &rows)?; - projection_splits(&boundaries, rows) - }; + let splits = row_chunks(builder.row_range.unwrap_or(0..row_count)); let io = Arc::new(SegmentScanIo::new( file.segments, segment_ranges(&file.locations), @@ -328,11 +293,9 @@ impl PreparedScan { session: builder.session, locations: file.locations, projection, - projection_starts, row_offset: builder.row_offset, values: Default::default(), }, - value_scopes, io, splits, selection: builder.selection, @@ -355,17 +318,12 @@ impl PreparedScan { file_ordinal: 0, rows: row_mask.row_range(), }; - let mut plans = self.plans.clone(); - if let Some(values) = &self.value_scopes { - plans.values = self.plans.values.for_plans( - values - .iter() - .filter(|(used, _)| used.start < rows.end && rows.start < used.end) - .map(|(_, plan)| plan), - ); - } Some(SplitPlan { - root: plan_projection(plans, scope.clone(), row_mask.mask().clone()), + root: plan_projection( + self.plans.clone(), + scope.clone(), + row_mask.mask().clone(), + ), scope, }) }) diff --git a/vortex-layout/src/scan/v2/splits.rs b/vortex-layout/src/scan/v2/splits.rs deleted file mode 100644 index 6e11e223d15..00000000000 --- a/vortex-layout/src/scan/v2/splits.rs +++ /dev/null @@ -1,175 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Splits a scan from its plans rather than from the layout reader. -//! -//! A scan is cut where the chunks of the columns its projection reads start, so each split reads -//! at most one chunk of every column. - -use std::iter; -use std::ops::Range; - -use itertools::Itertools; -use vortex_error::VortexResult; - -use crate::plan::physical::Concat; -use crate::plan::physical::Eval; -use crate::plan::physical::Filter; -use crate::plan::physical::LegacyStats; -use crate::plan::physical::Pack; -use crate::plan::physical::PlanRef; -use crate::plan::physical::Take; -use crate::plan::physical::Zoned; - -/// The row positions where the chunks of `plans` start, over their shared row domain, sorted and -/// without duplicates. -/// -/// This follows plans whose children cover the same rows, or known parts of them, and stops at any -/// other plan: a take's values cover other rows, and zone statistics cover zones. -/// For a lowered identity projection, `values` also records the root rows that need each whole -/// dictionary, so its cached value can be released after those splits finish. -pub(super) fn chunk_starts<'a>( - plans: impl IntoIterator, - mut values: Option<&mut Vec<(Range, PlanRef)>>, -) -> VortexResult> { - let mut starts = Vec::new(); - for plan in plans { - collect_starts(plan, 0, &mut starts, values.as_deref_mut())?; - } - starts.sort_unstable(); - starts.dedup(); - Ok(starts) -} - -fn collect_starts( - plan: &PlanRef, - offset: u64, - starts: &mut Vec, - mut values: Option<&mut Vec<(Range, PlanRef)>>, -) -> VortexResult<()> { - if let Some(concat) = plan.as_opt::() { - for (index, &child_offset) in concat.row_offsets().iter().enumerate() { - starts.push(offset + child_offset); - collect_starts( - &plan.child_required(index)?, - offset + child_offset, - starts, - values.as_deref_mut(), - )?; - } - } else if let Some(zoned) = plan.as_opt::() { - if let Some(data) = zoned.data_plan()? { - collect_starts(&data, offset, starts, values)?; - } - } else if let Some(legacy) = plan.as_opt::() { - collect_starts(&legacy.data_plan()?, offset, starts, values)?; - } else if plan.is::() { - if let Some(values) = values.as_deref_mut() { - let whole = plan.child_required(1)?; - let rows = offset..offset + plan.row_count(); - values.push((rows.clone(), whole.clone())); - collect_nested_values(&whole, &rows, values)?; - } - collect_starts(&plan.child_required(0)?, offset, starts, values)?; - } else if plan.is::() || plan.is::() || plan.is::() { - for child in plan.children().iter_refs() { - collect_starts(child?, offset, starts, values.as_deref_mut())?; - } - } - Ok(()) -} - -/// A dictionary's whole values can contain further dictionaries. They are needed over the -/// outer dictionary's row domain, regardless of their own row counts. -fn collect_nested_values( - plan: &PlanRef, - rows: &Range, - values: &mut Vec<(Range, PlanRef)>, -) -> VortexResult<()> { - if let Some(zoned) = plan.as_opt::() { - if let Some(data) = zoned.data_plan()? { - collect_nested_values(&data, rows, values)?; - } - } else if let Some(legacy) = plan.as_opt::() { - collect_nested_values(&legacy.data_plan()?, rows, values)?; - } else { - if plan.is::() { - values.push((rows.clone(), plan.child_required(1)?)); - } - for child in plan.children().iter_refs() { - collect_nested_values(child?, rows, values)?; - } - } - Ok(()) -} - -/// The projection splits of `rows`: cut at every one of `starts` inside it. -pub(in crate::scan) fn projection_splits(starts: &[u64], rows: Range) -> Vec> { - if rows.is_empty() { - return Vec::new(); - } - let first = starts.partition_point(|&start| start <= rows.start); - let last = starts.partition_point(|&start| start < rows.end); - iter::once(rows.start) - .chain(starts[first..last].iter().copied()) - .chain(iter::once(rows.end)) - .tuple_windows() - .map(|(start, end)| start..end) - .filter(|range| !range.is_empty()) - .collect() -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - use vortex_array::dtype::DType; - use vortex_array::dtype::PType; - use vortex_session::registry::ReadContext; - - use super::*; - use crate::plan::physical::ConcatPlan; - use crate::plan::physical::SegmentScanPlan; - use crate::plan::physical::TakePlan; - use crate::segments::SegmentId; - - #[rstest] - #[case::whole(&[], 10..20, vec![10..20])] - #[case::cut(&[0, 12, 15, 30], 10..20, vec![10..12, 12..15, 15..20])] - #[case::start_on_boundary(&[10, 15], 10..20, vec![10..15, 15..20])] - #[case::empty(&[10], 5..5, vec![])] - #[case::empty_on_boundary(&[5], 5..5, vec![])] - fn projection( - #[case] starts: &[u64], - #[case] rows: Range, - #[case] expected: Vec>, - ) { - assert_eq!(projection_splits(starts, rows), expected); - } - - #[test] - fn dictionary_values_use_the_outer_codes_row_domain() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, false.into()); - let segment = |id, rows| { - SegmentScanPlan::new( - dtype.clone(), - rows, - SegmentId::from(id), - ReadContext::new([]), - None, - ) - .into_plan() - }; - let inner_values = segment(0_u32, 3); - let outer_values = TakePlan::new(segment(1, 2), inner_values.clone()).into_plan(); - let dict = TakePlan::new(segment(2, 4), outer_values.clone()).into_plan(); - let plan = ConcatPlan::try_new(dtype.clone(), vec![segment(3, 5), dict])?.into_plan(); - let mut values = Vec::new(); - assert_eq!(chunk_starts([&plan], Some(&mut values))?, vec![0, 5]); - assert_eq!(values.len(), 2); - for ((rows, actual), expected) in values.iter().zip([outer_values, inner_values]) { - assert_eq!(rows, &(5..9)); - assert!(PlanRef::ptr_eq(actual, &expected)); - } - Ok(()) - } -} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index cb3ef8f9de3..e5763e21c85 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -33,14 +33,15 @@ use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; -use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use super::PreparedScan; +use super::ROWS_PER_CHUNK; use super::ScanBuilder; use super::ScanFile; +use super::row_chunks; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; @@ -50,9 +51,8 @@ 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::planning::SegmentLocation; +use crate::scan::execution::SegmentLocation; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; -use crate::scan::split_by::SplitBy; use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; @@ -197,14 +197,10 @@ async fn assert_same_as_v1( .with_selection(builder.selection.clone()) .with_row_offset(builder.row_offset) .with_ordered(builder.ordered) - .with_split_by(builder.split_by) .with_concurrency(builder.concurrency); if let Some(rows) = builder.row_range.clone() { v1 = v1.with_row_range(rows); } - if let Some(boundaries) = builder.natural_splits.clone() { - v1 = v1.with_natural_splits(boundaries); - } let expected = collect(dtype.clone(), v1.into_stream()?).await?; let actual = collect(dtype, builder.into_stream()?).await?; assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); @@ -306,54 +302,14 @@ async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { Ok(()) } -#[tokio::test] -async fn zoned_projection_keeps_data_splits() -> VortexResult<()> { - let session = new_session().with_tokio(); - let (segments, layout) = write(&session, zoned_table().as_ref()).await?; - let projection = get_item("ints", root()).bind(layout.dtype())?; - let scan = PreparedScan::try_new( - builder(&session, &segments, &layout)? - .with_projection(projection) - .with_row_range(250..3100), - )?; - assert_eq!( - scan.splits, - vec![250..1000, 1000..2000, 2000..3000, 3000..3100] - ); - Ok(()) -} - #[rstest] -#[case::layout(SplitBy::Layout, vec![0..1000, 1000..2000, 2000..3000, 3000..4000])] -#[case::bounded( - SplitBy::LayoutSubSplitting { max_rows: 600 }, - vec![0..500, 500..1000, 1000..1500, 1500..2000, 2000..2500, 2500..3000, 3000..3500, 3500..4000] -)] -#[case::rows(SplitBy::RowCount(1500), vec![0..1500, 1500..3000, 3000..4000])] -#[tokio::test] -async fn split_policy_is_preserved( - #[case] policy: SplitBy, - #[case] expected: Vec>, -) -> 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_split_by(policy))?; - assert_eq!(scan.splits, expected); - Ok(()) -} - -#[tokio::test] -async fn supplied_splits_take_precedence_and_are_clamped() -> 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_split_by(SplitBy::RowCount(1000)) - .with_natural_splits(Arc::from([0, 2500, 4000])) - .with_row_range(500..3000), - )?; - assert_eq!(scan.splits, vec![500..2500, 2500..3000]); - Ok(()) +#[case::whole(0..150_000, vec![0..65_536, 65_536..131_072, 131_072..150_000])] +#[case::unaligned(10..140_000, vec![10..65_536, 65_536..131_072, 131_072..140_000])] +#[case::boundary(65_536..131_072, vec![65_536..131_072])] +#[case::empty(5..5, vec![])] +#[case::near_max(u64::MAX - 3..u64::MAX, vec![u64::MAX - 3..u64::MAX])] +fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) { + assert_eq!(row_chunks(rows), expected); } struct CountingSegments { @@ -371,9 +327,31 @@ impl SegmentSource for CountingSegments { #[tokio::test] async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let session = new_session().with_tokio(); - let (segments, layout) = write(&session, table().as_ref()).await?; - let ints = layout.slot(1)?.ok_or_else(|| vortex_err!("missing ints"))?; - // The codes 0..4000 index the same values, so each split returns its own row numbers. + const STORED_CHUNK_ROWS: usize = 70_000; + const STORED_CHUNKS: usize = 3; + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let arrays = (0..STORED_CHUNKS) + .map(|index| { + let start = i32::try_from(index * STORED_CHUNK_ROWS)?; + let end = start + i32::try_from(STORED_CHUNK_ROWS)?; + Ok(( + sequence.advance(), + PrimitiveArray::from_iter(start..end).into_array(), + )) + }) + .collect::>>()?; + let dtype = DType::Primitive(vortex_array::dtype::PType::I32, false.into()); + let ints = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))).sendable(), + eof, + &session, + ) + .await?; + // Codes index the same values across fixed work chunks and stored chunk boundaries. let layout = DictLayout::new(Arc::clone(&ints), ints).into_layout(); let counted = Arc::new(CountingSegments { inner: segments, @@ -382,14 +360,27 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let segments = Arc::clone(&counted) as Arc; let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; let splits = scan.split_plans(); - assert_eq!(splits.len(), CHUNKS); + assert_eq!( + splits.len(), + (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(ROWS_PER_CHUNK) + ); for (index, split) in splits.into_iter().enumerate() { + let rows = split.scope.rows.clone(); let before = counted.reads.load(Ordering::Relaxed); let arrays = split.run(Arc::clone(&scan.io)).await?; let reads = counted.reads.load(Ordering::Relaxed) - before; - assert_eq!(reads, if index == 0 { CHUNKS } else { 1 }); - let start = i32::try_from(index * CHUNK_ROWS)?; - let expected = PrimitiveArray::from_iter(start..start + i32::try_from(CHUNK_ROWS)?); + let code_chunks = (rows.end - 1) / STORED_CHUNK_ROWS as u64 + - rows.start / STORED_CHUNK_ROWS as u64 + + 1; + assert_eq!( + reads, + if index == 0 { + STORED_CHUNKS + } else { + code_chunks as usize + } + ); + let expected = PrimitiveArray::from_iter(i32::try_from(rows.start)?..i32::try_from(rows.end)?); assert_arrays_eq!( ChunkedArray::try_new(arrays, layout.dtype().clone())?.into_array(), expected, @@ -449,18 +440,18 @@ async fn primitive_column() -> VortexResult<()> { Ok(()) } -/// Every chunk of the default zoned table is its own split, so the scan runs them in parallel. +/// Stored chunk boundaries do not subdivide a fixed work chunk. #[rstest] #[tokio::test] -async fn splits_at_chunk_starts( +async fn work_chunks_ignore_layout_boundaries( #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, ) -> 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)?)?; - let expected = (0..CHUNKS as u64) - .map(|chunk| chunk * CHUNK_ROWS as u64..(chunk + 1) * CHUNK_ROWS as u64) - .collect::>(); - assert_eq!(scan.splits, expected); + assert_eq!(scan.splits, vec![0..(CHUNKS * CHUNK_ROWS) as u64]); + let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; + assert_eq!(arrays.len(), 1); + assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) } From e39866a38d14f81a763f801039b20a71ec4888d4 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 16 Oct 2026 17:03:19 +0100 Subject: [PATCH 13/14] perf(scan): bound V2 row chunks by serialized size Signed-off-by: Joe Isaacs --- vortex-layout/src/scan/v2/mod.rs | 37 +++++++++++++++++++++++------- vortex-layout/src/scan/v2/tests.rs | 31 ++++++++++++++++++++++--- 2 files changed, 57 insertions(+), 11 deletions(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index d500fb29403..46cbde987fc 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -11,7 +11,8 @@ //! 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 -//! 65,536-row chunks, independently of the stored layout's chunk boundaries. +//! 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. mod io; mod split; @@ -63,19 +64,34 @@ use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; use crate::segments::SegmentSource; -/// Eight of the default writer's 8,192-row blocks per scan task. -const ROWS_PER_CHUNK: usize = 8 * 8192; +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; + +/// Limit estimated bytes per task without changing the writer-aligned row granularity. +fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { + if row_count <= MAX_ROWS_PER_CHUNK as u64 { + return MAX_ROWS_PER_CHUNK; + } + let bytes: u64 = locations + .iter() + .map(|location| u64::from(location.length)) + .sum(); + let bytes_per_row = bytes.div_ceil(row_count).max(1); + let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; + blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE +} /// Fixed boundaries in the file's row domain, clipped to the requested range. -fn row_chunks(rows: Range) -> Vec> { +fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { if rows.is_empty() { return Vec::new(); } - let first = rows.start / ROWS_PER_CHUNK as u64 * ROWS_PER_CHUNK as u64; + let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64; (first..rows.end) - .step_by(ROWS_PER_CHUNK) + .step_by(rows_per_chunk) .map(|start| { - start.max(rows.start)..start.saturating_add(ROWS_PER_CHUNK as u64).min(rows.end) + start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end) }) .collect() } @@ -282,7 +298,12 @@ impl PreparedScan { } else { optimize(plan_row_idx_expression(builder.projection, root)?)? }; - let splits = row_chunks(builder.row_range.unwrap_or(0..row_count)); + 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), diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index e5763e21c85..2e463d03150 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -38,10 +38,11 @@ use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use super::PreparedScan; -use super::ROWS_PER_CHUNK; +use super::MAX_ROWS_PER_CHUNK; use super::ScanBuilder; use super::ScanFile; use super::row_chunks; +use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; @@ -309,7 +310,31 @@ async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { #[case::empty(5..5, vec![])] #[case::near_max(u64::MAX - 3..u64::MAX, vec![u64::MAX - 3..u64::MAX])] fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) { - assert_eq!(row_chunks(rows), expected); + assert_eq!(row_chunks(rows, 65_536), expected); +} + +#[rstest] +#[case::empty(0, 0, 65_536)] +#[case::single_chunk(1000, 100_000_000, 65_536)] +#[case::no_bytes(100_000, 0, 65_536)] +#[case::small_rows(1_000_000, 64_000_000, 65_536)] +#[case::large_rows(1_000_000, 384_000_000, 16_384)] +#[case::minimum(1_000_000, 2_000_000_000, 8192)] +fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: usize) { + let locations = [SegmentLocation { + offset: 0, + length: bytes, + alignment: Alignment::none(), + }]; + assert_eq!(rows_per_chunk(rows, &locations), expected); +} + +#[test] +fn smaller_chunks_keep_file_alignment() { + assert_eq!( + row_chunks(7500..33_000, 16_384), + [7500..16_384, 16_384..32_768, 32_768..33_000] + ); } struct CountingSegments { @@ -362,7 +387,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let splits = scan.split_plans(); assert_eq!( splits.len(), - (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(ROWS_PER_CHUNK) + (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(); From 740e1ef080bf7eb268c5dfb186650648b34f94a5 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 16 Oct 2026 17:09:16 +0100 Subject: [PATCH 14/14] test(scan): fix V2 chunk test lint and formatting Signed-off-by: Joe Isaacs --- vortex-layout/src/scan/v2/tests.rs | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 2e463d03150..a910172f87b 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -37,8 +37,8 @@ use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; -use super::PreparedScan; use super::MAX_ROWS_PER_CHUNK; +use super::PreparedScan; use super::ScanBuilder; use super::ScanFile; use super::row_chunks; @@ -371,7 +371,8 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { .write_stream( ArrayContext::empty().into(), Arc::::clone(&segments), - SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))).sendable(), + SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))) + .sendable(), eof, &session, ) @@ -394,18 +395,18 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let before = counted.reads.load(Ordering::Relaxed); let arrays = split.run(Arc::clone(&scan.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; + let code_chunks = + (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; assert_eq!( reads, if index == 0 { STORED_CHUNKS } else { - code_chunks as usize + usize::try_from(code_chunks)? } ); - let expected = PrimitiveArray::from_iter(i32::try_from(rows.start)?..i32::try_from(rows.end)?); + let expected = + PrimitiveArray::from_iter(i32::try_from(rows.start)?..i32::try_from(rows.end)?); assert_arrays_eq!( ChunkedArray::try_new(arrays, layout.dtype().clone())?.into_array(), expected,