vortex_layout/plan/physical/
operator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Contracts and execution context for pipeline sources and operators.
5
6use std::collections::VecDeque;
7
8use smallvec::SmallVec;
9use vortex_array::ArrayRef;
10use vortex_array::ExecutionCtx;
11use vortex_array::buffer::BufferHandle;
12use vortex_error::VortexResult;
13use vortex_session::VortexSession;
14
15use super::pipeline::Arena;
16use super::pipeline::Inlet;
17use super::pipeline::PortId;
18use crate::segments::SegmentId;
19
20/// The output of one non-blocking compute call.
21#[derive(Debug)]
22pub enum Step {
23    /// No output remains for this input. A source has finished entirely.
24    End,
25    /// A batch with more computation available without another input or readiness check.
26    More(ArrayRef),
27    /// The last batch for this input. A stage needs another input; a source needs readiness
28    /// checked again.
29    Final(ArrayRef),
30}
31
32/// Whether a source can compute or needs the scheduler to wait.
33#[derive(Clone, Copy, Debug, PartialEq, Eq)]
34pub enum SourceReady {
35    /// The next compute call can produce output or finish without blocking.
36    Ready,
37    /// Nothing can happen until the condition clears.
38    Blocked(SourceBlocked),
39}
40
41/// Why a source cannot compute yet.
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43pub enum SourceBlocked {
44    /// Wait for bytes from newly queued or outstanding reads.
45    Io,
46    /// The inlet needs more batches or its sink to close it. It must be open and have room
47    /// for its sink to push; debug builds check that the sink can still wake this source.
48    Inlet(usize),
49}
50
51/// The start of a pipeline: produces batches from its inlets and the segments it asks to have
52/// read.
53pub trait Source: Operator {
54    /// Prepares the next [`compute`](Operator::compute) call, after the driver checks outlets.
55    /// Returning [`SourceReady::Ready`] guarantees that call can proceed without blocking.
56    /// Any input or choice made here can be retained by the source for that call.
57    ///
58    /// Queue reads with [`Cx::request_read`] and return
59    /// `SourceReady::Blocked(SourceBlocked::Io)`. The driver issues newly queued reads
60    /// before waiting. Readiness is checked again when bytes are available, including bytes
61    /// already delivered; they are available through [`Cx::take_segment`] in arrival order.
62    fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult<SourceReady> {
63        Ok(SourceReady::Ready)
64    }
65}
66
67/// What an operator is handed by the stage below it.
68#[derive(Debug)]
69pub enum Input {
70    /// No new input: a source is ready, or an operator is continuing after [`Step::More`].
71    None,
72    /// The next batch.
73    Batch(ArrayRef),
74    /// The stage below has finished: no more batches come. Handed over once; [`Input::None`]
75    /// may follow to drain any [`Step::More`] outputs.
76    End,
77}
78
79/// Computes batches without blocking, either as a source or as a stage transforming input.
80pub trait Operator: Send {
81    /// Computes one output. [`Step::More`] requests another call with [`Input::None`] before
82    /// any new input is supplied. [`Step::End`] and [`Step::Final`] finish the current input;
83    /// after [`Input::End`], they finish the stage as well.
84    ///
85    /// Sources receive only [`Input::None`]. Their first call and each call after
86    /// [`Step::Final`] follow [`Source::ready`]; [`Step::End`] finishes the source.
87    fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult<Step>;
88}
89
90/// Distinct inlet indices, with constant-time insertion and removal.
91#[derive(Default)]
92pub(super) struct InletSet {
93    indices: Vec<usize>,
94    positions: Vec<usize>,
95}
96
97impl InletSet {
98    pub(super) fn new(inlets: usize) -> Self {
99        Self {
100            indices: Vec::with_capacity(inlets),
101            positions: vec![usize::MAX; inlets],
102        }
103    }
104
105    pub(super) fn all(inlets: usize) -> Self {
106        Self {
107            indices: (0..inlets).collect(),
108            positions: (0..inlets).collect(),
109        }
110    }
111
112    pub(super) fn insert(&mut self, index: usize) {
113        if self.positions[index] == usize::MAX {
114            self.positions[index] = self.indices.len();
115            self.indices.push(index);
116        }
117    }
118
119    pub(super) fn remove(&mut self, index: usize) {
120        let position = self.positions[index];
121        if position == usize::MAX {
122            return;
123        }
124        self.indices.swap_remove(position);
125        if let Some(&moved) = self.indices.get(position) {
126            self.positions[moved] = position;
127        }
128        self.positions[index] = usize::MAX;
129    }
130
131    fn pop(&mut self) -> Option<usize> {
132        let index = *self.indices.last()?;
133        self.remove(index);
134        Some(index)
135    }
136
137    pub(super) fn indices(&self) -> &[usize] {
138        &self.indices
139    }
140
141    pub(super) fn clear(&mut self) {
142        for &index in &self.indices {
143            self.positions[index] = usize::MAX;
144        }
145        self.indices.clear();
146    }
147}
148
149/// What a source or stage may touch during one call: its pipeline's inlets, the bytes the
150/// source asked for, and the session.
151pub struct Cx<'a> {
152    pub(super) arena: &'a mut Arena,
153    pub(super) inlets: &'a [PortId],
154    pub(super) bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>,
155    pub(super) reads: &'a mut SmallVec<[SegmentId; 1]>,
156    pub(super) session: &'a VortexSession,
157    pub(super) exec: &'a mut ExecutionCtx,
158    pub(super) touched: &'a mut InletSet,
159    pub(super) changed: &'a mut InletSet,
160}
161
162impl Cx<'_> {
163    /// Takes an inlet index whose sink appended data or closed since the last readiness
164    /// check. Each index is reported once, and all inlets are reported on the first check.
165    /// Remaining notifications are cleared after [`Source::ready`] returns. A source must
166    /// account for its own consumption when caching inlet state.
167    pub fn changed_inlet(&mut self) -> Option<usize> {
168        self.changed.pop()
169    }
170
171    /// Queues a segment read during [`Source::ready`]. Return
172    /// `SourceReady::Blocked(SourceBlocked::Io)` after queuing the reads needed for
173    /// this call.
174    pub fn request_read(&mut self, segment: SegmentId) {
175        self.reads.push(segment);
176    }
177
178    /// Whether a requested segment has arrived and can be taken.
179    pub fn has_bytes(&self) -> bool {
180        !self.bytes.is_empty()
181    }
182
183    /// The reading side of inlet `index` of this pipeline's source.
184    pub fn inlet(&mut self, index: usize) -> Inlet<'_> {
185        self.touched.insert(index);
186        self.arena.inlet(self.inlets[index])
187    }
188
189    /// Takes an arrived segment's ID and bytes, oldest arrival first.
190    pub fn take_segment(&mut self) -> Option<(SegmentId, BufferHandle)> {
191        self.bytes.pop_front()
192    }
193
194    /// The session used for decoding and evaluation.
195    pub fn session(&self) -> &VortexSession {
196        self.session
197    }
198
199    /// The execution context for executing arrays.
200    pub fn exec(&mut self) -> &mut ExecutionCtx {
201        self.exec
202    }
203}
vortex_layout/plan/physical/operator.rs

vortex_layout/plan/physical/
operator.rs
