vortex_layout/plan/physical/pipeline/
compile.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Compiling plans into pipelines, and counting the readers of each segment so segments read
5//! more than once are decoded once.
6
7use std::ops::Range;
8
9use rustc_hash::FxHashMap;
10use smallvec::SmallVec;
11use vortex_error::VortexResult;
12use vortex_mask::Mask;
13use vortex_session::VortexSession;
14
15use super::DEFAULT_CAPACITY;
16use super::port::PortId;
17use super::port::SourceEnd;
18use super::scan::Core;
19use super::values::ValueCache;
20use crate::plan::physical::ConcatPlan;
21use crate::plan::physical::PlanRef;
22use crate::plan::physical::SegmentScanPlan;
23use crate::plan::physical::operator::Operator;
24use crate::plan::physical::operator::Source;
25use crate::plan::physical::operators::ConcatSource;
26use crate::plan::physical::operators::ScanSource;
27use crate::plan::physical::operators::SelectStage;
28use crate::segments::SegmentId;
29
30/// A pipeline under construction: a source and the stages after it, not yet given outlets, so a
31/// parent plan can add stages before it becomes a pipeline.
32pub struct Chain {
33    pub(crate) source: Box<dyn Source>,
34    pub(crate) stages: Vec<Box<dyn Operator>>,
35    pub(crate) inlets: Vec<PortId>,
36}
37
38impl Chain {
39    /// A chain of `source` alone.
40    pub fn new(source: impl Source + 'static) -> Self {
41        Self {
42            source: Box::new(source),
43            stages: Vec::new(),
44            inlets: Vec::new(),
45        }
46    }
47
48    /// A chain passing the batches of `port` through.
49    pub(crate) fn port(port: PortId) -> Self {
50        Self {
51            source: Box::new(ConcatSource::new(1)),
52            stages: Vec::new(),
53            inlets: vec![port],
54        }
55    }
56
57    /// This chain with `stage` after its last stage.
58    pub fn with(mut self, stage: impl Operator + 'static) -> Self {
59        self.stages.push(Box::new(stage));
60        self
61    }
62}
63
64impl Core {
65    /// Compiles `plan` over `rows`, restricted to `mask`. See [`Compiler::compile`].
66    pub(crate) fn compile(
67        &mut self,
68        plan: &PlanRef,
69        rows: Range<u64>,
70        mask: &Mask,
71    ) -> VortexResult<Option<Chain>> {
72        Compiler { core: self }.compile(plan, rows, mask)
73    }
74}
75
76/// What a plan compiles its children through. See
77/// [`PlanVTable::compile`](crate::plan::physical::PlanVTable::compile).
78pub struct Compiler<'a> {
79    pub(crate) core: &'a mut Core,
80}
81
82impl Compiler<'_> {
83    /// Compiles `plan` over `rows` of its domain, restricted to `mask`, into a chain producing
84    /// the rows `mask` selects, in order, or `None` when that is no row.
85    pub fn compile(
86        &mut self,
87        plan: &PlanRef,
88        rows: Range<u64>,
89        mask: &Mask,
90    ) -> VortexResult<Option<Chain>> {
91        if rows.start >= rows.end {
92            return Ok(None);
93        }
94        plan.compile(rows, mask, self)
95    }
96
97    /// A chain reading `chains` through `source`: each chain becomes a pipeline writing a new
98    /// port, inlet `i` of `source` reading chain `i`, holding [`DEFAULT_CAPACITY`] batches.
99    pub fn join(&mut self, chains: Vec<Chain>, source: impl Source + 'static) -> Chain {
100        self.join_with(chains, |_| DEFAULT_CAPACITY, source)
101    }
102
103    /// [`join`](Self::join), with inlet `i` holding `capacity(i)` batches before its sink is
104    /// blocked: `usize::MAX` for an inlet the source needs whole. An empty inlet always has
105    /// room.
106    pub fn join_with(
107        &mut self,
108        chains: Vec<Chain>,
109        capacity: impl Fn(usize) -> usize,
110        source: impl Source + 'static,
111    ) -> Chain {
112        let inlets = chains
113            .into_iter()
114            .enumerate()
115            .map(|(index, chain)| {
116                let port = self
117                    .core
118                    .arena
119                    .create(capacity(index), None, SourceEnd::Unclaimed);
120                self.core.add_pipeline(chain, smallvec::smallvec![port]);
121                port
122            })
123            .collect();
124        Chain {
125            source: Box::new(source),
126            stages: Vec::new(),
127            inlets,
128        }
129    }
130
131    /// The global row index of the scanned plan's row zero.
132    pub fn row_offset(&self) -> u64 {
133        self.core.row_offset
134    }
135
136    /// The session used for decoding and evaluation.
137    pub fn session(&self) -> &VortexSession {
138        &self.core.session
139    }
140
141    /// Whole values kept between the scans of the split compiling.
142    pub(crate) fn values(&self) -> &ValueCache {
143        &self.core.values
144    }
145
146    /// The rows `rows` of `scan`'s segment, filtered by `filter`, read by the split compiling.
147    /// A segment several readers read is decoded once and shared.
148    pub(crate) fn scan(
149        &mut self,
150        scan: &SegmentScanPlan,
151        rows: Range<u64>,
152        filter: Option<Mask>,
153    ) -> VortexResult<Option<Chain>> {
154        let filter = match filter {
155            Some(filter) if filter.all_false() => return Ok(None),
156            Some(filter) if filter.all_true() => None,
157            filter => filter,
158        };
159        let slice = if rows.start == 0 && rows.end == scan.row_count() {
160            None
161        } else {
162            Some(usize::try_from(rows.start)?..usize::try_from(rows.end)?)
163        };
164        let split = self.core.split;
165        Ok(Some(match self.claim(scan, split) {
166            None => Chain::new(ScanSource::new(scan.clone(), slice, filter)),
167            Some(port) => {
168                let chain = Chain::port(port);
169                if slice.is_none() && filter.is_none() {
170                    chain
171                } else {
172                    chain.with(SelectStage::new(slice, filter))
173                }
174            }
175        }))
176    }
177
178    /// Claims split `split`'s port for `scan`'s segment when the segment is shared, building
179    /// the pipeline that decodes it for the first claimant with one port per reader in every
180    /// split not yet finished, or `None` when this split's reader is its only one.
181    fn claim(&mut self, scan: &SegmentScanPlan, split: usize) -> Option<PortId> {
182        if split == usize::MAX {
183            return None;
184        }
185        let segment = scan.segment_id();
186        let shares = &mut self.core.shares;
187        let ports = &mut shares.ports[split];
188        if let Some(waiting) = ports.get_mut(&segment) {
189            let port = waiting.pop();
190            if waiting.is_empty() {
191                ports.remove(&segment);
192            }
193            return port;
194        }
195        let ranges = shares.shared.remove(&segment)?;
196        let mut claimants: SmallVec<[usize; 4]> = SmallVec::new();
197        for rows in &ranges {
198            claimants.extend(
199                shares
200                    .overlapping(rows)
201                    .filter(|&split| !shares.finished[split] && shares.overlaps(split, rows)),
202            );
203        }
204        if claimants.len() <= 1 {
205            return None;
206        }
207        let mut outlets: SmallVec<[PortId; 1]> = SmallVec::with_capacity(claimants.len());
208        let mut mine = None;
209        for claimant in claimants {
210            let port = self
211                .core
212                .arena
213                .create(usize::MAX, None, SourceEnd::Unclaimed);
214            outlets.push(port);
215            if claimant == split && mine.is_none() {
216                mine = Some(port);
217            } else {
218                self.core.shares.ports[claimant]
219                    .entry(segment)
220                    .or_default()
221                    .push(port);
222            }
223        }
224        self.core.add_pipeline(
225            Chain::new(ScanSource::new(scan.clone(), None, None)),
226            outlets,
227        );
228        mine
229    }
230}
231
232/// Where each segment read more than once is read from, and the ports of the segments whose
233/// decoding pipeline is built.
234#[derive(Default)]
235pub(crate) struct Shares {
236    /// The splits' row ranges, in split order.
237    splits: Vec<Range<u64>>,
238    /// Whether the splits are sorted and disjoint, so the splits overlapping a range are found
239    /// by binary search.
240    ordered: bool,
241    /// Splits that have finished, and so will claim nothing.
242    finished: Vec<bool>,
243    /// Segment readers collected before the scan starts. A split may read only a few of a
244    /// file's segments, so recording a reader must not allocate by the file-wide segment id.
245    reaches: Vec<(SegmentId, Range<u64>)>,
246    /// The row ranges of segments read more than once, until their decoding pipeline is built.
247    shared: FxHashMap<SegmentId, SmallVec<[Range<u64>; 1]>>,
248    /// By split: the unclaimed ports of each shared segment the split reads.
249    ports: Vec<FxHashMap<SegmentId, SmallVec<[PortId; 1]>>>,
250}
251
252impl Shares {
253    pub(crate) fn new(splits: Vec<Range<u64>>) -> Self {
254        let ordered = splits.windows(2).all(|pair| pair[0].end <= pair[1].start);
255        Self {
256            finished: vec![false; splits.len()],
257            ports: (0..splits.len()).map(|_| FxHashMap::default()).collect(),
258            splits,
259            ordered,
260            reaches: Vec::new(),
261            shared: FxHashMap::default(),
262        }
263    }
264
265    /// The splits overlapping `rows`.
266    fn overlapping(&self, rows: &Range<u64>) -> Range<usize> {
267        if self.ordered {
268            let first = self.splits.partition_point(|split| split.end <= rows.start);
269            let end = self.splits.partition_point(|split| split.start < rows.end);
270            return first..end.max(first);
271        }
272        let mut overlapping = self
273            .splits
274            .iter()
275            .enumerate()
276            .filter(|(_, split)| split.start < rows.end && rows.start < split.end)
277            .map(|(index, _)| index);
278        match overlapping.next() {
279            Some(first) => first..overlapping.next_back().unwrap_or(first) + 1,
280            None => 0..0,
281        }
282    }
283
284    fn overlaps(&self, split: usize, rows: &Range<u64>) -> bool {
285        let split = &self.splits[split];
286        split.start < rows.end && rows.start < split.end
287    }
288
289    /// Records the ranges `plan` reads each of its segments for, read over `rows` by every
290    /// split overlapping them. Call once per plan a split stage runs, before the scan starts,
291    /// then [`retain_shared`](Self::retain_shared).
292    pub(crate) fn add(&mut self, plan: &PlanRef, rows: Range<u64>) -> VortexResult<()> {
293        let reaches = &mut self.reaches;
294        plan.reach(rows, &Reach::Offset(0), &mut |segment, rows| {
295            reaches.push((segment, rows))
296        })
297    }
298
299    /// Keeps only the segments more than one reader reads. The rest are read by the one
300    /// pipeline that needs them.
301    pub(crate) fn retain_shared(&mut self) {
302        let mut reaches = std::mem::take(&mut self.reaches);
303        reaches.sort_unstable_by_key(|(segment, _)| *segment);
304        for readers in reaches.chunk_by(|(left, _), (right, _)| left == right) {
305            let (segment, rows) = &readers[0];
306            // One reader per split overlapping the range: shared when the range crosses a
307            // split boundary, or when multiple plan nodes read the segment.
308            if readers.len() == 1
309                && self
310                    .overlapping(rows)
311                    .filter(|&split| self.overlaps(split, rows))
312                    .nth(1)
313                    .is_none()
314            {
315                continue;
316            }
317            self.shared.insert(
318                *segment,
319                readers.iter().map(|(_, rows)| rows.clone()).collect(),
320            );
321        }
322    }
323
324    /// Every port made for a reader that has not claimed it yet.
325    #[cfg(debug_assertions)]
326    pub(crate) fn waiting(&self) -> impl Iterator<Item = PortId> + '_ {
327        self.ports
328            .iter()
329            .flat_map(|ports| ports.values().flatten().copied())
330    }
331
332    /// Marks split `split` finished and returns the ports it never claimed: those of its
333    /// readers that never came, as a chunk its mask ruled out or a stage it never ran. Each is
334    /// still an outlet of the pipeline that shares it, which must stop pushing into it: a port
335    /// has a consumer until it is closed and drained, so the caller detaches each port from
336    /// its sink before discarding it. See [`Core::finish_split`].
337    pub(crate) fn finish(&mut self, split: usize) -> SmallVec<[PortId; 4]> {
338        self.finished[split] = true;
339        std::mem::take(&mut self.ports[split])
340            .into_values()
341            .flatten()
342            .collect()
343    }
344}
345
346/// How a plan's rows map to the rows of the plan a scan's splits range over.
347#[derive(Clone, Debug)]
348pub enum Reach {
349    /// Row `r` is scanned row `r + offset`.
350    Offset(u64),
351    /// Every row is read for these scanned rows, as a dictionary's values are for its codes.
352    Fixed(Range<u64>),
353}
354
355impl Reach {
356    /// The scanned rows `rows` are read for.
357    pub fn root(&self, rows: &Range<u64>) -> Range<u64> {
358        match self {
359            Reach::Offset(offset) => rows.start + offset..rows.end + offset,
360            Reach::Fixed(root) => root.clone(),
361        }
362    }
363
364    /// The mapping of a child whose row zero is this plan's row `by`.
365    pub fn shift(&self, by: u64) -> Self {
366        match self {
367            Reach::Offset(offset) => Reach::Offset(offset + by),
368            Reach::Fixed(root) => Reach::Fixed(root.clone()),
369        }
370    }
371
372    /// The mapping of a child read whole for `rows` of this plan, as dictionary values are.
373    pub fn fixed(&self, rows: &Range<u64>) -> Self {
374        Reach::Fixed(self.root(rows))
375    }
376}
377
378/// The chunks of `concat` overlapping `rows`: each chunk's index, where it starts, and the rows
379/// of `rows` it holds, in the concatenation's domain. Found by binary search, so a split pays
380/// for the chunks it overlaps, not for the chunks of the file.
381pub(crate) fn overlapping<'a>(
382    concat: &'a ConcatPlan,
383    rows: &Range<u64>,
384) -> impl Iterator<Item = (usize, u64, Range<u64>)> + 'a {
385    let offsets = concat.row_offsets();
386    let first = offsets
387        .partition_point(|&offset| offset <= rows.start)
388        .saturating_sub(1);
389    let end = offsets.partition_point(|&offset| offset < rows.end);
390    let rows = rows.clone();
391    (first..end).filter_map(move |index| {
392        let start = offsets[index];
393        let chunk_end = offsets
394            .get(index + 1)
395            .copied()
396            .unwrap_or_else(|| concat.row_count());
397        let local = rows.start.max(start)..rows.end.min(chunk_end);
398        (local.start < local.end).then_some((index, start, local))
399    })
400}
401
402/// Unclaimed shared ports, for tests.
403#[cfg(test)]
404impl Shares {
405    pub(crate) fn pending(&self) -> usize {
406        self.ports
407            .iter()
408            .flat_map(|ports| ports.values())
409            .map(SmallVec::len)
410            .sum()
411    }
412}
vortex_layout/plan/physical/pipeline/compile.rs

vortex_layout/plan/physical/pipeline/
compile.rs
