From da97422fd1b03a2dcfbe780a1cbae8592d5e2be2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 13:42:41 +0000 Subject: [PATCH] feat(scan): a planning driver for planners and morsels Runs planners and morsels on the planning protocol, publishes their IO through an `IoService`, and traces their work. Nothing drives it until the V2 scan. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hai26WGsNERxRHXpA3JboE Signed-off-by: Claude --- Cargo.lock | 3 + vortex-scan/Cargo.toml | 6 + vortex-scan/src/planning/driver/mod.rs | 644 ++++++++++++++++ vortex-scan/src/planning/driver/tests.rs | 814 +++++++++++++++++++++ vortex-scan/src/planning/driver/trace.rs | 75 ++ vortex-scan/src/planning/mod.rs | 4 + vortex-scan/src/planning/morsel.rs | 6 + vortex-scan/src/planning/planner.rs | 12 + vortex-scan/src/planning/tests/mod.rs | 6 + vortex-scan/src/planning/tests/scripted.rs | 418 +++++++++++ 10 files changed, 1988 insertions(+) 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/driver/trace.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 411e87b3d8c..0f08930dbcc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12379,7 +12379,9 @@ version = "0.1.0" dependencies = [ "async-trait", "futures", + "parking_lot", "roaring", + "rstest", "tracing", "vortex-array", "vortex-buffer", @@ -12387,6 +12389,7 @@ dependencies = [ "vortex-io", "vortex-mask", "vortex-session", + "vortex-utils", ] [[package]] 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..ea948dff6b4 --- /dev/null +++ b/vortex-scan/src/planning/driver/mod.rs @@ -0,0 +1,644 @@ +// 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 std::time::Instant; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +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_io::request::trace::timestamp_ns; +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; + +mod trace; + +/// 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 { + trace: Option, + 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) { + if let Some(trace) = &self.trace { + trace.event("end", self.queue.work.len(), self.parked.len()); + } + for root in self.roots.values() { + root.io.clear(); + } + } +} + +impl Run { + /// Creates a run with nothing admitted. + pub fn new() -> Self { + let mut run = Self::default(); + if let Some(trace) = trace::Trace::new() { + run.trace = Some(trace); + } + run + } + + /// 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 }); + if let Some(trace) = &self.trace { + tracing::debug!(target: "vortex_scan::driver", run = trace.id, + ts_ns = timestamp_ns(), event = "admit", root = id.0, + file = scope.file_ordinal, row_start = scope.rows.start, row_end = scope.rows.end, + session = self.roots[&id].io.trace_id().unwrap_or(u64::MAX), "scan driver"); + } + 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; + }; + if let Some(trace) = &self.trace { + tracing::debug!(target: "vortex_scan::driver", run = trace.id, + ts_ns = timestamp_ns(), event = "cancel", root = root.0, "scan driver"); + } + 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 { + if let Some(trace) = &self.trace { + trace.event("advance_begin", self.queue.work.len(), self.parked.len()); + } + let progress = self.advance_inner(); + if let Some(trace) = &self.trace { + let event = match &progress { + Ok(Progress::Batch(_)) => "batch", + Ok(Progress::RootDone(_)) => "root_done", + Ok(Progress::Waiting) => "waiting", + Ok(Progress::Idle) => "idle", + Err(_) => "error", + }; + trace.event(event, self.queue.work.len(), self.parked.len()); + } + progress + } + + fn advance_inner(&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 Some(trace) = &self.trace { + tracing::debug!(target: "vortex_scan::driver", run = trace.id, + ts_ns = timestamp_ns(), event = "completion", owner = owner.0, + request = request.0, "scan driver"); + } + 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(); + if let Some(trace) = &self.trace { + trace.work("unpark", &work); + } + 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(), + }); + if let Some(trace) = &self.trace { + trace.work("spawn", &work); + } + self.queue.insert(work); + } + + /// Retires a finished item, and its root once nothing descended from it is live. + fn retire(&mut self, work: &Work) { + if let Some(trace) = &self.trace { + trace.work("retire", 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 + ); + } + if let Some(trace) = &self.trace { + trace.work("park", &work); + } + 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 started = + tracing::enabled!(target: "vortex_scan::compute_timing", tracing::Level::DEBUG) + .then(Instant::now); + let selection = self.trace.as_ref().and_then(|_| match &work.item { + Item::Planner(planner) => planner.trace_selection(), + Item::Morsel(_) => None, + }); + let start = self.trace.as_ref().map(|_| timestamp_ns()); + 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), + }, + }; + if let Some(started) = started { + let compute_ns = u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX); + tracing::debug!( + target: "vortex_scan::compute_timing", + completed_unix_ns = u64::try_from( + SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos() + ).unwrap_or(u64::MAX), + compute_ns, + "scan compute step" + ); + } + if let Some(trace) = &self.trace { + trace.compute(&work, start.unwrap_or_default(), &output); + if let (Some((revision, before_rows)), Item::Planner(planner)) = (selection, &work.item) + && let Some((after_revision, selected_rows)) = planner.trace_selection() + && revision != after_revision + { + tracing::debug!(target: "vortex_scan::driver", run = trace.id, + ts_ns = timestamp_ns(), event = "selection", owner = work.id.0, + root = work.root.0, revision = after_revision, before_rows, selected_rows, + stage = planner.trace_name(), "scan driver"); + } + } + 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 let Some(trace) = &self.trace { + trace.request(work, request); + } + 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..4a936c6d7cb --- /dev/null +++ b/vortex-scan/src/planning/driver/tests.rs @@ -0,0 +1,814 @@ +// 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)] +#[case(IoIntent::Forget)] +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/driver/trace.rs b/vortex-scan/src/planning/driver/trace.rs new file mode 100644 index 00000000000..b67a1c00e9f --- /dev/null +++ b/vortex-scan/src/planning/driver/trace.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use tracing::debug; +use vortex_io::request::IoRequest; +use vortex_io::request::IoTarget; +use vortex_io::request::trace::next_id; +use vortex_io::request::trace::timestamp_ns; + +use super::Item; +use super::Output; +use super::Work; +use crate::planning::morsel::MorselOutput; +use crate::planning::planner::PlannerOutput; + +pub(super) struct Trace { + pub(super) id: u64, +} + +impl Trace { + pub(super) fn new() -> Option { + if !tracing::enabled!(target: "vortex_scan::driver", tracing::Level::DEBUG) { + return None; + } + let trace = Self { id: next_id() }; + trace.event("start", 0, 0); + Some(trace) + } + + pub(super) fn event(&self, event: &str, ready: usize, parked: usize) { + debug!(target: "vortex_scan::driver", run = self.id, ts_ns = timestamp_ns(), event, + ready, parked, "scan driver"); + } + + pub(super) fn work(&self, event: &str, work: &Work) { + let (kind, stage) = match &work.item { + Item::Planner(planner) => ("planner", planner.trace_name()), + Item::Morsel(morsel) => ("morsel", morsel.trace_name()), + }; + debug!(target: "vortex_scan::driver", run = self.id, ts_ns = timestamp_ns(), event, + owner = work.id.0, root = work.root.0, file = work.scope.file_ordinal, + row_start = work.scope.rows.start, row_end = work.scope.rows.end, kind, stage, + outstanding = work.outstanding.len(), state = ?work.state(), "scan driver"); + } + + pub(super) fn compute(&self, work: &Work, start_ns: u64, output: &Output) { + let (outcome, rows) = match output { + Output::Planner(PlannerOutput::Done) | Output::Morsel(MorselOutput::Done) => { + ("done", 0) + } + Output::Planner(PlannerOutput::Continue) | Output::Morsel(MorselOutput::Continue) => { + ("continue", 0) + } + Output::Planner(PlannerOutput::NeedsIO(_)) + | Output::Morsel(MorselOutput::NeedsIO(_)) => ("io", 0), + Output::Planner(PlannerOutput::Planner(..)) => ("planner", 0), + Output::Planner(PlannerOutput::Morsel(..)) => ("morsel", 0), + Output::Morsel(MorselOutput::Batch(array)) => ("batch", array.len()), + }; + debug!(target: "vortex_scan::driver", run = self.id, ts_ns = timestamp_ns(), + event = "compute", owner = work.id.0, root = work.root.0, start_ns, outcome, rows, + state = ?work.state(), "scan driver"); + } + + pub(super) fn request(&self, work: &Work, request: &IoRequest) { + let (target, offset, length) = match request.target { + IoTarget::Range { offset, len, .. } => ("range", offset, len), + IoTarget::Size => ("size", 0, 0), + }; + debug!(target: "vortex_scan::driver", run = self.id, ts_ns = timestamp_ns(), + event = "request", owner = work.id.0, root = work.root.0, + request = request.request.0, intent = ?request.intent, target, offset, length, + "scan driver"); + } +} 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/morsel.rs b/vortex-scan/src/planning/morsel.rs index 407c58567e0..3cd92d34f3e 100644 --- a/vortex-scan/src/planning/morsel.rs +++ b/vortex-scan/src/planning/morsel.rs @@ -3,6 +3,7 @@ //! The morsel side of the protocol: array-producing work and its outputs. +use std::any::type_name; use std::fmt; use vortex_array::ArrayRef; @@ -35,6 +36,11 @@ pub enum MorselOutput { /// `NeedsCompute`. Batches preserve selected input-row order within and across calls. A morsel /// moves between threads with the run that owns it. pub trait Morsel: IoConsumer + Send { + /// Identifies the stage in diagnostic traces without inspecting its inputs. + fn trace_name(&self) -> &'static str { + type_name::() + } + /// Reports what the morsel needs next without doing work. fn state(&self) -> State; diff --git a/vortex-scan/src/planning/planner.rs b/vortex-scan/src/planning/planner.rs index 91d875ff317..4043dcf21e3 100644 --- a/vortex-scan/src/planning/planner.rs +++ b/vortex-scan/src/planning/planner.rs @@ -3,6 +3,7 @@ //! The planner side of the protocol: work scopes, object state, and planner outputs. +use std::any::type_name; use std::fmt; use std::ops::Range; @@ -61,6 +62,17 @@ pub enum PlannerOutput { /// when `state()` reports `NeedsCompute` and may produce one child, one IO batch, a checkpoint, /// or finish. A planner moves between threads with the run that owns it. pub trait Planner: IoConsumer + Send { + /// Identifies the stage in diagnostic traces without inspecting its inputs. + fn trace_name(&self) -> &'static str { + type_name::() + } + + /// Optional selection revision and surviving row count for diagnostic traces. + /// The revision changes whenever a filtering or pruning decision completes. + fn trace_selection(&self) -> Option<(usize, usize)> { + None + } + /// Reports what the planner needs next without doing work. fn state(&self) -> State; 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, + } +}