Skip to main content

rudb_exec/
prepared.rs

1//! An expression prepared once for a pipeline and then evaluated over every chunk.
2//!
3//! `spec/engine/04-expressions.md`. [`evaluate`](crate::evaluate) walks the plan's expression tree
4//! on every chunk, which means it does four things per chunk that depend on nothing about the
5//! chunk: it recurses, it resolves every column reference by a linear search through the schema, it
6//! clones a [`LogicalType`] for every node, and it copies the whole column a [`Expr::Column`] names.
7//! Over `hits` at a hundred thousand chunks that is a hundred thousand schema searches per column
8//! reference and a hundred thousand copies of every column any expression mentions.
9//!
10//! This type does all four once. The tree is flattened into a post order array, so evaluating it is
11//! a loop over that array and the recursion is gone with it. Column references are resolved to
12//! positions when the pipeline is built. Types are held here rather than cloned out of the plan.
13//! And a column reference is not a step that produces anything: it is read straight out of the chunk
14//! at the point an operand is wanted, so the column is never copied at all.
15//!
16//! # What is shared and what is not
17//!
18//! [`Prepared`] is immutable after it is built and is `Send` and `Sync`, so one of them serves every
19//! thread running a copy of the pipeline. [`Scratch`] is the per chunk working space and there is
20//! one per pipeline instance. That split is not for this layer's benefit. It is the same split every
21//! operator needs at layer eight, where the scheduler runs one pipeline on as many threads as it has
22//! morsels for, and building it here means the operators above are written against it from the start
23//! rather than retrofitted onto it.
24//!
25//! # What is still allocated per chunk
26//!
27//! Two things, and both are named rather than hidden. A node with four or more operands gathers
28//! references to them into a `Vec<&Vector>` so a kernel can take a slice, which is one allocation of
29//! pointers rather than a copy of any data, and which a node of one, two or three operands does on
30//! the stack instead. And every kernel allocates the vector it returns, because no kernel in
31//! `rudb-kernels` takes an output parameter. The second is much the larger of the two and it is the
32//! one tier 1 fusion removes, which is scheduled after layer six for the reason
33//! `spec/engine/04-expressions.md` gives: once the tree walk is gone what is left to save is pass
34//! count, and at 1024 rows the intermediate vectors are eight kilobytes and stay in L1.
35
36use rudb_common::{Error, LogicalType, Result, Value};
37use rudb_kernels::{
38    Comparison, Connective, cast, combine, compare, is_true, refine, refine_flags, selection,
39};
40use rudb_plan::{CompareOp, ConjunctionOp, Expr, ExprRef, Plan};
41use rudb_vector::{Chunk, Selection, Vector};
42
43use crate::schema::Schema;
44
45/// The scheduler's half of the expression contract, imposed now rather than at layer eight.
46///
47/// A prepared expression is the immutable half of a pipeline and layer eight hands one of them to
48/// every thread running that pipeline. That is only sound if it holds nothing thread local, and the
49/// way to find out on the commit that breaks it rather than eight layers later is to ask the
50/// compiler here, exactly as [`Chunk`] does for the data plane.
51const _: () = {
52    const fn assert_shareable<T: Send + Sync>() {}
53    assert_shareable::<Prepared>();
54};
55
56/// One or more bound expressions, flattened and resolved against a schema.
57///
58/// Built once per pipeline with [`Prepared::new`] and evaluated per chunk with
59/// [`Prepared::evaluate`] or [`Prepared::evaluate_one`], each of which wants the [`Scratch`] that
60/// [`Prepared::scratch`] hands out.
61#[derive(Debug)]
62pub struct Prepared {
63    /// The nodes in post order, so every node's operands have already been computed when it runs.
64    steps: Vec<Step>,
65    /// The type each step produces, indexed the same way as `steps`.
66    ///
67    /// A parallel array rather than a field in the variant, for the reason [`Expr`] gives: a
68    /// [`LogicalType`] owns a `Vec` for its nested cases and putting one in every variant would make
69    /// the common variants several times larger for the benefit of the rare ones.
70    types: Vec<LogicalType>,
71    /// The operand lists of the steps that have one, as runs of step indices.
72    operands: Vec<usize>,
73    /// The last step that reads each step's slot, or `usize::MAX` for one nothing reads.
74    ///
75    /// A slot is emptied as soon as the step that was the last to read it has run. Keeping every
76    /// intermediate alive to the end of the array instead is what the first measured version of this
77    /// did, and a chain of eight additions was slower prepared than walked because of it: nine live
78    /// intermediates at eight kilobytes each is seventy two kilobytes of working set where the tree
79    /// walk has two, and two is the pair the allocator hands back and forth and that stays in L1.
80    /// Everything else about the prepared form was faster and this one thing paid all of it back.
81    last_use: Vec<usize>,
82    /// The step index each expression this was built from ends at.
83    roots: Vec<usize>,
84}
85
86/// One node of a flattened expression.
87///
88/// A step refers to its operands by their index in [`Prepared::steps`], which is always smaller than
89/// its own because the array is in post order.
90#[derive(Debug)]
91enum Step {
92    /// A column of the chunk, by resolved position.
93    ///
94    /// This step computes nothing. Its slot stays empty and an operand that names it is read out of
95    /// the chunk, which is the whole of what makes a column reference free rather than a copy.
96    Column(usize),
97    /// A literal, materialized into a constant vector as long as the chunk.
98    Constant(Value),
99    /// A cast to this step's own type.
100    Cast {
101        /// The step being cast.
102        input: usize,
103        /// Whether a failed cast yields null instead of raising.
104        try_cast: bool,
105    },
106    /// A binary comparison.
107    Compare {
108        /// Which comparison.
109        op: Comparison,
110        /// The left operand's step.
111        left: usize,
112        /// The right operand's step.
113        right: usize,
114    },
115    /// An `AND` or `OR` over a run of [`Prepared::operands`].
116    Conjunction {
117        /// Which connective.
118        op: Connective,
119        /// Where the operand list starts.
120        start: usize,
121        /// How many operands it has.
122        len: usize,
123    },
124    /// A scalar function over a run of [`Prepared::operands`].
125    Function {
126        /// The resolved function name, held here so the plan is not consulted per chunk.
127        name: String,
128        /// Where the argument list starts.
129        start: usize,
130        /// How many arguments it has.
131        len: usize,
132    },
133    /// A searched `CASE`, whose branches are prepared expressions of their own.
134    ///
135    /// Nested rather than flattened into the same array because a branch is not evaluated over the
136    /// chunk, it is evaluated over the rows no earlier arm claimed, and a step in the outer array
137    /// would have no way to say that. The selection threaded form in #57 replaces this whole
138    /// variant, and when it does the branches stop being separate arrays.
139    Case {
140        /// The `WHEN`/`THEN` pairs, in order.
141        arms: Vec<PreparedArm>,
142        /// The `ELSE`, if there is one. Absent means null.
143        otherwise: Option<Prepared>,
144    },
145}
146
147/// One `WHEN`/`THEN` pair of a prepared [`Step::Case`].
148#[derive(Debug)]
149struct PreparedArm {
150    /// The condition.
151    when: Prepared,
152    /// The result if the condition is true.
153    then: Prepared,
154}
155
156/// The per chunk working space of one [`Prepared`].
157///
158/// One per pipeline instance and never shared, which is the mutable half of the split the module
159/// documentation describes. It is handed back in rather than made inside [`Prepared::evaluate`] so
160/// that the array of slots survives from one chunk to the next instead of being allocated a hundred
161/// thousand times over a scan.
162#[derive(Debug)]
163pub struct Scratch {
164    /// What each step produced, or `None` for a step that produces nothing and for one that has not
165    /// run yet.
166    slots: Vec<Option<Vector>>,
167}
168
169impl Prepared {
170    /// Prepares `exprs` against `schema`.
171    ///
172    /// # Errors
173    ///
174    /// If a column reference names a binding the schema does not have, or if an aggregate appears
175    /// where an ordinary expression was expected. Both are failures of the plan rather than of the
176    /// data, which is why they are found here, once, rather than on some chunk in the middle of a
177    /// scan.
178    pub fn new(plan: &Plan, exprs: &[ExprRef], schema: &Schema) -> Result<Self> {
179        let mut prepared = Self {
180            steps: Vec::new(),
181            types: Vec::new(),
182            operands: Vec::new(),
183            last_use: Vec::new(),
184            roots: Vec::new(),
185        };
186        for &expr in exprs {
187            let root = prepared.push(plan, expr, schema)?;
188            prepared.roots.push(root);
189        }
190        prepared.last_use = prepared.last_uses();
191        Ok(prepared)
192    }
193
194    /// Which step is the last to read each step, computed once when the expression is prepared.
195    ///
196    /// A root is never freed, because the whole point of running the array was to produce it. A
197    /// step nothing reads and that is not a root cannot happen, since every step is pushed by the
198    /// node that wanted it, but saying `usize::MAX` rather than asserting that keeps this a fact
199    /// about the array rather than a claim about the builder.
200    fn last_uses(&self) -> Vec<usize> {
201        let mut last = vec![usize::MAX; self.steps.len()];
202        for index in 0..self.steps.len() {
203            self.for_each_operand(index, |operand| last[operand] = index);
204        }
205        for &root in &self.roots {
206            last[root] = usize::MAX;
207        }
208        last
209    }
210
211    /// Visits the steps one step reads, whatever shape its operands are held in.
212    fn for_each_operand(&self, index: usize, mut visit: impl FnMut(usize)) {
213        match &self.steps[index] {
214            // A case's branches are arrays of their own and read nothing out of this one.
215            Step::Column(_) | Step::Constant(_) | Step::Case { .. } => {}
216            Step::Cast { input, .. } => visit(*input),
217            Step::Compare { left, right, .. } => {
218                visit(*left);
219                visit(*right);
220            }
221            Step::Conjunction { start, len, .. } | Step::Function { start, len, .. } => {
222                for &operand in &self.operands[*start..*start + *len] {
223                    visit(operand);
224                }
225            }
226        }
227    }
228
229    /// Prepares one expression, which is the common case and saves the caller a slice.
230    ///
231    /// # Errors
232    ///
233    /// Whatever [`Prepared::new`] reports.
234    pub fn one(plan: &Plan, expr: ExprRef, schema: &Schema) -> Result<Self> {
235        Self::new(plan, &[expr], schema)
236    }
237
238    /// Working space sized for this expression.
239    #[must_use]
240    pub fn scratch(&self) -> Scratch {
241        Scratch { slots: (0..self.steps.len()).map(|_| None).collect() }
242    }
243
244    /// How many expressions this was built from.
245    #[must_use]
246    pub fn len(&self) -> usize {
247        self.roots.len()
248    }
249
250    /// Whether it was built from no expressions at all.
251    #[must_use]
252    pub fn is_empty(&self) -> bool {
253        self.roots.is_empty()
254    }
255
256    /// Evaluates every expression over `chunk`, appending one vector each to `out`.
257    ///
258    /// Appends rather than returns a `Vec`, so a caller in a loop reuses one buffer.
259    ///
260    /// # Errors
261    ///
262    /// Anything a kernel reports, on the first expression that reports it.
263    pub fn evaluate(
264        &self,
265        chunk: &Chunk,
266        scratch: &mut Scratch,
267        out: &mut Vec<Vector>,
268    ) -> Result<()> {
269        self.run(chunk, scratch)?;
270        for &root in &self.roots {
271            // The one place a column is copied, and it is copied because the caller is taking
272            // ownership of a vector that has to outlive the chunk it came from. `SELECT a` is that
273            // shape and a projection of a bare column is the only expression where it happens.
274            match self.steps[root] {
275                Step::Column(position) => out.push(chunk.column(position)?.clone()),
276                _ => out.push(scratch.slots[root].take().ok_or_else(|| missing(root))?),
277            }
278        }
279        Ok(())
280    }
281
282    /// Evaluates a single expression over `chunk`, handing back a reference to the answer.
283    ///
284    /// A reference rather than a vector, because the caller of this is a filter, which reads the
285    /// flags to build a selection and then drops them. Nothing about that wants ownership, and a
286    /// predicate that is a bare column reference, which `WHERE flag` is, would otherwise copy the
287    /// column to hand it over.
288    ///
289    /// # Errors
290    ///
291    /// Anything a kernel reports, and an internal error if this was not built from exactly one
292    /// expression.
293    pub fn evaluate_one<'s>(
294        &'s self,
295        chunk: &'s Chunk,
296        scratch: &'s mut Scratch,
297    ) -> Result<&'s Vector> {
298        let [root] = self.roots[..] else {
299            return Err(Error::internal(format!(
300                "evaluate_one over a prepared expression of {} roots",
301                self.roots.len()
302            )));
303        };
304        self.run(chunk, scratch)?;
305        self.operand(root, chunk, &scratch.slots)
306    }
307
308    /// Evaluates a single expression as a filter, handing back the rows it keeps.
309    ///
310    /// The difference between this and [`evaluate_one`](Self::evaluate_one) followed by
311    /// [`selection`] is the whole of what a threaded filter is. An `AND` evaluated as an expression
312    /// runs every conjunct over every row and then combines the flag vectors, so a predicate of four
313    /// conjuncts that each pass a fifth of the rows does five times the work of one that stops
314    /// looking at a row as soon as a conjunct rejects it. TPC-H Q6 is exactly that predicate.
315    ///
316    /// So the conjuncts of a top level `AND` are run one at a time, each over the rows the ones
317    /// before it left, and the moment nothing is left the rest of the predicate is not run at all.
318    /// The order is the order the plan gives, which is the optimizer's business rather than this
319    /// one's until the adaptive reordering of #57 lands.
320    ///
321    /// What is threaded is the conjunct's own comparison rather than the whole of its subtree. A
322    /// conjunct of `a + b > 5` still adds over the whole chunk, because the scalar kernels take a
323    /// vector rather than a selection, and it is the comparison and everything downstream of it that
324    /// reads only the rows still in play. A conjunct that is a bare column, a function or a nested
325    /// `OR` produces flags over the chunk and is intersected with [`refine_flags`], which is what
326    /// keeps one awkward conjunct from putting the others back on the unthreaded path.
327    ///
328    /// # Errors
329    ///
330    /// Anything a kernel reports, and an internal error if this was not built from exactly one
331    /// expression.
332    pub fn evaluate_filter(&self, chunk: &Chunk, scratch: &mut Scratch) -> Result<Selection> {
333        let [root] = self.roots[..] else {
334            return Err(Error::internal(format!(
335                "evaluate_filter over a prepared expression of {} roots",
336                self.roots.len()
337            )));
338        };
339        let Step::Conjunction { op: Connective::And, start, len } = self.steps[root] else {
340            let flags = self.evaluate_one(chunk, scratch)?;
341            return Ok(selection(flags, chunk.len()));
342        };
343
344        scratch.slots.clear();
345        scratch.slots.resize_with(self.steps.len(), || None);
346        // The array is in post order and this expression's steps are the whole of it, so the subtree
347        // of the first conjunct starts at zero and the subtree of every other one starts just after
348        // the conjunct before it ends. That is what makes running a conjunct at a time a matter of
349        // walking the same array in the same order rather than of holding a second structure.
350        let mut begin = 0;
351        let mut kept: Option<Selection> = None;
352        for at in 0..len {
353            let conjunct = self.operands[start + at];
354            if kept.as_ref().is_some_and(Selection::is_empty) {
355                break;
356            }
357            for index in begin..conjunct {
358                self.run_step(index, chunk, scratch)?;
359            }
360            let next = self.thread(conjunct, chunk, scratch, kept.as_ref())?;
361            kept = Some(next);
362            // A conjunct's subtree is its own, because nothing here looks for a common subexpression
363            // and so no step outside the range is reading one inside it.
364            for index in begin..=conjunct {
365                scratch.slots[index] = None;
366            }
367            begin = conjunct + 1;
368        }
369        Ok(kept.unwrap_or_else(|| Selection::identity(chunk.len())))
370    }
371
372    /// One conjunct, over the rows the conjuncts before it left, or over all of them for the first.
373    fn thread(
374        &self,
375        index: usize,
376        chunk: &Chunk,
377        scratch: &mut Scratch,
378        kept: Option<&Selection>,
379    ) -> Result<Selection> {
380        if let Step::Compare { op, left, right } = self.steps[index] {
381            let left = self.operand(left, chunk, &scratch.slots)?;
382            let right = self.operand(right, chunk, &scratch.slots)?;
383            return match kept {
384                // The first conjunct has every row in play, and asking the threaded kernel for that
385                // would be a pass over an identity selection the unthreaded one does not need.
386                None => Ok(selection(&compare(op, left, right)?, chunk.len())),
387                Some(kept) => refine(op, left, right, kept),
388            };
389        }
390        self.run_step(index, chunk, scratch)?;
391        let flags = self.operand(index, chunk, &scratch.slots)?;
392        match kept {
393            None => Ok(selection(flags, chunk.len())),
394            Some(kept) => refine_flags(flags, kept),
395        }
396    }
397
398    /// Runs every step in order, filling the slots.
399    fn run(&self, chunk: &Chunk, scratch: &mut Scratch) -> Result<()> {
400        scratch.slots.clear();
401        scratch.slots.resize_with(self.steps.len(), || None);
402        for index in 0..self.steps.len() {
403            self.run_step(index, chunk, scratch)?;
404        }
405        Ok(())
406    }
407
408    /// Runs one step and empties the slot of every operand this was the last step to read.
409    fn run_step(&self, index: usize, chunk: &Chunk, scratch: &mut Scratch) -> Result<()> {
410        let produced = self.step(index, chunk, &scratch.slots)?;
411        scratch.slots[index] = produced;
412        let slots = &mut scratch.slots;
413        self.for_each_operand(index, |operand| {
414            if self.last_use[operand] == index {
415                slots[operand] = None;
416            }
417        });
418        Ok(())
419    }
420
421    /// Runs one step, given what the steps before it produced.
422    fn step(
423        &self,
424        index: usize,
425        chunk: &Chunk,
426        slots: &[Option<Vector>],
427    ) -> Result<Option<Vector>> {
428        let ty = &self.types[index];
429        let produced = match &self.steps[index] {
430            Step::Column(_) => None,
431            Step::Constant(value) => Some(Vector::constant(ty.clone(), value.clone(), chunk.len())),
432            Step::Cast { input, try_cast } => {
433                Some(cast(self.operand(*input, chunk, slots)?, ty, *try_cast)?)
434            }
435            Step::Compare { op, left, right } => Some(compare(
436                *op,
437                self.operand(*left, chunk, slots)?,
438                self.operand(*right, chunk, slots)?,
439            )?),
440            Step::Conjunction { op, start, len } => {
441                Some(
442                    self.with_operands(*start, *len, chunk, slots, |children| {
443                        combine(*op, children)
444                    })?,
445                )
446            }
447            Step::Function { name, start, len } => {
448                Some(self.with_operands(*start, *len, chunk, slots, |args| {
449                    rudb_kernels::call(name, args, ty)
450                })?)
451            }
452            Step::Case { arms, otherwise } => {
453                Some(self.case(chunk, arms, otherwise.as_ref(), ty)?)
454            }
455        };
456        Ok(produced)
457    }
458
459    /// The vector a step produced, or the chunk's column if the step is a column reference.
460    fn operand<'v>(
461        &self,
462        index: usize,
463        chunk: &'v Chunk,
464        slots: &'v [Option<Vector>],
465    ) -> Result<&'v Vector> {
466        if let Step::Column(position) = self.steps[index] {
467            return chunk.column(position);
468        }
469        slots[index].as_ref().ok_or_else(|| missing(index))
470    }
471
472    /// Hands a kernel the references to an operand list, without allocating for the usual widths.
473    ///
474    /// One, two and three because those are what a bound tree is made of: every scalar function in
475    /// the catalog is unary or binary, a comparison is binary, and a conjunction is two or three
476    /// often enough to be worth a line. A stack array for those means a chain of eight additions
477    /// makes zero allocations for its operand lists over a chunk instead of eight, and eight
478    /// allocations a chunk at the rate a pipeline produces chunks is a real number rather than a
479    /// tidiness argument. Anything wider falls back to [`gather`](Self::gather), which is a `Vec`
480    /// of pointers and still moves no data.
481    fn with_operands<'v, T>(
482        &self,
483        start: usize,
484        len: usize,
485        chunk: &'v Chunk,
486        slots: &'v [Option<Vector>],
487        run: impl FnOnce(&[&'v Vector]) -> Result<T>,
488    ) -> Result<T> {
489        match self.operands[start..start + len] {
490            [a] => run(&[self.operand(a, chunk, slots)?]),
491            [a, b] => run(&[self.operand(a, chunk, slots)?, self.operand(b, chunk, slots)?]),
492            [a, b, c] => run(&[
493                self.operand(a, chunk, slots)?,
494                self.operand(b, chunk, slots)?,
495                self.operand(c, chunk, slots)?,
496            ]),
497            _ => {
498                let gathered = self.gather(start, len, chunk, slots)?;
499                run(&gathered)
500            }
501        }
502    }
503
504    /// References to an operand list, for a kernel that takes a slice of them.
505    ///
506    /// The `Vec` here is the allocation the module documentation names: it holds pointers rather
507    /// than vectors, so it is a dozen bytes an operand and no data moves.
508    fn gather<'v>(
509        &self,
510        start: usize,
511        len: usize,
512        chunk: &'v Chunk,
513        slots: &'v [Option<Vector>],
514    ) -> Result<Vec<&'v Vector>> {
515        let mut gathered = Vec::with_capacity(len);
516        for &operand in &self.operands[start..start + len] {
517            gathered.push(self.operand(operand, chunk, slots)?);
518        }
519        Ok(gathered)
520    }
521
522    /// A searched `CASE` over the rows no earlier arm claimed.
523    ///
524    /// The same shape [`evaluate`](crate::evaluate) has, because the thing that makes it that shape
525    /// is a correctness rule rather than a performance one: `CASE WHEN x <> 0 THEN 1 / x ELSE 0 END`
526    /// divides by zero on the rows the arm excludes if the arm is evaluated for them. What is left
527    /// of it after #57 is the same rule expressed as a selection rather than as a narrowed chunk,
528    /// with the answers scattered back instead of assembled out of a `Vec<Value>`.
529    fn case(
530        &self,
531        chunk: &Chunk,
532        arms: &[PreparedArm],
533        otherwise: Option<&Prepared>,
534        ty: &LogicalType,
535    ) -> Result<Vector> {
536        let mut answers = vec![Value::Null; chunk.len()];
537        let mut pending: Vec<usize> = (0..chunk.len()).collect();
538        for arm in arms {
539            if pending.is_empty() {
540                break;
541            }
542            let narrowed = narrow(chunk, &pending)?;
543            let mut scratch = arm.when.scratch();
544            let flags = arm.when.evaluate_one(&narrowed, &mut scratch)?;
545            let mut taken = Vec::new();
546            let mut still = Vec::new();
547            // row at a time: the scatter that replaces these three loops is #57, and this variant
548            // goes with it.
549            for (at, &row) in pending.iter().enumerate() {
550                if is_true(&flags.value_at(at)) {
551                    taken.push((at, row));
552                } else {
553                    still.push(row);
554                }
555            }
556            if !taken.is_empty() {
557                let positions: Vec<usize> = taken.iter().map(|&(at, _)| at).collect();
558                let matched = narrow(&narrowed, &positions)?;
559                let mut scratch = arm.then.scratch();
560                let results = arm.then.evaluate_one(&matched, &mut scratch)?;
561                // row at a time: the scatter this wants is #57, same as the loop above.
562                for (slot, &(_, row)) in taken.iter().enumerate() {
563                    answers[row] = results.value_at(slot);
564                }
565            }
566            pending = still;
567        }
568        if let Some(otherwise) = otherwise {
569            if !pending.is_empty() {
570                let narrowed = narrow(chunk, &pending)?;
571                let mut scratch = otherwise.scratch();
572                let results = otherwise.evaluate_one(&narrowed, &mut scratch)?;
573                // row at a time: the scatter this wants is #57, same as the two above.
574                for (slot, &row) in pending.iter().enumerate() {
575                    answers[row] = results.value_at(slot);
576                }
577            }
578        }
579        Vector::from_values(ty.clone(), &answers)
580    }
581
582    /// Flattens one expression, appending its steps and returning the index of its last one.
583    fn push(&mut self, plan: &Plan, expr: ExprRef, schema: &Schema) -> Result<usize> {
584        let ty = plan.expr_type(expr).clone();
585        let step = match *plan.expr(expr) {
586            Expr::Column(binding) => {
587                let position = schema.position_of(binding).ok_or_else(|| {
588                    Error::internal(format!(
589                        "column #{}.{} is not in the schema this operator was given",
590                        binding.table, binding.column
591                    ))
592                })?;
593                Step::Column(position)
594            }
595            Expr::Constant(reference) => Step::Constant(plan.value(reference).clone()),
596            Expr::Cast { input, try_cast } => {
597                Step::Cast { input: self.push(plan, input, schema)?, try_cast }
598            }
599            Expr::Compare { op, left, right } => Step::Compare {
600                op: comparison(op),
601                left: self.push(plan, left, schema)?,
602                right: self.push(plan, right, schema)?,
603            },
604            Expr::Conjunction { op, children } => {
605                let (start, len) = self.push_list(plan, plan.expr_list(children), schema)?;
606                Step::Conjunction { op: connective(op), start, len }
607            }
608            Expr::Function { name, args } => {
609                let (start, len) = self.push_list(plan, plan.expr_list(args), schema)?;
610                Step::Function { name: plan.string(name).to_string(), start, len }
611            }
612            Expr::Aggregate { name, .. } => {
613                return Err(Error::internal(format!(
614                    "the {} aggregate was evaluated as an ordinary expression",
615                    plan.string(name)
616                )));
617            }
618            Expr::Case { arms, otherwise } => {
619                let mut prepared = Vec::new();
620                for &arm in plan.arm_list(arms) {
621                    prepared.push(PreparedArm {
622                        when: Self::one(plan, arm.when, schema)?,
623                        then: Self::one(plan, arm.then, schema)?,
624                    });
625                }
626                let otherwise = match otherwise {
627                    Some(otherwise) => Some(Self::one(plan, otherwise, schema)?),
628                    None => None,
629                };
630                Step::Case { arms: prepared, otherwise }
631            }
632        };
633        self.steps.push(step);
634        self.types.push(ty);
635        Ok(self.steps.len() - 1)
636    }
637
638    /// Flattens a list of expressions and records where its operand run starts and how long it is.
639    ///
640    /// The operand run is written after every child has been flattened rather than as they go,
641    /// because a child that is itself a list would otherwise interleave its run with this one.
642    fn push_list(
643        &mut self,
644        plan: &Plan,
645        exprs: &[ExprRef],
646        schema: &Schema,
647    ) -> Result<(usize, usize)> {
648        let mut indices = Vec::with_capacity(exprs.len());
649        for &expr in exprs {
650            indices.push(self.push(plan, expr, schema)?);
651        }
652        let start = self.operands.len();
653        let len = indices.len();
654        self.operands.extend(indices);
655        Ok((start, len))
656    }
657}
658
659/// The error for a slot that should have held something and did not.
660///
661/// This cannot happen while the array is in post order, since every operand's index is smaller than
662/// the index of the step using it and every step runs in order. It is an error rather than a panic
663/// because the property it depends on is a property of [`Prepared::push`], and the day somebody
664/// writes a pass that reorders the array is the day it stops holding.
665fn missing(index: usize) -> Error {
666    Error::internal(format!("step {index} was used as an operand before it produced anything"))
667}
668
669/// The chunk cut down to the given rows.
670///
671/// The reason `CASE` is written with this rather than by evaluating every arm over the whole chunk
672/// and picking afterwards. `CASE WHEN x <> 0 THEN 1 / x ELSE 0 END` divides by zero on the rows the
673/// arm does not apply to if the arm is evaluated for them, and a `CASE` that raises on a row it was
674/// written to exclude is the classic wrong answer this shape prevents.
675pub(crate) fn narrow(chunk: &Chunk, rows: &[usize]) -> Result<Chunk> {
676    let mut selection = Selection::with_capacity(rows.len());
677    for &row in rows {
678        selection.push(row);
679    }
680    chunk.clone().select(&selection)
681}
682
683/// The kernels' comparison for the plan's.
684///
685/// A translation rather than one shared enum, because the kernels are rank 3 and the plan is rank
686/// 9. This function is the whole of what that separation costs.
687pub(crate) fn comparison(op: CompareOp) -> Comparison {
688    match op {
689        CompareOp::Equal => Comparison::Equal,
690        CompareOp::NotEqual => Comparison::NotEqual,
691        CompareOp::Less => Comparison::Less,
692        CompareOp::LessOrEqual => Comparison::LessOrEqual,
693        CompareOp::Greater => Comparison::Greater,
694        CompareOp::GreaterOrEqual => Comparison::GreaterOrEqual,
695        CompareOp::DistinctFrom => Comparison::DistinctFrom,
696        CompareOp::NotDistinctFrom => Comparison::NotDistinctFrom,
697    }
698}
699
700/// The kernels' connective for the plan's.
701pub(crate) fn connective(op: ConjunctionOp) -> Connective {
702    match op {
703        ConjunctionOp::And => Connective::And,
704        ConjunctionOp::Or => Connective::Or,
705    }
706}
707
708#[cfg(test)]
709mod tests {
710    use rudb_common::{Field, LogicalType, Value};
711    use rudb_kernels::is_true;
712    use rudb_plan::{ExprRef, Node, Plan};
713    use rudb_vector::{Chunk, Selection, Vector};
714
715    use super::{Prepared, narrow};
716    use crate::expr::evaluate;
717    use crate::schema::Schema;
718
719    /// Two columns with a null in each, because every disagreement between these two evaluators
720    /// that is worth finding is a disagreement about which rows are null.
721    fn input() -> (Schema, Chunk) {
722        let schema = Schema::numbered(
723            vec![Field::new("x", LogicalType::Integer), Field::new("s", LogicalType::Varchar)],
724            0,
725        );
726        let x = Vector::from_values(
727            LogicalType::Integer,
728            &[Value::Integer(3), Value::Integer(1), Value::Null, Value::Integer(2)],
729        )
730        .expect("four integers");
731        let s = Vector::from_values(
732            LogicalType::Varchar,
733            &[
734                Value::Varchar("a".to_string()),
735                Value::Null,
736                Value::Varchar("c".to_string()),
737                Value::Varchar("a".to_string()),
738            ],
739        )
740        .expect("four strings");
741        (schema, Chunk::new(vec![x, s]).expect("two columns of four rows"))
742    }
743
744    /// The expressions of a projection written in the plan's textual form, over the two columns
745    /// [`input`] produces.
746    ///
747    /// Going through the text rather than the arena builders for the reason the other test module
748    /// gives: a test that says what it evaluates in the notation a plan dump uses is a test whose
749    /// failure can be pasted into a plan and vice versa.
750    fn projection(exprs: &str) -> (Plan, Vec<ExprRef>) {
751        let text =
752            format!("Project #1 [{exprs}]\n  Get memory.main.t AS t #0 [x::INTEGER, s::VARCHAR]");
753        let plan = Plan::parse(&text).expect("a well formed plan");
754        let Node::Project { exprs, .. } = *plan.node(plan.root()) else {
755            panic!("the root of that text is a projection");
756        };
757        let list = plan.expr_list(exprs).to_vec();
758        (plan, list)
759    }
760
761    /// Every expression shape, evaluated both ways over the same chunk.
762    ///
763    /// This is the agreement the module documentation claims and it is the only thing that makes
764    /// the prepared form safe to put in front of the tree walk. The generated well typed trees the
765    /// test gate of #57 asks for are a wider version of this and are worth building once the
766    /// selection threaded shapes exist to disagree about.
767    fn agrees(exprs: &str) {
768        let (schema, chunk) = input();
769        let (plan, list) = projection(exprs);
770        let prepared = Prepared::new(&plan, &list, &schema).expect("the expressions resolve");
771        let mut scratch = prepared.scratch();
772        let mut fast = Vec::new();
773        prepared.evaluate(&chunk, &mut scratch, &mut fast).expect("the prepared form runs");
774        for (at, &expr) in list.iter().enumerate() {
775            let slow = evaluate(&plan, expr, &schema, &chunk).expect("the tree walk runs");
776            for row in 0..chunk.len() {
777                assert_eq!(
778                    fast[at].value_at(row),
779                    slow.value_at(row),
780                    "expression {at} of `{exprs}` at row {row}"
781                );
782            }
783        }
784    }
785
786    #[test]
787    fn a_column_reference_agrees() {
788        agrees("#0.0::INTEGER AS a, #0.1::VARCHAR AS b");
789    }
790
791    #[test]
792    fn a_constant_agrees() {
793        agrees("7::INTEGER AS a, NULL::INTEGER AS b");
794    }
795
796    #[test]
797    fn a_cast_agrees() {
798        agrees("CAST(#0.0::INTEGER)::BIGINT AS a, CAST(#0.0::INTEGER)::VARCHAR AS b");
799    }
800
801    #[test]
802    fn a_comparison_agrees() {
803        agrees("(#0.0::INTEGER > 1::INTEGER)::BOOLEAN AS a");
804    }
805
806    #[test]
807    fn a_conjunction_agrees() {
808        agrees(
809            "((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND (#0.0::INTEGER < 3::INTEGER)::BOOLEAN)\
810             ::BOOLEAN AS a",
811        );
812    }
813
814    #[test]
815    fn a_function_agrees() {
816        agrees("\"+\"(#0.0::INTEGER, 1::INTEGER)::INTEGER AS a");
817    }
818
819    #[test]
820    fn a_case_agrees() {
821        agrees(
822            "CASE WHEN (#0.0::INTEGER > 1::INTEGER)::BOOLEAN THEN 10::INTEGER \
823             ELSE 20::INTEGER END::INTEGER AS a",
824        );
825    }
826
827    /// The same expression twice, which is where the tree walk copies the column twice and this
828    /// does not, and the answers still have to be identical.
829    #[test]
830    fn a_column_mentioned_three_times_agrees() {
831        agrees("\"+\"(\"+\"(#0.0::INTEGER, #0.0::INTEGER)::INTEGER, #0.0::INTEGER)::INTEGER AS a");
832    }
833
834    /// The intermediates of a chain are not all held to the end of it.
835    ///
836    /// This is the whole difference between the prepared form being faster than the tree walk on a
837    /// deep chain and being slower than it, and it is a property of the slot array rather than of
838    /// any answer, so it is asserted here rather than left to the benchmark to catch.
839    #[test]
840    fn a_chain_holds_one_intermediate_at_a_time() {
841        let (schema, chunk) = input();
842        let mut expr = "#0.0::INTEGER".to_string();
843        for _ in 0..8 {
844            expr = format!("\"+\"({expr}, 1::INTEGER)::INTEGER");
845        }
846        let (plan, list) = projection(&format!("{expr} AS a"));
847        let prepared = Prepared::new(&plan, &list, &schema).expect("the chain resolves");
848        let mut scratch = prepared.scratch();
849        prepared.run(&chunk, &mut scratch).expect("the chain runs");
850        let live = scratch.slots.iter().filter(|slot| slot.is_some()).count();
851        assert_eq!(live, 1, "a chain that has run should be holding its answer and nothing else");
852    }
853
854    /// The rows a threaded filter keeps are the rows the tree walk says the predicate is true for.
855    ///
856    /// Every threaded conjunct is a chance to disagree with the unthreaded answer about a null,
857    /// about a row an earlier conjunct had already dropped, or about a chunk nothing survives, and
858    /// the answer is a set of row numbers rather than a vector, so this is checked against the tree
859    /// walk read a row at a time rather than against the prepared form it is part of.
860    fn filters(predicate: &str) {
861        let (schema, chunk) = input();
862        let (plan, list) = projection(&format!("{predicate} AS p"));
863        let prepared = Prepared::new(&plan, &list, &schema).expect("the predicate resolves");
864        let mut scratch = prepared.scratch();
865        let threaded = prepared.evaluate_filter(&chunk, &mut scratch).expect("the filter runs");
866        let flags = evaluate(&plan, list[0], &schema, &chunk).expect("the tree walk runs");
867        let expected = Selection::from_predicate(chunk.len(), |row| is_true(&flags.value_at(row)));
868        assert_eq!(threaded, expected, "`{predicate}`");
869        // And running it again over the same scratch is the same answer, because a pipeline calls
870        // this once a chunk and a slot left behind by the conjunct before would show up here.
871        let again = prepared.evaluate_filter(&chunk, &mut scratch).expect("the filter runs again");
872        assert_eq!(again, expected, "`{predicate}` a second time");
873    }
874
875    /// A predicate with no `AND` in it is not threaded and has to keep saying the same thing.
876    #[test]
877    fn a_single_comparison_filters_the_same_rows() {
878        filters("(#0.0::INTEGER > 1::INTEGER)::BOOLEAN");
879        filters("(#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN");
880        filters("(#0.0::INTEGER IS NOT DISTINCT FROM NULL::INTEGER)::BOOLEAN");
881    }
882
883    #[test]
884    fn a_chain_of_conjuncts_keeps_what_all_of_them_keep() {
885        filters(
886            "((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND (#0.0::INTEGER < 3::INTEGER)::BOOLEAN)\
887             ::BOOLEAN",
888        );
889        filters(
890            "((#0.0::INTEGER >= 1::INTEGER)::BOOLEAN AND (#0.0::INTEGER <= 3::INTEGER)::BOOLEAN \
891             AND (#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN AND (#0.0::INTEGER <> 2::INTEGER)\
892             ::BOOLEAN)::BOOLEAN",
893        );
894    }
895
896    /// A conjunct that rejects every row, in front of one that would have kept some. The rows are
897    /// the same either way and the point of the shape is that the second conjunct never runs.
898    #[test]
899    fn a_conjunct_that_keeps_nothing_ends_the_predicate() {
900        filters(
901            "((#0.0::INTEGER > 9::INTEGER)::BOOLEAN AND (#0.0::INTEGER < 9::INTEGER)::BOOLEAN)\
902             ::BOOLEAN",
903        );
904    }
905
906    /// A conjunct whose operands are computed rather than read, which is the shape where the
907    /// comparison is threaded and the arithmetic under it is not.
908    #[test]
909    fn a_conjunct_over_a_computed_operand_keeps_the_same_rows() {
910        filters(
911            "((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND \
912             (\"+\"(#0.0::INTEGER, 1::INTEGER)::INTEGER < 4::INTEGER)::BOOLEAN)::BOOLEAN",
913        );
914    }
915
916    /// A conjunct that is not a comparison at all, which is the one that goes through the flag
917    /// kernel rather than the comparison kernel.
918    #[test]
919    fn a_conjunct_that_is_not_a_comparison_is_threaded_too() {
920        filters(
921            "((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND ((#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN \
922             OR (#0.0::INTEGER = 1::INTEGER)::BOOLEAN)::BOOLEAN)::BOOLEAN",
923        );
924        filters(
925            "(((#0.1::VARCHAR = 'c'::VARCHAR)::BOOLEAN OR (#0.0::INTEGER = 3::INTEGER)::BOOLEAN)\
926             ::BOOLEAN AND (#0.0::INTEGER <> 1::INTEGER)::BOOLEAN)::BOOLEAN",
927        );
928    }
929
930    /// An `OR` at the top is not threaded, because a row the left side rejects is a row the right
931    /// side may still keep. Threading it would be the wrong answer rather than a slower one.
932    #[test]
933    fn an_or_at_the_top_is_not_threaded() {
934        filters(
935            "((#0.0::INTEGER > 2::INTEGER)::BOOLEAN OR (#0.1::VARCHAR = 'c'::VARCHAR)::BOOLEAN)\
936             ::BOOLEAN",
937        );
938    }
939
940    /// A filter over a chunk that has already been narrowed, which is what a second filter in a
941    /// pipeline sees and is the form pair the threaded kernels have to handle rather than fall
942    /// through on.
943    #[test]
944    fn a_filter_over_a_selected_chunk_keeps_the_same_rows() {
945        let (schema, chunk) = input();
946        let predicate = "((#0.0::INTEGER >= 1::INTEGER)::BOOLEAN AND \
947                         (#0.1::VARCHAR = 'a'::VARCHAR)::BOOLEAN)::BOOLEAN";
948        let (plan, list) = projection(&format!("{predicate} AS p"));
949        let prepared = Prepared::new(&plan, &list, &schema).expect("the predicate resolves");
950        let mut scratch = prepared.scratch();
951        let narrowed = narrow(&chunk, &[0, 3]).expect("two of the four rows");
952        let threaded = prepared.evaluate_filter(&narrowed, &mut scratch).expect("the filter runs");
953        let flags = evaluate(&plan, list[0], &schema, &narrowed).expect("the tree walk runs");
954        let expected =
955            Selection::from_predicate(narrowed.len(), |row| is_true(&flags.value_at(row)));
956        assert_eq!(threaded, expected);
957    }
958
959    /// Preparing is per pipeline and evaluating is per chunk, so the scratch has to survive being
960    /// used again and give the same answer the second time.
961    #[test]
962    fn a_scratch_used_twice_gives_the_same_answer_twice() {
963        let (schema, chunk) = input();
964        let (plan, list) = projection("\"+\"(#0.0::INTEGER, 1::INTEGER)::INTEGER AS a");
965        let prepared = Prepared::new(&plan, &list, &schema).expect("the expressions resolve");
966        let mut scratch = prepared.scratch();
967        let mut once = Vec::new();
968        prepared.evaluate(&chunk, &mut scratch, &mut once).expect("the first chunk runs");
969        let mut twice = Vec::new();
970        prepared.evaluate(&chunk, &mut scratch, &mut twice).expect("the second chunk runs");
971        assert_eq!(once, twice);
972    }
973
974    /// A chunk shorter than the last one, because a scan's final chunk is that and a constant
975    /// materialized to the wrong length would be an out of range read rather than a wrong answer.
976    #[test]
977    fn a_shorter_chunk_after_a_longer_one_is_evaluated_at_its_own_length() {
978        let (schema, chunk) = input();
979        let (plan, list) = projection("7::INTEGER AS a");
980        let prepared = Prepared::new(&plan, &list, &schema).expect("the expressions resolve");
981        let mut scratch = prepared.scratch();
982        let mut full = Vec::new();
983        prepared.evaluate(&chunk, &mut scratch, &mut full).expect("the full chunk runs");
984        assert_eq!(full[0].len(), 4);
985        let short = chunk
986            .clone()
987            .select(&{
988                let mut selection = Selection::with_capacity(2);
989                selection.push(0);
990                selection.push(2);
991                selection
992            })
993            .expect("two of the four rows");
994        let mut cut = Vec::new();
995        prepared.evaluate(&short, &mut scratch, &mut cut).expect("the short chunk runs");
996        assert_eq!(cut[0].len(), 2);
997    }
998
999    /// An aggregate is not an expression and saying so when the pipeline is built is better than
1000    /// saying it on the first chunk.
1001    #[test]
1002    fn an_aggregate_is_refused_when_it_is_prepared() {
1003        let (schema, _) = input();
1004        let text = "Aggregate #1 groups=[] aggregates=[sum(#0.0::INTEGER)::HUGEINT]\n  \
1005                    Get memory.main.t AS t #0 [x::INTEGER, s::VARCHAR]";
1006        let plan = Plan::parse(text).expect("a well formed plan");
1007        let Node::Aggregate { aggregates, .. } = *plan.node(plan.root()) else {
1008            panic!("the root of that text is an aggregate");
1009        };
1010        let list = plan.expr_list(aggregates).to_vec();
1011        let error = Prepared::new(&plan, &list, &schema).expect_err("sum is not a scalar");
1012        assert!(error.message().contains("sum"), "{error}");
1013    }
1014}