vortex_scan/planning/driver/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! A single-worker driver for the protocol.
5//!
6//! A [`Run`] drives any number of root planners, each with its own IO session, and hands their
7//! batches to its caller as they are produced. Runnable items are prioritised in DFS preorder,
8//! with roots ordered by admission and children by emission. Each visit takes at most one CPU
9//! step. A parent's continuation follows its emitted child's entire subtree. Waiting work is
10//! skipped and regains its priority when IO makes it runnable. Published IO is registered
11//! immediately and state is inspected again before parking. Output follows emission order,
12//! not global row order. An earlier runnable branch can starve later branches.
13//!
14//! A run never blocks: [`Run::advance`] returns [`Progress::Waiting`] when every live item waits
15//! for IO, and the caller decides how to wait, by awaiting [`Run::poll_completion`] or by handing
16//! a completion it took from a source to [`Run::complete`]. [`Driver::run`] is the blocking
17//! convenience for one root.
18
19use std::collections::BTreeMap;
20use std::collections::VecDeque;
21use std::sync::Arc;
22use std::task::Context;
23use std::task::Poll;
24
25use vortex_array::ArrayRef;
26use vortex_error::VortexResult;
27use vortex_error::vortex_bail;
28use vortex_error::vortex_err;
29use vortex_io::request::Completion;
30use vortex_io::request::IoBatch;
31use vortex_io::request::IoConsumer;
32use vortex_io::request::IoIntent;
33use vortex_io::request::IoOwnerId;
34use vortex_io::request::IoRequestId;
35use vortex_io::request::IoResult;
36use vortex_io::request::IoSource;
37use vortex_io::request::IoTarget;
38use vortex_utils::aliases::hash_map::Entry;
39use vortex_utils::aliases::hash_map::HashMap;
40
41use crate::planning::morsel::Morsel;
42use crate::planning::morsel::MorselOutput;
43use crate::planning::planner::Planner;
44use crate::planning::planner::PlannerOutput;
45use crate::planning::planner::State;
46use crate::planning::planner::WorkScope;
47
48/// Identifies one root admitted to a [`Run`]. Unique within the run, in admission order.
49#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
50pub struct RootId(pub u64);
51
52/// One array produced by a morsel, with the scope the morsel was created for.
53pub struct Batch {
54    /// The root the morsel descends from.
55    pub root: RootId,
56    /// Scope of the morsel that produced the array.
57    pub scope: WorkScope,
58    /// The non-empty array.
59    pub array: ArrayRef,
60}
61
62/// What a call to [`Run::advance`] left behind.
63pub enum Progress {
64    /// A morsel's batch. One morsel's batches arrive in the order it emitted them.
65    Batch(Batch),
66    /// Every item descended from the root has retired, and its IO session was cleared.
67    RootDone(RootId),
68    /// Every live item waits for IO. Await [`Run::poll_completion`], or pass the next completion
69    /// to [`Run::complete`].
70    Waiting,
71    /// Nothing is live: every admitted root has finished.
72    Idle,
73}
74
75enum Item {
76    Planner(Box<dyn Planner>),
77    Morsel(Box<dyn Morsel>),
78}
79
80/// What one item's `compute()` produced.
81enum Output {
82    Planner(PlannerOutput),
83    Morsel(MorselOutput),
84}
85
86/// One live planner or morsel. Boxed in the run queue and the parked map, so scheduling moves a
87/// pointer.
88struct Work {
89    id: IoOwnerId,
90    root: RootId,
91    scope: WorkScope,
92    item: Item,
93    /// Stable child-ordinal path, beginning with the root's admission ordinal.
94    path: Vec<u64>,
95    /// The parent continuation sorts after all previously emitted child subtrees.
96    next_child: u64,
97    /// Published `Fetch` requests not yet delivered, with the target each must be answered with.
98    outstanding: HashMap<IoRequestId, IoTarget>,
99}
100
101impl Work {
102    fn priority(&self) -> Vec<u64> {
103        let mut priority = self.path.clone();
104        priority.push(self.next_child);
105        priority
106    }
107
108    // Child k takes the continuation's old key as its path. Its subtree then sorts before
109    // the parent's new continuation at k + 1, even if descendants are discovered after IO.
110    fn child_path(&mut self) -> Vec<u64> {
111        let path = self.priority();
112        self.next_child += 1;
113        path
114    }
115
116    fn state(&self) -> State {
117        match &self.item {
118            Item::Planner(planner) => planner.state(),
119            Item::Morsel(morsel) => morsel.state(),
120        }
121    }
122
123    fn consumer(&mut self) -> &mut dyn IoConsumer {
124        match &mut self.item {
125            Item::Planner(planner) => &mut **planner,
126            Item::Morsel(morsel) => &mut **morsel,
127        }
128    }
129
130    /// Checks `result` against the outstanding request it answers and hands it to the item.
131    fn deliver(&mut self, request: IoRequestId, result: IoResult) -> VortexResult<()> {
132        let Some(target) = self.outstanding.remove(&request) else {
133            vortex_bail!(
134                "completion for {request:?} of {:?}, which is not outstanding",
135                self.id
136            );
137        };
138        if !result.matches(&target) {
139            vortex_bail!("IO source answered {target:?} with {}", result.kind());
140        }
141        self.consumer().set_io_result(request, result);
142        Ok(())
143    }
144}
145
146/// Keeps scheduling order separate from lookup for deliveries to runnable work.
147#[derive(Default)]
148struct ReadyQueue {
149    priorities: BTreeMap<Vec<u64>, IoOwnerId>,
150    work: HashMap<IoOwnerId, Box<Work>>,
151}
152
153impl ReadyQueue {
154    fn insert(&mut self, work: Box<Work>) {
155        self.priorities.insert(work.priority(), work.id);
156        self.work.insert(work.id, work);
157    }
158
159    fn pop(&mut self) -> Option<Box<Work>> {
160        let (_, owner) = self.priorities.pop_first()?;
161        self.work.remove(&owner)
162    }
163
164    fn is_empty(&self) -> bool {
165        self.priorities.is_empty()
166    }
167
168    fn cancel(&mut self, root: RootId) {
169        self.work.retain(|_, work| work.root != root);
170        self.priorities
171            .retain(|_, owner| self.work.contains_key(owner));
172    }
173}
174
175/// Poll IO even when one high-priority branch keeps producing CPU work.
176const IO_POLL_INTERVAL: usize = 64;
177
178/// One admitted root: the session its descendants read through, and how many of them are live.
179struct Root {
180    io: Arc<dyn IoSource>,
181    live: usize,
182}
183
184/// A set of root planners advanced by its caller on one thread at a time.
185///
186/// Dropping a run clears the IO session of every root still live, even when it ends in an error.
187#[derive(Default)]
188pub struct Run {
189    roots: HashMap<RootId, Root>,
190    next_root: u64,
191    next_owner: u64,
192    queue: ReadyQueue,
193    parked: HashMap<IoOwnerId, Box<Work>>,
194    /// Batches and finished roots not yet returned by [`advance`](Self::advance).
195    events: VecDeque<Progress>,
196    /// Visits left before the sources are polled for completions again.
197    visits_until_poll: usize,
198    step_limit: Option<usize>,
199    steps: usize,
200}
201
202impl Drop for Run {
203    fn drop(&mut self) {
204        for root in self.roots.values() {
205            root.io.clear();
206        }
207    }
208}
209
210impl Run {
211    /// Creates a run with nothing admitted.
212    pub fn new() -> Self {
213        Self::default()
214    }
215
216    /// Fails with an error naming the limit if more than `steps` visits are needed. Tests use it
217    /// to catch loops.
218    pub fn with_step_limit(mut self, steps: usize) -> Self {
219        self.step_limit = Some(steps);
220        self
221    }
222
223    /// Admits `root` and everything it spawns, reading through `io`.
224    ///
225    /// The session belongs to the root: the run submits the requests of every descendant to it,
226    /// and clears it once they have all retired or the root is cancelled.
227    pub fn admit(
228        &mut self,
229        root: Box<dyn Planner>,
230        scope: WorkScope,
231        io: Arc<dyn IoSource>,
232    ) -> RootId {
233        let id = RootId(self.next_root);
234        self.next_root += 1;
235        self.roots.insert(id, Root { io, live: 0 });
236        self.spawn(id, scope, Item::Planner(root), vec![id.0]);
237        id
238    }
239
240    /// Drops every live item descended from `root` and clears its session. Batches the root
241    /// already produced are still returned, followed by no [`Progress::RootDone`].
242    pub fn cancel(&mut self, root: RootId) {
243        let Some(cancelled) = self.roots.remove(&root) else {
244            return;
245        };
246        self.queue.cancel(root);
247        self.parked.retain(|_, work| work.root != root);
248        cancelled.io.clear();
249    }
250
251    /// The number of admitted roots that have not finished.
252    pub fn live_roots(&self) -> usize {
253        self.roots.len()
254    }
255
256    /// Runs ready work, taking completions the sources already have, until there is something to
257    /// hand back or everything left waits for IO. Never blocks.
258    pub fn advance(&mut self) -> VortexResult<Progress> {
259        loop {
260            if let Some(event) = self.events.pop_front() {
261                return Ok(event);
262            }
263            if self.visits_until_poll == 0 || self.queue.is_empty() {
264                self.take_completions()?;
265                self.visits_until_poll = IO_POLL_INTERVAL;
266                if self.queue.is_empty() {
267                    return Ok(if self.parked.is_empty() {
268                        Progress::Idle
269                    } else {
270                        Progress::Waiting
271                    });
272                }
273            }
274            self.visits_until_poll -= 1;
275            let Some(work) = self.queue.pop() else {
276                continue;
277            };
278            self.steps += 1;
279            if let Some(limit) = self.step_limit
280                && self.steps > limit
281            {
282                vortex_bail!("driver exceeded its step limit of {limit}");
283            }
284            self.visit(work)?;
285        }
286    }
287
288    /// Delivers a completion taken from the session of one of the run's roots.
289    pub fn complete(&mut self, completion: Completion) -> VortexResult<()> {
290        let Completion {
291            owner,
292            request,
293            result,
294        } = completion;
295        let result = result?;
296        if let Entry::Occupied(mut parked) = self.parked.entry(owner) {
297            let work = parked.get_mut();
298            work.deliver(request, result)?;
299            // An item that still waits for other requests stays parked.
300            if work.state() != State::Waiting {
301                let work = parked.remove();
302                self.queue.insert(work);
303            } else if work.outstanding.is_empty() {
304                vortex_bail!(
305                    "work for scope {:?} waits with no outstanding fetch",
306                    work.scope
307                );
308            }
309            return Ok(());
310        }
311        // An item that published fetches and had CPU work left is still in the run queue.
312        let Some(work) = self.queue.work.get_mut(&owner) else {
313            vortex_bail!("completion for unknown work {owner:?}");
314        };
315        work.deliver(request, result)
316    }
317
318    /// Waits for a completion from any root's session, registering `cx` to be woken by the next
319    /// one, and delivers every completion that is ready.
320    ///
321    /// Call it after [`advance`](Self::advance) returned [`Progress::Waiting`].
322    pub fn poll_completion(&mut self, cx: &mut Context<'_>) -> Poll<VortexResult<()>> {
323        let sources: Vec<_> = self
324            .roots
325            .values()
326            .map(|root| Arc::clone(&root.io))
327            .collect();
328        if sources.is_empty() {
329            return Poll::Ready(Err(vortex_err!(
330                "the run is waiting for IO with no live root"
331            )));
332        }
333        let mut delivered = false;
334        for io in sources {
335            match io.poll_completion(cx) {
336                Poll::Ready(Ok(completion)) => {
337                    if let Err(error) = self.complete(completion) {
338                        return Poll::Ready(Err(error));
339                    }
340                    delivered = true;
341                }
342                Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
343                Poll::Pending => {}
344            }
345        }
346        if delivered {
347            Poll::Ready(Ok(()))
348        } else {
349            Poll::Pending
350        }
351    }
352
353    /// Delivers the completions every root's session has ready, without blocking.
354    fn take_completions(&mut self) -> VortexResult<()> {
355        if self.roots.len() == 1 {
356            // The common case, without collecting the sessions.
357            let Some(io) = self.roots.values().next().map(|root| Arc::clone(&root.io)) else {
358                return Ok(());
359            };
360            while let Some(completion) = io.poll()? {
361                self.complete(completion)?;
362            }
363            return Ok(());
364        }
365        let sources: Vec<_> = self
366            .roots
367            .values()
368            .map(|root| Arc::clone(&root.io))
369            .collect();
370        for io in sources {
371            while let Some(completion) = io.poll()? {
372                self.complete(completion)?;
373            }
374        }
375        Ok(())
376    }
377
378    fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec<u64>) {
379        let id = IoOwnerId(self.next_owner);
380        self.next_owner += 1;
381        if let Some(root) = self.roots.get_mut(&root) {
382            root.live += 1;
383        }
384        let work = Box::new(Work {
385            id,
386            root,
387            scope,
388            item,
389            path,
390            next_child: 0,
391            outstanding: HashMap::default(),
392        });
393        self.queue.insert(work);
394    }
395
396    /// Retires a finished item, and its root once nothing descended from it is live.
397    fn retire(&mut self, work: &Work) {
398        let Entry::Occupied(mut root) = self.roots.entry(work.root) else {
399            return;
400        };
401        root.get().io.release(work.id);
402        root.get_mut().live -= 1;
403        if root.get().live == 0 {
404            root.remove().io.clear();
405            self.events.push_back(Progress::RootDone(work.root));
406        }
407    }
408
409    /// Parks an item that waits for IO, which is a protocol error with no fetch outstanding.
410    fn park(&mut self, work: Box<Work>) -> VortexResult<()> {
411        if work.outstanding.is_empty() {
412            vortex_bail!(
413                "work for scope {:?} waits with no outstanding fetch",
414                work.scope
415            );
416        }
417        self.parked.insert(work.id, work);
418        Ok(())
419    }
420
421    /// Requeues, parks, or retires an item according to the state it reports.
422    fn settle(&mut self, work: Box<Work>) -> VortexResult<()> {
423        match work.state() {
424            State::Done => self.retire(&work),
425            State::NeedsCompute => self.queue.insert(work),
426            State::Waiting => self.park(work)?,
427        }
428        Ok(())
429    }
430
431    /// Visits one item: computes once if it is ready, and settles it otherwise.
432    fn visit(&mut self, mut work: Box<Work>) -> VortexResult<()> {
433        let output = match &mut work.item {
434            Item::Planner(planner) => match planner.state() {
435                State::NeedsCompute => Output::Planner(planner.compute()?),
436                _ => return self.settle(work),
437            },
438            Item::Morsel(morsel) => match morsel.state() {
439                State::NeedsCompute => Output::Morsel(morsel.compute()?),
440                _ => return self.settle(work),
441            },
442        };
443        match output {
444            Output::Planner(PlannerOutput::Done) | Output::Morsel(MorselOutput::Done) => {
445                self.retire(&work);
446            }
447            Output::Planner(PlannerOutput::Continue) | Output::Morsel(MorselOutput::Continue) => {
448                self.queue.insert(work);
449            }
450            Output::Planner(PlannerOutput::NeedsIO(batch))
451            | Output::Morsel(MorselOutput::NeedsIO(batch)) => {
452                self.submit(&mut work, batch)?;
453                self.settle(work)?;
454            }
455            Output::Planner(PlannerOutput::Planner(scope, child)) => {
456                let path = work.child_path();
457                self.spawn(work.root, scope, Item::Planner(child), path);
458                self.queue.insert(work);
459            }
460            Output::Planner(PlannerOutput::Morsel(scope, morsel)) => {
461                let path = work.child_path();
462                self.spawn(work.root, scope, Item::Morsel(morsel), path);
463                self.queue.insert(work);
464            }
465            Output::Morsel(MorselOutput::Batch(array)) => {
466                if array.is_empty() {
467                    vortex_bail!("morsel for scope {:?} returned an empty batch", work.scope);
468                }
469                self.events.push_back(Progress::Batch(Batch {
470                    root: work.root,
471                    scope: work.scope.clone(),
472                    array,
473                }));
474                self.queue.insert(work);
475            }
476        }
477        Ok(())
478    }
479
480    /// Registers the requests an item published with its root's session.
481    fn submit(&mut self, work: &mut Work, batch: IoBatch) -> VortexResult<()> {
482        if batch.is_empty() {
483            vortex_bail!("compute() published an empty batch");
484        }
485        for request in &batch {
486            if request.intent == IoIntent::Fetch
487                && work
488                    .outstanding
489                    .insert(request.request, request.target)
490                    .is_some()
491            {
492                vortex_bail!(
493                    "fetch {:?} of {:?} was published while it is outstanding",
494                    request.request,
495                    work.id
496                );
497            }
498        }
499        let Some(root) = self.roots.get(&work.root) else {
500            vortex_bail!("work {:?} outlived its root {:?}", work.id, work.root);
501        };
502        root.io.submit(work.id, batch)
503    }
504}
505
506/// Drives one root to completion on the calling thread, blocking while it waits for IO.
507pub struct Driver {
508    io: Arc<dyn IoSource>,
509    step_limit: Option<usize>,
510}
511
512impl Driver {
513    /// Creates a driver that submits every request to `io`.
514    pub fn new(io: Arc<dyn IoSource>) -> Self {
515        Self {
516            io,
517            step_limit: None,
518        }
519    }
520
521    /// Fails with an error naming the limit if more than `steps` visits are needed. Tests use it
522    /// to catch loops.
523    pub fn with_step_limit(self, steps: usize) -> Self {
524        Self {
525            step_limit: Some(steps),
526            ..self
527        }
528    }
529
530    /// Runs `root` and everything it spawns, returning batches in emission order.
531    pub fn run(&self, root: Box<dyn Planner>) -> VortexResult<Vec<Batch>> {
532        let mut run = Run::new();
533        run.step_limit = self.step_limit;
534        let scope = WorkScope {
535            file_ordinal: 0,
536            rows: 0..0,
537        };
538        run.admit(root, scope, Arc::clone(&self.io));
539        let mut batches = Vec::new();
540        loop {
541            match run.advance()? {
542                Progress::Batch(batch) => batches.push(batch),
543                Progress::RootDone(_) | Progress::Idle => return Ok(batches),
544                Progress::Waiting => run.complete(self.io.wait()?)?,
545            }
546        }
547    }
548}
549
550#[cfg(test)]
551mod tests;
vortex_scan/planning/driver/mod.rs

vortex_scan/planning/driver/
mod.rs
