reifydb_engine/expression/
compile.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2026 ReifyDB
3
4use std::{mem::discriminant, slice::from_ref};
5
6use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
7use reifydb_rql::expression::{Expression, name::display_label};
8use reifydb_value::{
9	error::{BinaryOp, Error, IntoDiagnostic, LogicalOp, RuntimeErrorKind, TypeError},
10	fragment::Fragment,
11	value::{Value, value_type::ValueType},
12};
13
14use super::{
15	context::CompileContext,
16	option::{binary_op_unwrap_option, unary_op_unwrap_option},
17};
18use crate::{
19	Result,
20	error::CastError,
21	expression::{
22		access::access_lookup,
23		arith::{add::add_columns, div::div_columns, mul::mul_columns, rem::rem_columns, sub::sub_columns},
24		call::call_builtin,
25		cast::cast_column_data,
26		compare::{Equal, GreaterThan, GreaterThanEqual, LessThan, LessThanEqual, NotEqual, compare_columns},
27		constant::{constant_value, constant_value_of},
28		context::EvalContext,
29		logic::{execute_logical_op, try_short_circuit_and, try_short_circuit_or},
30		lookup::column_lookup,
31		parameter::parameter_lookup,
32		prefix::prefix_apply,
33	},
34	vm::stack::Variable,
35};
36
37type SingleExprFn = Box<dyn Fn(&EvalContext) -> Result<ColumnWithName> + Send + Sync>;
38type MultiExprFn = Box<dyn Fn(&EvalContext) -> Result<Vec<ColumnWithName>> + Send + Sync>;
39
40pub struct CompiledExpr {
41	inner: CompiledExprInner,
42	access_column_name: Option<String>,
43}
44
45enum CompiledExprInner {
46	Single(SingleExprFn),
47	Multi(MultiExprFn),
48}
49
50impl CompiledExpr {
51	pub fn new(f: impl Fn(&EvalContext) -> Result<ColumnWithName> + Send + Sync + 'static) -> Self {
52		Self {
53			inner: CompiledExprInner::Single(Box::new(f)),
54			access_column_name: None,
55		}
56	}
57
58	pub fn new_multi(f: impl Fn(&EvalContext) -> Result<Vec<ColumnWithName>> + Send + Sync + 'static) -> Self {
59		Self {
60			inner: CompiledExprInner::Multi(Box::new(f)),
61			access_column_name: None,
62		}
63	}
64
65	pub fn new_access(
66		name: String,
67		f: impl Fn(&EvalContext) -> Result<ColumnWithName> + Send + Sync + 'static,
68	) -> Self {
69		Self {
70			inner: CompiledExprInner::Single(Box::new(f)),
71			access_column_name: Some(name),
72		}
73	}
74
75	pub fn access_column_name(&self) -> Option<&str> {
76		self.access_column_name.as_deref()
77	}
78
79	pub fn execute(&self, ctx: &EvalContext) -> Result<ColumnWithName> {
80		match &self.inner {
81			CompiledExprInner::Single(f) => f(ctx),
82			CompiledExprInner::Multi(f) => {
83				let columns = f(ctx)?;
84				Ok(columns.into_iter().next().unwrap_or_else(|| ColumnWithName {
85					name: Fragment::internal("none"),
86					data: ColumnBuffer::with_capacity(
87						ValueType::Option(Box::new(ValueType::Boolean)),
88						0,
89					),
90				}))
91			}
92		}
93	}
94
95	pub fn execute_multi(&self, ctx: &EvalContext) -> Result<Vec<ColumnWithName>> {
96		match &self.inner {
97			CompiledExprInner::Single(f) => Ok(vec![f(ctx)?]),
98			CompiledExprInner::Multi(f) => f(ctx),
99		}
100	}
101}
102
103macro_rules! compile_arith {
104	($ctx:expr, $parent:expr, $e:expr, $op_fn:path) => {{
105		let left = compile_expression($ctx, &$e.left)?;
106		let right = compile_expression($ctx, &$e.right)?;
107		let fragment = $e.full_fragment_owned();
108		let label = display_label($parent);
109		CompiledExpr::new(move |ctx| {
110			let l = left.execute(ctx)?;
111			let r = right.execute(ctx)?;
112			let mut col = $op_fn(ctx, &l, &r, || fragment.clone())?;
113			col.name = label.clone();
114			Ok(col)
115		})
116	}};
117}
118
119macro_rules! compile_compare {
120	($ctx:expr, $parent:expr, $e:expr, $cmp_type:ty, $binary_op:expr) => {{
121		let left = compile_expression($ctx, &$e.left)?;
122		let right = compile_expression($ctx, &$e.right)?;
123		let fragment = $e.full_fragment_owned();
124		let label = display_label($parent);
125		CompiledExpr::new(move |ctx| {
126			let l = left.execute(ctx)?;
127			let r = right.execute(ctx)?;
128			let mut col = compare_columns::<$cmp_type>(&l, &r, fragment.clone(), |f, l, r| {
129				TypeError::BinaryOperatorNotApplicable {
130					operator: $binary_op,
131					left: l,
132					right: r,
133					fragment: f,
134				}
135				.into_diagnostic()
136			})?;
137			col.name = label.clone();
138			Ok(col)
139		})
140	}};
141}
142
143pub fn compile_expression(_ctx: &CompileContext, expr: &Expression) -> Result<CompiledExpr> {
144	Ok(match expr {
145		Expression::Constant(e) => {
146			let constant = e.clone();
147			let label = display_label(expr);
148			CompiledExpr::new(move |ctx| {
149				let row_count = ctx.take.unwrap_or(ctx.row_count);
150				Ok(ColumnWithName {
151					name: label.clone(),
152					data: constant_value(&constant, row_count)?,
153				})
154			})
155		}
156
157		Expression::Column(e) => {
158			let expr = e.clone();
159			CompiledExpr::new(move |ctx| column_lookup(ctx, &expr))
160		}
161
162		Expression::Variable(e) => {
163			let expr = e.clone();
164			CompiledExpr::new(move |ctx| {
165				let variable_name = expr.name();
166
167				if variable_name == "env" {
168					return Err(TypeError::Runtime {
169						kind: RuntimeErrorKind::VariableIsDataframe {
170							name: variable_name.to_string(),
171						},
172						message: format!(
173							"Variable '{}' contains a dataframe and cannot be used directly in scalar expressions",
174							variable_name
175						),
176					}
177					.into());
178				}
179
180				match ctx.symbols.get(variable_name) {
181					Some(Variable::Columns {
182						columns,
183					}) if columns.is_scalar() => {
184						let value = columns.scalar_value();
185						let mut data =
186							ColumnBuffer::with_capacity(value.get_type(), ctx.row_count);
187						for _ in 0..ctx.row_count {
188							data.push_value(value.clone());
189						}
190						Ok(ColumnWithName {
191							name: Fragment::internal(variable_name),
192							data,
193						})
194					}
195					Some(Variable::Columns {
196						..
197					})
198					| Some(Variable::ForIterator {
199						..
200					})
201					| Some(Variable::Closure(_)) => Err(TypeError::Runtime {
202						kind: RuntimeErrorKind::VariableIsDataframe {
203							name: variable_name.to_string(),
204						},
205						message: format!(
206							"Variable '{}' contains a dataframe and cannot be used directly in scalar expressions",
207							variable_name
208						),
209					}
210					.into()),
211					None => {
212						if let Some(value) = ctx.params.get_named(variable_name) {
213							let mut data = ColumnBuffer::with_capacity(
214								value.get_type(),
215								ctx.row_count,
216							);
217							for _ in 0..ctx.row_count {
218								data.push_value(value.clone());
219							}
220							return Ok(ColumnWithName {
221								name: Fragment::internal(variable_name),
222								data,
223							});
224						}
225						Err(TypeError::Runtime {
226							kind: RuntimeErrorKind::VariableNotFound {
227								name: variable_name.to_string(),
228							},
229							message: format!("Variable '{}' is not defined", variable_name),
230						}
231						.into())
232					}
233				}
234			})
235		}
236
237		Expression::Parameter(e) => {
238			let expr = e.clone();
239			CompiledExpr::new(move |ctx| parameter_lookup(ctx, &expr))
240		}
241
242		Expression::Alias(e) => {
243			let inner = compile_expression(_ctx, &e.expression)?;
244			let alias = e.alias.0.clone();
245			CompiledExpr::new(move |ctx| {
246				let mut column = inner.execute(ctx)?;
247				column.name = alias.clone();
248				Ok(column)
249			})
250		}
251
252		Expression::Add(e) => compile_arith!(_ctx, expr, e, add_columns),
253		Expression::Sub(e) => compile_arith!(_ctx, expr, e, sub_columns),
254		Expression::Mul(e) => compile_arith!(_ctx, expr, e, mul_columns),
255		Expression::Div(e) => compile_arith!(_ctx, expr, e, div_columns),
256		Expression::Rem(e) => compile_arith!(_ctx, expr, e, rem_columns),
257
258		Expression::Equal(e) => compile_compare!(_ctx, expr, e, Equal, BinaryOp::Equal),
259		Expression::NotEqual(e) => compile_compare!(_ctx, expr, e, NotEqual, BinaryOp::NotEqual),
260		Expression::GreaterThan(e) => compile_compare!(_ctx, expr, e, GreaterThan, BinaryOp::GreaterThan),
261		Expression::GreaterThanEqual(e) => {
262			compile_compare!(_ctx, expr, e, GreaterThanEqual, BinaryOp::GreaterThanEqual)
263		}
264		Expression::LessThan(e) => compile_compare!(_ctx, expr, e, LessThan, BinaryOp::LessThan),
265		Expression::LessThanEqual(e) => compile_compare!(_ctx, expr, e, LessThanEqual, BinaryOp::LessThanEqual),
266
267		Expression::And(e) => {
268			let left = compile_expression(_ctx, &e.left)?;
269			let right = compile_expression(_ctx, &e.right)?;
270			let fragment = e.full_fragment_owned();
271			let label = display_label(expr);
272			CompiledExpr::new(move |ctx| {
273				let l = left.execute(ctx)?;
274				if let Some(mut short) = try_short_circuit_and(&l, &fragment, l.data().len()) {
275					short.name = label.clone();
276					return Ok(short);
277				}
278				let r = right.execute(ctx)?;
279				let mut col = execute_logical_op(&l, &r, &fragment, LogicalOp::And, |a, b| a && b)?;
280				col.name = label.clone();
281				Ok(col)
282			})
283		}
284
285		Expression::Or(e) => {
286			let left = compile_expression(_ctx, &e.left)?;
287			let right = compile_expression(_ctx, &e.right)?;
288			let fragment = e.full_fragment_owned();
289			let label = display_label(expr);
290			CompiledExpr::new(move |ctx| {
291				let l = left.execute(ctx)?;
292				if let Some(mut short) = try_short_circuit_or(&l, &fragment, l.data().len()) {
293					short.name = label.clone();
294					return Ok(short);
295				}
296				let r = right.execute(ctx)?;
297				let mut col = execute_logical_op(&l, &r, &fragment, LogicalOp::Or, |a, b| a || b)?;
298				col.name = label.clone();
299				Ok(col)
300			})
301		}
302
303		Expression::Xor(e) => {
304			let left = compile_expression(_ctx, &e.left)?;
305			let right = compile_expression(_ctx, &e.right)?;
306			let fragment = e.full_fragment_owned();
307			let label = display_label(expr);
308			CompiledExpr::new(move |ctx| {
309				let l = left.execute(ctx)?;
310				let r = right.execute(ctx)?;
311				let mut col = execute_logical_op(&l, &r, &fragment, LogicalOp::Xor, |a, b| a != b)?;
312				col.name = label.clone();
313				Ok(col)
314			})
315		}
316
317		Expression::Prefix(e) => {
318			let inner = compile_expression(_ctx, &e.expression)?;
319			let operator = e.operator.clone();
320			let fragment = e.full_fragment_owned();
321			let label = display_label(expr);
322			CompiledExpr::new(move |ctx| {
323				let column = inner.execute(ctx)?;
324				let mut col = prefix_apply(&column, &operator, &fragment)?;
325				col.name = label.clone();
326				Ok(col)
327			})
328		}
329
330		Expression::Type(e) => {
331			let ty = e.ty.clone();
332			let fragment = e.fragment.clone();
333			CompiledExpr::new(move |ctx| {
334				let row_count = ctx.take.unwrap_or(ctx.row_count);
335				let values: Vec<Box<Value>> =
336					(0..row_count).map(|_| Box::new(Value::Type(ty.clone()))).collect();
337				Ok(ColumnWithName::new(fragment.text(), ColumnBuffer::any(values)))
338			})
339		}
340
341		Expression::AccessSource(e) => {
342			let col_name = e.column.name.text().to_string();
343			let expr = e.clone();
344			CompiledExpr::new_access(col_name, move |ctx| access_lookup(ctx, &expr))
345		}
346
347		Expression::Tuple(e) => {
348			if e.expressions.len() == 1 {
349				let inner = compile_expression(_ctx, &e.expressions[0])?;
350				CompiledExpr::new(move |ctx| inner.execute(ctx))
351			} else {
352				let compiled: Vec<CompiledExpr> = e
353					.expressions
354					.iter()
355					.map(|expr| compile_expression(_ctx, expr))
356					.collect::<Result<Vec<_>>>()?;
357				let fragment = e.fragment.clone();
358				CompiledExpr::new(move |ctx| {
359					let columns: Vec<ColumnWithName> = compiled
360						.iter()
361						.map(|expr| expr.execute(ctx))
362						.collect::<Result<Vec<_>>>()?;
363
364					let len = columns.first().map_or(1, |c| c.data().len());
365					let mut data: Vec<Box<Value>> = Vec::with_capacity(len);
366
367					for i in 0..len {
368						let items: Vec<Value> =
369							columns.iter().map(|col| col.data().get_value(i)).collect();
370						data.push(Box::new(Value::Tuple(items)));
371					}
372
373					Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::any(data)))
374				})
375			}
376		}
377
378		Expression::List(e) => {
379			let compiled: Vec<CompiledExpr> = e
380				.expressions
381				.iter()
382				.map(|expr| compile_expression(_ctx, expr))
383				.collect::<Result<Vec<_>>>()?;
384			let fragment = e.fragment.clone();
385			CompiledExpr::new(move |ctx| {
386				let columns: Vec<ColumnWithName> =
387					compiled.iter().map(|expr| expr.execute(ctx)).collect::<Result<Vec<_>>>()?;
388
389				let len = columns.first().map_or(1, |c| c.data().len());
390				let mut data: Vec<Box<Value>> = Vec::with_capacity(len);
391
392				for i in 0..len {
393					let items: Vec<Value> =
394						columns.iter().map(|col| col.data().get_value(i)).collect();
395					data.push(Box::new(Value::List(items)));
396				}
397
398				Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::any(data)))
399			})
400		}
401
402		Expression::Between(e) => {
403			let value = compile_expression(_ctx, &e.value)?;
404			let lower = compile_expression(_ctx, &e.lower)?;
405			let upper = compile_expression(_ctx, &e.upper)?;
406			let fragment = e.fragment.clone();
407			CompiledExpr::new(move |ctx| {
408				let value_col = value.execute(ctx)?;
409				let lower_col = lower.execute(ctx)?;
410				let upper_col = upper.execute(ctx)?;
411
412				let ge_result = compare_columns::<GreaterThanEqual>(
413					&value_col,
414					&lower_col,
415					fragment.clone(),
416					|f, l, r| {
417						TypeError::BinaryOperatorNotApplicable {
418							operator: BinaryOp::Between,
419							left: l,
420							right: r,
421							fragment: f,
422						}
423						.into_diagnostic()
424					},
425				)?;
426				let le_result = compare_columns::<LessThanEqual>(
427					&value_col,
428					&upper_col,
429					fragment.clone(),
430					|f, l, r| {
431						TypeError::BinaryOperatorNotApplicable {
432							operator: BinaryOp::Between,
433							left: l,
434							right: r,
435							fragment: f,
436						}
437						.into_diagnostic()
438					},
439				)?;
440
441				if !matches!(ge_result.data(), ColumnBuffer::Bool(_))
442					|| !matches!(le_result.data(), ColumnBuffer::Bool(_))
443				{
444					return Err(TypeError::BinaryOperatorNotApplicable {
445						operator: BinaryOp::Between,
446						left: value_col.get_type(),
447						right: lower_col.get_type(),
448						fragment: fragment.clone(),
449					}
450					.into());
451				}
452
453				match (ge_result.data(), le_result.data()) {
454					(ColumnBuffer::Bool(ge_container), ColumnBuffer::Bool(le_container)) => {
455						let mut data = Vec::with_capacity(ge_container.len());
456						let mut bitvec = Vec::with_capacity(ge_container.len());
457
458						for i in 0..ge_container.len() {
459							if ge_container.is_defined(i) && le_container.is_defined(i) {
460								data.push(ge_container.data().get(i)
461									&& le_container.data().get(i));
462								bitvec.push(true);
463							} else {
464								data.push(false);
465								bitvec.push(false);
466							}
467						}
468
469						Ok(ColumnWithName {
470							name: fragment.clone(),
471							data: ColumnBuffer::bool_with_bitvec(data, bitvec),
472						})
473					}
474					_ => unreachable!(
475						"Both comparison results should be boolean after the check above"
476					),
477				}
478			})
479		}
480
481		Expression::In(e) => {
482			let list_expressions = match e.list.as_ref() {
483				Expression::Tuple(tuple) => &tuple.expressions,
484				Expression::List(list) => &list.expressions,
485				_ => from_ref(e.list.as_ref()),
486			};
487			let value = compile_expression(_ctx, &e.value)?;
488			let list: Vec<CompiledExpr> = list_expressions
489				.iter()
490				.map(|expr| compile_expression(_ctx, expr))
491				.collect::<Result<Vec<_>>>()?;
492			let negated = e.negated;
493			let fragment = e.fragment.clone();
494			CompiledExpr::new(move |ctx| {
495				if list.is_empty() {
496					let value_col = value.execute(ctx)?;
497					let len = value_col.data().len();
498					let result = vec![negated; len];
499					return Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(result)));
500				}
501
502				let value_col = value.execute(ctx)?;
503
504				let first_col = list[0].execute(ctx)?;
505				let mut result = compare_columns::<Equal>(
506					&value_col,
507					&first_col,
508					fragment.clone(),
509					|f, l, r| {
510						TypeError::BinaryOperatorNotApplicable {
511							operator: BinaryOp::Equal,
512							left: l,
513							right: r,
514							fragment: f,
515						}
516						.into_diagnostic()
517					},
518				)?;
519
520				for list_expr in list.iter().skip(1) {
521					let list_col = list_expr.execute(ctx)?;
522					let eq_result = compare_columns::<Equal>(
523						&value_col,
524						&list_col,
525						fragment.clone(),
526						|f, l, r| {
527							TypeError::BinaryOperatorNotApplicable {
528								operator: BinaryOp::Equal,
529								left: l,
530								right: r,
531								fragment: f,
532							}
533							.into_diagnostic()
534						},
535					)?;
536					result = combine_bool_columns(result, eq_result, fragment.clone(), |l, r| {
537						l || r
538					})?;
539				}
540
541				if negated {
542					result = negate_column(result, fragment.clone());
543				}
544
545				Ok(result)
546			})
547		}
548
549		Expression::Contains(e) => {
550			let list_expressions = match e.list.as_ref() {
551				Expression::Tuple(tuple) => &tuple.expressions,
552				Expression::List(list) => &list.expressions,
553				_ => from_ref(e.list.as_ref()),
554			};
555			let value = compile_expression(_ctx, &e.value)?;
556			let list: Vec<CompiledExpr> = list_expressions
557				.iter()
558				.map(|expr| compile_expression(_ctx, expr))
559				.collect::<Result<Vec<_>>>()?;
560			let fragment = e.fragment.clone();
561			CompiledExpr::new(move |ctx| {
562				let value_col = value.execute(ctx)?;
563
564				if list.is_empty() {
565					let len = value_col.data().len();
566					let result = vec![true; len];
567					return Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(result)));
568				}
569
570				let first_col = list[0].execute(ctx)?;
571				let mut result = list_contains_element(&value_col, &first_col, &fragment)?;
572
573				for list_expr in list.iter().skip(1) {
574					let list_col = list_expr.execute(ctx)?;
575					let element_result = list_contains_element(&value_col, &list_col, &fragment)?;
576					result = combine_bool_columns(
577						result,
578						element_result,
579						fragment.clone(),
580						|l, r| l && r,
581					)?;
582				}
583
584				Ok(result)
585			})
586		}
587
588		Expression::Cast(e) => {
589			let label = display_label(expr);
590			if let Expression::Constant(const_expr) = e.expression.as_ref() {
591				let const_expr = const_expr.clone();
592				let target_type = e.to.ty.clone();
593				CompiledExpr::new(move |ctx| {
594					let row_count = ctx.take.unwrap_or(ctx.row_count);
595					let data = constant_value(&const_expr, row_count)?;
596					let casted = if data.get_type() == target_type {
597						data
598					} else {
599						constant_value_of(&const_expr, target_type.clone(), row_count)?
600					};
601					Ok(ColumnWithName::new(label.clone(), casted))
602				})
603			} else {
604				let inner = compile_expression(_ctx, &e.expression)?;
605				let target_type = e.to.ty.clone();
606				let inner_fragment = e.expression.full_fragment_owned();
607				CompiledExpr::new(move |ctx| {
608					let column = inner.execute(ctx)?;
609					let frag = inner_fragment.clone();
610					let casted = cast_column_data(ctx, column.data(), target_type.clone(), &|| {
611						inner_fragment.clone()
612					})
613					.map_err(|e| {
614						Error::from(CastError::InvalidNumber {
615							fragment: frag,
616							target: target_type.clone(),
617							cause: e.diagnostic(),
618						})
619					})?;
620					Ok(ColumnWithName::new(label.clone(), casted))
621				})
622			}
623		}
624
625		Expression::If(e) => {
626			let condition = compile_expression(_ctx, &e.condition)?;
627			let then_expr = compile_expressions(_ctx, from_ref(e.then_expr.as_ref()))?;
628			let else_ifs: Vec<(CompiledExpr, Vec<CompiledExpr>)> = e
629				.else_ifs
630				.iter()
631				.map(|ei| {
632					Ok((
633						compile_expression(_ctx, &ei.condition)?,
634						compile_expressions(_ctx, from_ref(ei.then_expr.as_ref()))?,
635					))
636				})
637				.collect::<Result<Vec<_>>>()?;
638			let else_branch: Option<Vec<CompiledExpr>> = match &e.else_expr {
639				Some(expr) => Some(compile_expressions(_ctx, from_ref(expr.as_ref()))?),
640				None => None,
641			};
642			let fragment = e.fragment.clone();
643			CompiledExpr::new_multi(move |ctx| {
644				execute_if_multi(ctx, &condition, &then_expr, &else_ifs, &else_branch, &fragment)
645			})
646		}
647
648		Expression::Map(e) => {
649			let expressions = compile_expressions(_ctx, &e.expressions)?;
650			CompiledExpr::new_multi(move |ctx| execute_projection_multi(ctx, &expressions))
651		}
652
653		Expression::Extend(e) => {
654			let expressions = compile_expressions(_ctx, &e.expressions)?;
655			CompiledExpr::new_multi(move |ctx| execute_projection_multi(ctx, &expressions))
656		}
657
658		Expression::Call(e) => {
659			let compiled_args: Vec<CompiledExpr> =
660				e.args.iter().map(|arg| compile_expression(_ctx, arg)).collect::<Result<Vec<_>>>()?;
661			let expr = e.clone();
662			CompiledExpr::new(move |ctx| {
663				let mut arg_columns = Vec::with_capacity(compiled_args.len());
664				for compiled_arg in &compiled_args {
665					arg_columns.push(compiled_arg.execute(ctx)?);
666				}
667				let arguments = Columns::new(arg_columns);
668				call_builtin(ctx, &expr, arguments)
669			})
670		}
671
672		Expression::SumTypeConstructor(_) => {
673			panic!(
674				"SumTypeConstructor in expression context — constructors should be expanded by InlineDataNode before expression compilation"
675			);
676		}
677
678		Expression::IsVariant(e) => {
679			let col_name = match e.expression.as_ref() {
680				Expression::Column(c) => c.0.name.text().to_string(),
681				other => display_label(other).text().to_string(),
682			};
683			let tag_col_name = format!("{}_tag", col_name);
684			let tag = e.tag.expect("IS variant tag must be resolved before compilation");
685			let fragment = e.fragment.clone();
686			CompiledExpr::new(move |ctx| {
687				if let Some(tag_col) =
688					ctx.columns.iter().find(|c| c.name().text() == tag_col_name.as_str())
689				{
690					match tag_col.data() {
691						ColumnBuffer::Uint1(container) => {
692							let results: Vec<bool> = container
693								.iter()
694								.take(ctx.row_count)
695								.map(|v| v == Some(tag))
696								.collect();
697							Ok(ColumnWithName::new(
698								fragment.clone(),
699								ColumnBuffer::bool(results),
700							))
701						}
702						_ => Ok(ColumnWithName {
703							name: fragment.clone(),
704							data: ColumnBuffer::none_typed(
705								ValueType::Boolean,
706								ctx.row_count,
707							),
708						}),
709					}
710				} else {
711					Ok(ColumnWithName {
712						name: fragment.clone(),
713						data: ColumnBuffer::none_typed(ValueType::Boolean, ctx.row_count),
714					})
715				}
716			})
717		}
718
719		Expression::FieldAccess(e) => {
720			let field_name = e.field.text().to_string();
721
722			let var_name = match e.object.as_ref() {
723				Expression::Variable(var_expr) => Some(var_expr.name().to_string()),
724				_ => None,
725			};
726			let object = compile_expression(_ctx, &e.object)?;
727			CompiledExpr::new(move |ctx| {
728				if let Some(ref variable_name) = var_name {
729					match ctx.symbols.get(variable_name) {
730						Some(Variable::Columns {
731							columns,
732						}) if !columns.is_scalar() => {
733							let col_pos = columns
734								.names
735								.iter()
736								.position(|n| n.text() == field_name);
737							match col_pos {
738								Some(pos) => {
739									let value = columns.columns[pos].get_value(0);
740									let row_count =
741										ctx.take.unwrap_or(ctx.row_count);
742									let mut data = ColumnBuffer::with_capacity(
743										value.get_type(),
744										row_count,
745									);
746									for _ in 0..row_count {
747										data.push_value(value.clone());
748									}
749									Ok(ColumnWithName {
750										name: Fragment::internal(&field_name),
751										data,
752									})
753								}
754								None => {
755									let available: Vec<String> = columns
756										.names
757										.iter()
758										.map(|n| n.text().to_string())
759										.collect();
760									Err(TypeError::Runtime {
761										kind: RuntimeErrorKind::FieldNotFound {
762											variable: variable_name
763												.to_string(),
764											field: field_name.to_string(),
765											available,
766										},
767										message: format!(
768											"Field '{}' not found on variable '{}'",
769											field_name, variable_name
770										),
771									}
772									.into())
773								}
774							}
775						}
776						Some(Variable::Columns {
777							..
778						})
779						| Some(Variable::Closure(_)) => Err(TypeError::Runtime {
780							kind: RuntimeErrorKind::FieldNotFound {
781								variable: variable_name.to_string(),
782								field: field_name.to_string(),
783								available: vec![],
784							},
785							message: format!(
786								"Field '{}' not found on variable '{}'",
787								field_name, variable_name
788							),
789						}
790						.into()),
791						Some(Variable::ForIterator {
792							..
793						}) => Err(TypeError::Runtime {
794							kind: RuntimeErrorKind::VariableIsDataframe {
795								name: variable_name.to_string(),
796							},
797							message: format!(
798								"Variable '{}' contains a dataframe and cannot be used directly in scalar expressions",
799								variable_name
800							),
801						}
802						.into()),
803						None => Err(TypeError::Runtime {
804							kind: RuntimeErrorKind::VariableNotFound {
805								name: variable_name.to_string(),
806							},
807							message: format!("Variable '{}' is not defined", variable_name),
808						}
809						.into()),
810					}
811				} else {
812					let _obj_col = object.execute(ctx)?;
813					Err(TypeError::Runtime {
814						kind: RuntimeErrorKind::FieldNotFound {
815							variable: "<expression>".to_string(),
816							field: field_name.to_string(),
817							available: vec![],
818						},
819						message: format!(
820							"Field '{}' not found on variable '<expression>'",
821							field_name
822						),
823					}
824					.into())
825				}
826			})
827		}
828	})
829}
830
831fn compile_expressions(ctx: &CompileContext, exprs: &[Expression]) -> Result<Vec<CompiledExpr>> {
832	exprs.iter().map(|e| compile_expression(ctx, e)).collect()
833}
834
835fn combine_bool_columns(
836	left: ColumnWithName,
837	right: ColumnWithName,
838	fragment: Fragment,
839	combine_fn: fn(bool, bool) -> bool,
840) -> Result<ColumnWithName> {
841	binary_op_unwrap_option(&left, &right, fragment.clone(), |left, right| match (left.data(), right.data()) {
842		(ColumnBuffer::Bool(l), ColumnBuffer::Bool(r)) => {
843			let len = l.len();
844			let mut data = Vec::with_capacity(len);
845			let mut bitvec = Vec::with_capacity(len);
846
847			for i in 0..len {
848				let l_defined = l.is_defined(i);
849				let r_defined = r.is_defined(i);
850				let l_val = l.data().get(i);
851				let r_val = r.data().get(i);
852
853				if l_defined && r_defined {
854					data.push(combine_fn(l_val, r_val));
855					bitvec.push(true);
856				} else {
857					data.push(false);
858					bitvec.push(false);
859				}
860			}
861
862			Ok(ColumnWithName {
863				name: fragment.clone(),
864				data: ColumnBuffer::bool_with_bitvec(data, bitvec),
865			})
866		}
867		_ => {
868			unreachable!("combine_bool_columns should only be called with boolean columns")
869		}
870	})
871}
872
873fn list_items_contain(items: &[Value], element: &Value, fragment: &Fragment) -> bool {
874	if items.iter().any(|item| item == element) {
875		return true;
876	}
877	if items.is_empty() {
878		return false;
879	}
880
881	if let Some(items_buf) = build_homogeneous_buffer(items) {
882		let elems_buf = ColumnBuffer::from_many(element.clone(), items.len());
883		let items_col = ColumnWithName::new(fragment.clone(), items_buf);
884		let elems_col = ColumnWithName::new(fragment.clone(), elems_buf);
885		return compare_columns::<Equal>(&items_col, &elems_col, fragment.clone(), |f, l, r| {
886			TypeError::BinaryOperatorNotApplicable {
887				operator: BinaryOp::Equal,
888				left: l,
889				right: r,
890				fragment: f,
891			}
892			.into_diagnostic()
893		})
894		.map(|c| bool_column_has_true(&c))
895		.unwrap_or(false);
896	}
897
898	list_items_contain_per_item(items, element, fragment)
899}
900
901fn list_items_contain_per_item(items: &[Value], element: &Value, fragment: &Fragment) -> bool {
902	items.iter().any(|item| {
903		let item_col = ColumnWithName::new(fragment.clone(), ColumnBuffer::from(item.clone()));
904		let elem_col = ColumnWithName::new(fragment.clone(), ColumnBuffer::from(element.clone()));
905		compare_columns::<Equal>(&item_col, &elem_col, fragment.clone(), |f, l, r| {
906			TypeError::BinaryOperatorNotApplicable {
907				operator: BinaryOp::Equal,
908				left: l,
909				right: r,
910				fragment: f,
911			}
912			.into_diagnostic()
913		})
914		.ok()
915		.and_then(|c| match c.data() {
916			ColumnBuffer::Bool(b) => Some(b.data().get(0)),
917			_ => None,
918		})
919		.unwrap_or(false)
920	})
921}
922
923fn bool_column_has_true(col: &ColumnWithName) -> bool {
924	match col.data() {
925		ColumnBuffer::Bool(b) => b.data().any(),
926		ColumnBuffer::Option {
927			inner,
928			bitvec,
929		} => match inner.as_ref() {
930			ColumnBuffer::Bool(b) => {
931				let n = bitvec.len().min(b.len());
932				(0..n).any(|i| bitvec.get(i) && b.data().get(i))
933			}
934			_ => false,
935		},
936		_ => false,
937	}
938}
939
940fn build_homogeneous_buffer(items: &[Value]) -> Option<ColumnBuffer> {
941	let first = items.first()?;
942	let first_disc = discriminant(first);
943	if !items.iter().all(|v| discriminant(v) == first_disc) {
944		return None;
945	}
946
947	macro_rules! collect {
948		($variant:ident, $constructor:ident, |$x:ident| $convert:expr) => {{
949			let data: Vec<_> = items
950				.iter()
951				.map(|v| match v {
952					Value::$variant($x) => $convert,
953					_ => unreachable!("homogeneous check guarantees variant"),
954				})
955				.collect();
956			Some(ColumnBuffer::$constructor(data))
957		}};
958	}
959
960	match first {
961		Value::Boolean(_) => collect!(Boolean, bool, |x| *x),
962		Value::Float4(_) => collect!(Float4, float4, |x| x.value()),
963		Value::Float8(_) => collect!(Float8, float8, |x| x.value()),
964		Value::Int1(_) => collect!(Int1, int1, |x| *x),
965		Value::Int2(_) => collect!(Int2, int2, |x| *x),
966		Value::Int4(_) => collect!(Int4, int4, |x| *x),
967		Value::Int8(_) => collect!(Int8, int8, |x| *x),
968		Value::Int16(_) => collect!(Int16, int16, |x| *x),
969		Value::Uint1(_) => collect!(Uint1, uint1, |x| *x),
970		Value::Uint2(_) => collect!(Uint2, uint2, |x| *x),
971		Value::Uint4(_) => collect!(Uint4, uint4, |x| *x),
972		Value::Uint8(_) => collect!(Uint8, uint8, |x| *x),
973		Value::Uint16(_) => collect!(Uint16, uint16, |x| *x),
974		Value::Utf8(_) => collect!(Utf8, utf8, |x| x.clone()),
975		Value::Date(_) => collect!(Date, date, |x| *x),
976		Value::DateTime(_) => collect!(DateTime, datetime, |x| *x),
977		Value::Time(_) => collect!(Time, time, |x| *x),
978		Value::Duration(_) => collect!(Duration, duration, |x| *x),
979		Value::Uuid4(_) => collect!(Uuid4, uuid4, |x| *x),
980		Value::Uuid7(_) => collect!(Uuid7, uuid7, |x| *x),
981		Value::IdentityId(_) => collect!(IdentityId, identity_id, |x| *x),
982		Value::Blob(_) => collect!(Blob, blob, |x| x.clone()),
983		Value::Int(_) => collect!(Int, int, |x| x.clone()),
984		Value::Uint(_) => collect!(Uint, uint, |x| x.clone()),
985		Value::Decimal(_) => collect!(Decimal, decimal, |x| x.clone()),
986		Value::DictionaryId(_) => collect!(DictionaryId, dictionary_id, |x| *x),
987
988		_ => None,
989	}
990}
991
992fn list_contains_element(
993	list_col: &ColumnWithName,
994	element_col: &ColumnWithName,
995	fragment: &Fragment,
996) -> Result<ColumnWithName> {
997	let len = list_col.data().len();
998	let mut data = Vec::with_capacity(len);
999
1000	for i in 0..len {
1001		let list_value = list_col.data().get_value(i);
1002		let element_value = element_col.data().get_value(i);
1003
1004		let contained = match &list_value {
1005			Value::List(items) => list_items_contain(items, &element_value, fragment),
1006			Value::Tuple(items) => list_items_contain(items, &element_value, fragment),
1007			Value::Any(boxed) => match boxed.as_ref() {
1008				Value::List(items) => list_items_contain(items, &element_value, fragment),
1009				Value::Tuple(items) => list_items_contain(items, &element_value, fragment),
1010				_ => false,
1011			},
1012			_ => false,
1013		};
1014		data.push(contained);
1015	}
1016
1017	Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(data)))
1018}
1019
1020fn negate_column(col: ColumnWithName, fragment: Fragment) -> ColumnWithName {
1021	unary_op_unwrap_option(&col, |col| match col.data() {
1022		ColumnBuffer::Bool(container) => {
1023			let len = container.len();
1024			let mut data = Vec::with_capacity(len);
1025			let mut bitvec = Vec::with_capacity(len);
1026
1027			for i in 0..len {
1028				if container.is_defined(i) {
1029					data.push(!container.data().get(i));
1030					bitvec.push(true);
1031				} else {
1032					data.push(false);
1033					bitvec.push(false);
1034				}
1035			}
1036
1037			Ok(ColumnWithName {
1038				name: fragment.clone(),
1039				data: ColumnBuffer::bool_with_bitvec(data, bitvec),
1040			})
1041		}
1042		_ => unreachable!("negate_column should only be called with boolean columns"),
1043	})
1044	.unwrap()
1045}
1046
1047fn is_truthy(value: &Value) -> bool {
1048	match value {
1049		Value::Boolean(true) => true,
1050		Value::Boolean(false) => false,
1051		Value::None {
1052			..
1053		} => false,
1054		Value::Int1(0) | Value::Int2(0) | Value::Int4(0) | Value::Int8(0) | Value::Int16(0) => false,
1055		Value::Uint1(0) | Value::Uint2(0) | Value::Uint4(0) | Value::Uint8(0) | Value::Uint16(0) => false,
1056		Value::Int1(_) | Value::Int2(_) | Value::Int4(_) | Value::Int8(_) | Value::Int16(_) => true,
1057		Value::Uint1(_) | Value::Uint2(_) | Value::Uint4(_) | Value::Uint8(_) | Value::Uint16(_) => true,
1058		Value::Utf8(s) => !s.is_empty(),
1059		_ => true,
1060	}
1061}
1062
1063fn execute_if_multi(
1064	ctx: &EvalContext,
1065	condition: &CompiledExpr,
1066	then_expr: &[CompiledExpr],
1067	else_ifs: &[(CompiledExpr, Vec<CompiledExpr>)],
1068	else_branch: &Option<Vec<CompiledExpr>>,
1069	_fragment: &Fragment,
1070) -> Result<Vec<ColumnWithName>> {
1071	let condition_column = condition.execute(ctx)?;
1072
1073	let mut result_data: Option<Vec<ColumnBuffer>> = None;
1074	let mut result_names: Vec<Fragment> = Vec::new();
1075
1076	for row_idx in 0..ctx.row_count {
1077		let condition_value = condition_column.data().get_value(row_idx);
1078
1079		let branch_results = if is_truthy(&condition_value) {
1080			execute_multi_exprs(ctx, then_expr)?
1081		} else {
1082			let mut found_branch = false;
1083			let mut branch_columns = None;
1084
1085			for (else_if_condition, else_if_then) in else_ifs {
1086				let else_if_col = else_if_condition.execute(ctx)?;
1087				let else_if_value = else_if_col.data().get_value(row_idx);
1088
1089				if is_truthy(&else_if_value) {
1090					branch_columns = Some(execute_multi_exprs(ctx, else_if_then)?);
1091					found_branch = true;
1092					break;
1093				}
1094			}
1095
1096			if found_branch {
1097				branch_columns.unwrap()
1098			} else if let Some(else_exprs) = else_branch {
1099				execute_multi_exprs(ctx, else_exprs)?
1100			} else {
1101				vec![]
1102			}
1103		};
1104
1105		let is_empty_result = branch_results.is_empty();
1106		if is_empty_result {
1107			if let Some(data) = result_data.as_mut() {
1108				for col_data in data.iter_mut() {
1109					col_data.push_value(Value::none());
1110				}
1111			}
1112			continue;
1113		}
1114
1115		if result_data.is_none() {
1116			let mut data: Vec<ColumnBuffer> = branch_results
1117				.iter()
1118				.map(|col| ColumnBuffer::with_capacity(col.data().get_type(), ctx.row_count))
1119				.collect();
1120			for _ in 0..row_idx {
1121				for col_data in data.iter_mut() {
1122					col_data.push_value(Value::none());
1123				}
1124			}
1125			result_data = Some(data);
1126			result_names = branch_results.iter().map(|col| col.name.clone()).collect();
1127		}
1128
1129		let data = result_data.as_mut().unwrap();
1130		for (i, branch_col) in branch_results.iter().enumerate() {
1131			if i < data.len() {
1132				let branch_value = branch_col.data().get_value(row_idx);
1133				data[i].push_value(branch_value);
1134			}
1135		}
1136	}
1137
1138	let result_data = result_data.unwrap_or_default();
1139	let result: Vec<ColumnWithName> = result_data
1140		.into_iter()
1141		.enumerate()
1142		.map(|(i, data)| ColumnWithName {
1143			name: result_names.get(i).cloned().unwrap_or_else(|| Fragment::internal("column")),
1144			data,
1145		})
1146		.collect();
1147
1148	if result.is_empty() {
1149		Ok(vec![ColumnWithName {
1150			name: Fragment::internal("none"),
1151			data: ColumnBuffer::none_typed(ValueType::Boolean, ctx.row_count),
1152		}])
1153	} else {
1154		Ok(result)
1155	}
1156}
1157
1158fn execute_multi_exprs(ctx: &EvalContext, exprs: &[CompiledExpr]) -> Result<Vec<ColumnWithName>> {
1159	let mut result = Vec::new();
1160	for expr in exprs {
1161		result.extend(expr.execute_multi(ctx)?);
1162	}
1163	Ok(result)
1164}
1165
1166fn execute_projection_multi(ctx: &EvalContext, expressions: &[CompiledExpr]) -> Result<Vec<ColumnWithName>> {
1167	let mut result = Vec::with_capacity(expressions.len());
1168
1169	for expr in expressions {
1170		let column = expr.execute(ctx)?;
1171		let name = column.name.text().to_string();
1172		result.push(ColumnWithName::new(Fragment::internal(name), column.data));
1173	}
1174
1175	Ok(result)
1176}
reifydb_engine/expression/compile.rs

reifydb_engine/expression/
compile.rs
