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