vortex_layout/plan/physical/pipeline/
synthetic.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Plans whose chains are built by hand, for exercising pipelines without layouts.
5//!
6//! A [`SyntheticPlan`] carries a closure that compiles it. A scan treats it like any other
7//! plan, so a test or benchmark can wire hand-written sources under the real operators and run
8//! the whole thing through a [`Scan`](super::Scan), answering reads from memory.
9
10use std::fmt;
11use std::ops::Range;
12use std::sync::Arc;
13
14use parking_lot::Mutex;
15use vortex_array::EmptyMetadata;
16use vortex_array::IntoArray;
17use vortex_array::buffer::BufferHandle;
18use vortex_array::dtype::DType;
19use vortex_array::dtype::Nullability;
20use vortex_array::dtype::PType;
21use vortex_buffer::Buffer;
22use vortex_error::VortexResult;
23use vortex_error::vortex_err;
24use vortex_mask::Mask;
25use vortex_session::registry::CachedId;
26
27use super::Chain;
28use super::Compiler;
29use crate::plan::physical::Plan;
30use crate::plan::physical::PlanChildren;
31use crate::plan::physical::PlanId;
32use crate::plan::physical::PlanParts;
33use crate::plan::physical::PlanRef;
34use crate::plan::physical::PlanVTable;
35use crate::plan::physical::operator::Cx;
36use crate::plan::physical::operator::Input;
37use crate::plan::physical::operator::Operator;
38use crate::plan::physical::operator::Source;
39use crate::plan::physical::operator::SourceBlocked;
40use crate::plan::physical::operator::SourceReady;
41use crate::plan::physical::operator::Step;
42use crate::segments::SegmentId;
43
44/// Compiles a [`SyntheticPlan`] over some rows, restricted to a mask.
45pub type ChainBuilder =
46    dyn Fn(Range<u64>, &Mask, &mut Compiler<'_>) -> VortexResult<Option<Chain>> + Send + Sync;
47
48/// A plan operator whose chain is built by a closure.
49#[derive(Clone, Debug)]
50pub struct Synthetic;
51
52/// The chain builder of a [`Synthetic`] plan.
53#[derive(Clone)]
54pub struct SyntheticData {
55    builder: Arc<ChainBuilder>,
56}
57
58impl fmt::Debug for SyntheticData {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        f.write_str("SyntheticData")
61    }
62}
63
64/// A plan whose chain is built by a closure.
65pub type SyntheticPlan = Plan<Synthetic>;
66
67impl SyntheticPlan {
68    /// A plan of `row_count` rows of `dtype`, with `children`, compiled by `builder`.
69    pub fn new(
70        dtype: DType,
71        row_count: u64,
72        children: Vec<PlanRef>,
73        builder: impl Fn(Range<u64>, &Mask, &mut Compiler<'_>) -> VortexResult<Option<Chain>>
74        + Send
75        + Sync
76        + 'static,
77    ) -> Self {
78        PlanParts {
79            vtable: Synthetic,
80            dtype,
81            row_count,
82            children: children.into(),
83            data: SyntheticData {
84                builder: Arc::new(builder),
85            },
86        }
87        .into_typed()
88    }
89}
90
91impl PlanVTable for Synthetic {
92    type PlanData = SyntheticData;
93    type Metadata = EmptyMetadata;
94
95    fn id(&self) -> PlanId {
96        static ID: CachedId = CachedId::new("vortex.plan.synthetic");
97        *ID
98    }
99
100    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
101        None
102    }
103
104    fn with_children(
105        _plan: &Plan<Self>,
106        _children: &PlanChildren,
107        _data: &mut Self::PlanData,
108    ) -> VortexResult<()> {
109        Ok(())
110    }
111
112    fn compile(
113        plan: &Plan<Self>,
114        rows: Range<u64>,
115        mask: &Mask,
116        compiler: &mut Compiler<'_>,
117    ) -> VortexResult<Option<Chain>> {
118        (plan.data().builder)(rows, mask, compiler)
119    }
120}
121
122/// The dtype every synthetic source produces: its rows' indices, as `u64`.
123pub fn row_dtype() -> DType {
124    DType::Primitive(PType::U64, Nullability::NonNullable)
125}
126
127/// A source that emits its rows' indices, in order, cut into pieces of the given lengths, one
128/// piece per call, optionally after a read of a segment it does not look at.
129///
130/// A selected source emits only the rows its mask selects, as every plan does. A dense source
131/// ignores its mask and emits every row of its range, as a bare segment scan does. Piece
132/// lengths count emitted rows; the last piece takes whatever remains.
133pub struct RowSource {
134    /// Remaining indices to emit, in order.
135    remaining: Vec<u64>,
136    /// Lengths of the pieces still to emit, last first.
137    pieces: Vec<usize>,
138    /// A segment to read before emitting anything, standing in for a real read, until it has
139    /// been asked for.
140    segment: Option<SegmentId>,
141    /// The segment has been asked for and its bytes not yet taken.
142    waiting: bool,
143}
144
145impl RowSource {
146    /// A plan of `row_count` rows compiled to a [`RowSource`] emitting pieces of the given
147    /// lengths.
148    pub fn plan(
149        row_count: u64,
150        pieces: Vec<usize>,
151        segment: Option<SegmentId>,
152        dense: bool,
153    ) -> PlanRef {
154        SyntheticPlan::new(row_dtype(), row_count, Vec::new(), move |rows, mask, _| {
155            let remaining: Vec<u64> = if dense || mask.all_true() {
156                (rows.start..rows.end).collect()
157            } else {
158                (rows.start..rows.end)
159                    .zip(mask.iter())
160                    .filter_map(|(row, selected)| selected.then_some(row))
161                    .collect()
162            };
163            if remaining.is_empty() {
164                return Ok(None);
165            }
166            Ok(Some(Chain::new(RowSource {
167                remaining,
168                pieces: pieces.iter().rev().copied().collect(),
169                segment,
170                waiting: false,
171            })))
172        })
173        .into_plan()
174    }
175}
176
177impl Source for RowSource {
178    fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult<SourceReady> {
179        if let Some(segment) = self.segment.take() {
180            self.waiting = true;
181            cx.request_read(segment);
182            return Ok(SourceReady::Blocked(SourceBlocked::Io));
183        }
184        Ok(if self.waiting && !cx.has_bytes() {
185            SourceReady::Blocked(SourceBlocked::Io)
186        } else {
187            SourceReady::Ready
188        })
189    }
190}
191
192impl Operator for RowSource {
193    fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult<Step> {
194        if self.waiting {
195            self.waiting = false;
196            cx.take_segment()
197                .ok_or_else(|| vortex_err!("RowSource ran without its bytes"))?;
198        }
199        if self.remaining.is_empty() {
200            return Ok(Step::End);
201        }
202        let len = self
203            .pieces
204            .pop()
205            .unwrap_or(self.remaining.len())
206            .min(self.remaining.len());
207        let piece: Vec<u64> = self.remaining.drain(..len).collect();
208        Ok(Step::Final(Buffer::from(piece).into_array()))
209    }
210}
211
212/// What a [`Probe`] saw of its inlets on one call: for each, the batches queued and whether
213/// its sink had closed it.
214pub type ProbeRun = Vec<(usize, bool)>;
215
216/// A source over one inlet per child that records what it sees on each readiness check, then
217/// passes its first inlet through, one batch per call, and drops what the others hold.
218pub struct Probe {
219    inlets: usize,
220    runs: Arc<Mutex<Vec<ProbeRun>>>,
221}
222
223impl Probe {
224    /// A plan over `children` compiled to a probe whose inlets hold `capacity` batches and
225    /// which logs into `runs`.
226    pub fn plan(
227        children: Vec<PlanRef>,
228        capacity: usize,
229        runs: Arc<Mutex<Vec<ProbeRun>>>,
230    ) -> PlanRef {
231        let row_count = children.first().map_or(0, |child| child.row_count());
232        SyntheticPlan::new(
233            row_dtype(),
234            row_count,
235            children.clone(),
236            move |rows, mask, compiler| {
237                let mut chains = Vec::with_capacity(children.len());
238                for child in &children {
239                    match compiler.compile(child, rows.clone(), mask)? {
240                        Some(chain) => chains.push(chain),
241                        None => return Ok(None),
242                    }
243                }
244                let probe = Probe {
245                    inlets: chains.len(),
246                    runs: Arc::clone(&runs),
247                };
248                Ok(Some(compiler.join_with(chains, |_| capacity, probe)))
249            },
250        )
251        .into_plan()
252    }
253}
254
255impl Source for Probe {
256    fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult<SourceReady> {
257        let run: ProbeRun = (0..self.inlets)
258            .map(|index| {
259                let inlet = cx.inlet(index);
260                (inlet.len(), inlet.closed())
261            })
262            .collect();
263        self.runs.lock().push(run);
264        for index in 1..self.inlets {
265            while cx.inlet(index).take().is_some() {}
266        }
267        let first = cx.inlet(0);
268        if !first.is_empty() {
269            return Ok(SourceReady::Ready);
270        }
271        if first.closed() {
272            for index in 1..self.inlets {
273                if !cx.inlet(index).closed() {
274                    return Ok(SourceReady::Blocked(SourceBlocked::Inlet(index)));
275                }
276            }
277            return Ok(SourceReady::Ready);
278        }
279        Ok(SourceReady::Blocked(SourceBlocked::Inlet(0)))
280    }
281}
282
283impl Operator for Probe {
284    fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult<Step> {
285        Ok(match cx.inlet(0).take() {
286            Some(batch) => Step::Final(batch),
287            None => Step::End,
288        })
289    }
290}
291
292/// An empty segment, for sources whose read is only a scheduling event.
293pub fn empty_segment() -> BufferHandle {
294    BufferHandle::new_host(Buffer::<u8>::empty().into_byte_buffer())
295}
vortex_layout/plan/physical/pipeline/synthetic.rs

vortex_layout/plan/physical/pipeline/
synthetic.rs
