1use 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_type::{
9 error::{BinaryOp, Error, IntoDiagnostic, LogicalOp, RuntimeErrorKind, TypeError},
10 fragment::Fragment,
11 value::{Value, r#type::Type},
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(Type::Option(Box::new(Type::Boolean)), 0),
87 }))
88 }
89 }
90 }
91
92 pub fn execute_multi(&self, ctx: &EvalContext) -> Result<Vec<ColumnWithName>> {
93 match &self.inner {
94 CompiledExprInner::Single(f) => Ok(vec![f(ctx)?]),
95 CompiledExprInner::Multi(f) => f(ctx),
96 }
97 }
98}
99
100macro_rules! compile_arith {
101 ($ctx:expr, $parent:expr, $e:expr, $op_fn:path) => {{
102 let left = compile_expression($ctx, &$e.left)?;
103 let right = compile_expression($ctx, &$e.right)?;
104 let fragment = $e.full_fragment_owned();
105 let label = display_label($parent);
106 CompiledExpr::new(move |ctx| {
107 let l = left.execute(ctx)?;
108 let r = right.execute(ctx)?;
109 let mut col = $op_fn(ctx, &l, &r, || fragment.clone())?;
110 col.name = label.clone();
111 Ok(col)
112 })
113 }};
114}
115
116macro_rules! compile_compare {
117 ($ctx:expr, $parent:expr, $e:expr, $cmp_type:ty, $binary_op:expr) => {{
118 let left = compile_expression($ctx, &$e.left)?;
119 let right = compile_expression($ctx, &$e.right)?;
120 let fragment = $e.full_fragment_owned();
121 let label = display_label($parent);
122 CompiledExpr::new(move |ctx| {
123 let l = left.execute(ctx)?;
124 let r = right.execute(ctx)?;
125 let mut col = compare_columns::<$cmp_type>(&l, &r, fragment.clone(), |f, l, r| {
126 TypeError::BinaryOperatorNotApplicable {
127 operator: $binary_op,
128 left: l,
129 right: r,
130 fragment: f,
131 }
132 .into_diagnostic()
133 })?;
134 col.name = label.clone();
135 Ok(col)
136 })
137 }};
138}
139
140pub 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) {
214 let mut data = ColumnBuffer::with_capacity(
215 value.get_type(),
216 ctx.row_count,
217 );
218 for _ in 0..ctx.row_count {
219 data.push_value(value.clone());
220 }
221 return Ok(ColumnWithName {
222 name: Fragment::internal(variable_name),
223 data,
224 });
225 }
226 Err(TypeError::Runtime {
227 kind: RuntimeErrorKind::VariableNotFound {
228 name: variable_name.to_string(),
229 },
230 message: format!("Variable '{}' is not defined", variable_name),
231 }
232 .into())
233 }
234 }
235 })
236 }
237
238 Expression::Parameter(e) => {
239 let expr = e.clone();
240 CompiledExpr::new(move |ctx| parameter_lookup(ctx, &expr))
241 }
242
243 Expression::Alias(e) => {
244 let inner = compile_expression(_ctx, &e.expression)?;
245 let alias = e.alias.0.clone();
246 CompiledExpr::new(move |ctx| {
247 let mut column = inner.execute(ctx)?;
248 column.name = alias.clone();
249 Ok(column)
250 })
251 }
252
253 Expression::Add(e) => compile_arith!(_ctx, expr, e, add_columns),
254 Expression::Sub(e) => compile_arith!(_ctx, expr, e, sub_columns),
255 Expression::Mul(e) => compile_arith!(_ctx, expr, e, mul_columns),
256 Expression::Div(e) => compile_arith!(_ctx, expr, e, div_columns),
257 Expression::Rem(e) => compile_arith!(_ctx, expr, e, rem_columns),
258
259 Expression::Equal(e) => compile_compare!(_ctx, expr, e, Equal, BinaryOp::Equal),
260 Expression::NotEqual(e) => compile_compare!(_ctx, expr, e, NotEqual, BinaryOp::NotEqual),
261 Expression::GreaterThan(e) => compile_compare!(_ctx, expr, e, GreaterThan, BinaryOp::GreaterThan),
262 Expression::GreaterThanEqual(e) => {
263 compile_compare!(_ctx, expr, e, GreaterThanEqual, BinaryOp::GreaterThanEqual)
264 }
265 Expression::LessThan(e) => compile_compare!(_ctx, expr, e, LessThan, BinaryOp::LessThan),
266 Expression::LessThanEqual(e) => compile_compare!(_ctx, expr, e, LessThanEqual, BinaryOp::LessThanEqual),
267
268 Expression::And(e) => {
269 let left = compile_expression(_ctx, &e.left)?;
270 let right = compile_expression(_ctx, &e.right)?;
271 let fragment = e.full_fragment_owned();
272 let label = display_label(expr);
273 CompiledExpr::new(move |ctx| {
274 let l = left.execute(ctx)?;
275 if let Some(mut short) = try_short_circuit_and(&l, &fragment, l.data().len()) {
276 short.name = label.clone();
277 return Ok(short);
278 }
279 let r = right.execute(ctx)?;
280 let mut col = execute_logical_op(&l, &r, &fragment, LogicalOp::And, |a, b| a && b)?;
281 col.name = label.clone();
282 Ok(col)
283 })
284 }
285
286 Expression::Or(e) => {
287 let left = compile_expression(_ctx, &e.left)?;
288 let right = compile_expression(_ctx, &e.right)?;
289 let fragment = e.full_fragment_owned();
290 let label = display_label(expr);
291 CompiledExpr::new(move |ctx| {
292 let l = left.execute(ctx)?;
293 if let Some(mut short) = try_short_circuit_or(&l, &fragment, l.data().len()) {
294 short.name = label.clone();
295 return Ok(short);
296 }
297 let r = right.execute(ctx)?;
298 let mut col = execute_logical_op(&l, &r, &fragment, LogicalOp::Or, |a, b| a || b)?;
299 col.name = label.clone();
300 Ok(col)
301 })
302 }
303
304 Expression::Xor(e) => {
305 let left = compile_expression(_ctx, &e.left)?;
306 let right = compile_expression(_ctx, &e.right)?;
307 let fragment = e.full_fragment_owned();
308 let label = display_label(expr);
309 CompiledExpr::new(move |ctx| {
310 let l = left.execute(ctx)?;
311 let r = right.execute(ctx)?;
312 let mut col = execute_logical_op(&l, &r, &fragment, LogicalOp::Xor, |a, b| a != b)?;
313 col.name = label.clone();
314 Ok(col)
315 })
316 }
317
318 Expression::Prefix(e) => {
319 let inner = compile_expression(_ctx, &e.expression)?;
320 let operator = e.operator.clone();
321 let fragment = e.full_fragment_owned();
322 let label = display_label(expr);
323 CompiledExpr::new(move |ctx| {
324 let column = inner.execute(ctx)?;
325 let mut col = prefix_apply(&column, &operator, &fragment)?;
326 col.name = label.clone();
327 Ok(col)
328 })
329 }
330
331 Expression::Type(e) => {
332 let ty = e.ty.clone();
333 let fragment = e.fragment.clone();
334 CompiledExpr::new(move |ctx| {
335 let row_count = ctx.take.unwrap_or(ctx.row_count);
336 let values: Vec<Box<Value>> =
337 (0..row_count).map(|_| Box::new(Value::Type(ty.clone()))).collect();
338 Ok(ColumnWithName::new(fragment.text(), ColumnBuffer::any(values)))
339 })
340 }
341
342 Expression::AccessSource(e) => {
343 let col_name = e.column.name.text().to_string();
344 let expr = e.clone();
345 CompiledExpr::new_access(col_name, move |ctx| access_lookup(ctx, &expr))
346 }
347
348 Expression::Tuple(e) => {
349 if e.expressions.len() == 1 {
350 let inner = compile_expression(_ctx, &e.expressions[0])?;
351 CompiledExpr::new(move |ctx| inner.execute(ctx))
352 } else {
353 let compiled: Vec<CompiledExpr> = e
354 .expressions
355 .iter()
356 .map(|expr| compile_expression(_ctx, expr))
357 .collect::<Result<Vec<_>>>()?;
358 let fragment = e.fragment.clone();
359 CompiledExpr::new(move |ctx| {
360 let columns: Vec<ColumnWithName> = compiled
361 .iter()
362 .map(|expr| expr.execute(ctx))
363 .collect::<Result<Vec<_>>>()?;
364
365 let len = columns.first().map_or(1, |c| c.data().len());
366 let mut data: Vec<Box<Value>> = Vec::with_capacity(len);
367
368 for i in 0..len {
369 let items: Vec<Value> =
370 columns.iter().map(|col| col.data().get_value(i)).collect();
371 data.push(Box::new(Value::Tuple(items)));
372 }
373
374 Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::any(data)))
375 })
376 }
377 }
378
379 Expression::List(e) => {
380 let compiled: Vec<CompiledExpr> = e
381 .expressions
382 .iter()
383 .map(|expr| compile_expression(_ctx, expr))
384 .collect::<Result<Vec<_>>>()?;
385 let fragment = e.fragment.clone();
386 CompiledExpr::new(move |ctx| {
387 let columns: Vec<ColumnWithName> =
388 compiled.iter().map(|expr| expr.execute(ctx)).collect::<Result<Vec<_>>>()?;
389
390 let len = columns.first().map_or(1, |c| c.data().len());
391 let mut data: Vec<Box<Value>> = Vec::with_capacity(len);
392
393 for i in 0..len {
394 let items: Vec<Value> =
395 columns.iter().map(|col| col.data().get_value(i)).collect();
396 data.push(Box::new(Value::List(items)));
397 }
398
399 Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::any(data)))
400 })
401 }
402
403 Expression::Between(e) => {
404 let value = compile_expression(_ctx, &e.value)?;
405 let lower = compile_expression(_ctx, &e.lower)?;
406 let upper = compile_expression(_ctx, &e.upper)?;
407 let fragment = e.fragment.clone();
408 CompiledExpr::new(move |ctx| {
409 let value_col = value.execute(ctx)?;
410 let lower_col = lower.execute(ctx)?;
411 let upper_col = upper.execute(ctx)?;
412
413 let ge_result = compare_columns::<GreaterThanEqual>(
414 &value_col,
415 &lower_col,
416 fragment.clone(),
417 |f, l, r| {
418 TypeError::BinaryOperatorNotApplicable {
419 operator: BinaryOp::Between,
420 left: l,
421 right: r,
422 fragment: f,
423 }
424 .into_diagnostic()
425 },
426 )?;
427 let le_result = compare_columns::<LessThanEqual>(
428 &value_col,
429 &upper_col,
430 fragment.clone(),
431 |f, l, r| {
432 TypeError::BinaryOperatorNotApplicable {
433 operator: BinaryOp::Between,
434 left: l,
435 right: r,
436 fragment: f,
437 }
438 .into_diagnostic()
439 },
440 )?;
441
442 if !matches!(ge_result.data(), ColumnBuffer::Bool(_))
443 || !matches!(le_result.data(), ColumnBuffer::Bool(_))
444 {
445 return Err(TypeError::BinaryOperatorNotApplicable {
446 operator: BinaryOp::Between,
447 left: value_col.get_type(),
448 right: lower_col.get_type(),
449 fragment: fragment.clone(),
450 }
451 .into());
452 }
453
454 match (ge_result.data(), le_result.data()) {
455 (ColumnBuffer::Bool(ge_container), ColumnBuffer::Bool(le_container)) => {
456 let mut data = Vec::with_capacity(ge_container.len());
457 let mut bitvec = Vec::with_capacity(ge_container.len());
458
459 for i in 0..ge_container.len() {
460 if ge_container.is_defined(i) && le_container.is_defined(i) {
461 data.push(ge_container.data().get(i)
462 && le_container.data().get(i));
463 bitvec.push(true);
464 } else {
465 data.push(false);
466 bitvec.push(false);
467 }
468 }
469
470 Ok(ColumnWithName {
471 name: fragment.clone(),
472 data: ColumnBuffer::bool_with_bitvec(data, bitvec),
473 })
474 }
475 _ => unreachable!(
476 "Both comparison results should be boolean after the check above"
477 ),
478 }
479 })
480 }
481
482 Expression::In(e) => {
483 let list_expressions = match e.list.as_ref() {
484 Expression::Tuple(tuple) => &tuple.expressions,
485 Expression::List(list) => &list.expressions,
486 _ => from_ref(e.list.as_ref()),
487 };
488 let value = compile_expression(_ctx, &e.value)?;
489 let list: Vec<CompiledExpr> = list_expressions
490 .iter()
491 .map(|expr| compile_expression(_ctx, expr))
492 .collect::<Result<Vec<_>>>()?;
493 let negated = e.negated;
494 let fragment = e.fragment.clone();
495 CompiledExpr::new(move |ctx| {
496 if list.is_empty() {
497 let value_col = value.execute(ctx)?;
498 let len = value_col.data().len();
499 let result = vec![negated; len];
500 return Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(result)));
501 }
502
503 let value_col = value.execute(ctx)?;
504
505 let first_col = list[0].execute(ctx)?;
506 let mut result = compare_columns::<Equal>(
507 &value_col,
508 &first_col,
509 fragment.clone(),
510 |f, l, r| {
511 TypeError::BinaryOperatorNotApplicable {
512 operator: BinaryOp::Equal,
513 left: l,
514 right: r,
515 fragment: f,
516 }
517 .into_diagnostic()
518 },
519 )?;
520
521 for list_expr in list.iter().skip(1) {
522 let list_col = list_expr.execute(ctx)?;
523 let eq_result = compare_columns::<Equal>(
524 &value_col,
525 &list_col,
526 fragment.clone(),
527 |f, l, r| {
528 TypeError::BinaryOperatorNotApplicable {
529 operator: BinaryOp::Equal,
530 left: l,
531 right: r,
532 fragment: f,
533 }
534 .into_diagnostic()
535 },
536 )?;
537 result = combine_bool_columns(result, eq_result, fragment.clone(), |l, r| {
538 l || r
539 })?;
540 }
541
542 if negated {
543 result = negate_column(result, fragment.clone());
544 }
545
546 Ok(result)
547 })
548 }
549
550 Expression::Contains(e) => {
551 let list_expressions = match e.list.as_ref() {
552 Expression::Tuple(tuple) => &tuple.expressions,
553 Expression::List(list) => &list.expressions,
554 _ => from_ref(e.list.as_ref()),
555 };
556 let value = compile_expression(_ctx, &e.value)?;
557 let list: Vec<CompiledExpr> = list_expressions
558 .iter()
559 .map(|expr| compile_expression(_ctx, expr))
560 .collect::<Result<Vec<_>>>()?;
561 let fragment = e.fragment.clone();
562 CompiledExpr::new(move |ctx| {
563 let value_col = value.execute(ctx)?;
564
565 if list.is_empty() {
567 let len = value_col.data().len();
568 let result = vec![true; len];
569 return Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(result)));
570 }
571
572 let first_col = list[0].execute(ctx)?;
574 let mut result = list_contains_element(&value_col, &first_col, &fragment)?;
575
576 for list_expr in list.iter().skip(1) {
577 let list_col = list_expr.execute(ctx)?;
578 let element_result = list_contains_element(&value_col, &list_col, &fragment)?;
579 result = combine_bool_columns(
580 result,
581 element_result,
582 fragment.clone(),
583 |l, r| l && r,
584 )?;
585 }
586
587 Ok(result)
588 })
589 }
590
591 Expression::Cast(e) => {
592 let label = display_label(expr);
593 if let Expression::Constant(const_expr) = e.expression.as_ref() {
594 let const_expr = const_expr.clone();
595 let target_type = e.to.ty.clone();
596 CompiledExpr::new(move |ctx| {
597 let row_count = ctx.take.unwrap_or(ctx.row_count);
598 let data = constant_value(&const_expr, row_count)?;
599 let casted = if data.get_type() == target_type {
600 data
601 } else {
602 constant_value_of(&const_expr, target_type.clone(), row_count)?
603 };
604 Ok(ColumnWithName::new(label.clone(), casted))
605 })
606 } else {
607 let inner = compile_expression(_ctx, &e.expression)?;
608 let target_type = e.to.ty.clone();
609 let inner_fragment = e.expression.full_fragment_owned();
610 CompiledExpr::new(move |ctx| {
611 let column = inner.execute(ctx)?;
612 let frag = inner_fragment.clone();
613 let casted = cast_column_data(ctx, column.data(), target_type.clone(), &|| {
614 inner_fragment.clone()
615 })
616 .map_err(|e| {
617 Error::from(CastError::InvalidNumber {
618 fragment: frag,
619 target: target_type.clone(),
620 cause: e.diagnostic(),
621 })
622 })?;
623 Ok(ColumnWithName::new(label.clone(), casted))
624 })
625 }
626 }
627
628 Expression::If(e) => {
629 let condition = compile_expression(_ctx, &e.condition)?;
630 let then_expr = compile_expressions(_ctx, from_ref(e.then_expr.as_ref()))?;
631 let else_ifs: Vec<(CompiledExpr, Vec<CompiledExpr>)> = e
632 .else_ifs
633 .iter()
634 .map(|ei| {
635 Ok((
636 compile_expression(_ctx, &ei.condition)?,
637 compile_expressions(_ctx, from_ref(ei.then_expr.as_ref()))?,
638 ))
639 })
640 .collect::<Result<Vec<_>>>()?;
641 let else_branch: Option<Vec<CompiledExpr>> = match &e.else_expr {
642 Some(expr) => Some(compile_expressions(_ctx, from_ref(expr.as_ref()))?),
643 None => None,
644 };
645 let fragment = e.fragment.clone();
646 CompiledExpr::new_multi(move |ctx| {
647 execute_if_multi(ctx, &condition, &then_expr, &else_ifs, &else_branch, &fragment)
648 })
649 }
650
651 Expression::Map(e) => {
652 let expressions = compile_expressions(_ctx, &e.expressions)?;
653 CompiledExpr::new_multi(move |ctx| execute_projection_multi(ctx, &expressions))
654 }
655
656 Expression::Extend(e) => {
657 let expressions = compile_expressions(_ctx, &e.expressions)?;
658 CompiledExpr::new_multi(move |ctx| execute_projection_multi(ctx, &expressions))
659 }
660
661 Expression::Call(e) => {
662 let compiled_args: Vec<CompiledExpr> =
663 e.args.iter().map(|arg| compile_expression(_ctx, arg)).collect::<Result<Vec<_>>>()?;
664 let expr = e.clone();
665 CompiledExpr::new(move |ctx| {
666 let mut arg_columns = Vec::with_capacity(compiled_args.len());
667 for compiled_arg in &compiled_args {
668 arg_columns.push(compiled_arg.execute(ctx)?);
669 }
670 let arguments = Columns::new(arg_columns);
671 call_builtin(ctx, &expr, arguments)
672 })
673 }
674
675 Expression::SumTypeConstructor(_) => {
676 panic!(
677 "SumTypeConstructor in expression context — constructors should be expanded by InlineDataNode before expression compilation"
678 );
679 }
680
681 Expression::IsVariant(e) => {
682 let col_name = match e.expression.as_ref() {
683 Expression::Column(c) => c.0.name.text().to_string(),
684 other => display_label(other).text().to_string(),
685 };
686 let tag_col_name = format!("{}_tag", col_name);
687 let tag = e.tag.expect("IS variant tag must be resolved before compilation");
688 let fragment = e.fragment.clone();
689 CompiledExpr::new(move |ctx| {
690 if let Some(tag_col) =
691 ctx.columns.iter().find(|c| c.name().text() == tag_col_name.as_str())
692 {
693 match tag_col.data() {
694 ColumnBuffer::Uint1(container) => {
695 let results: Vec<bool> = container
696 .iter()
697 .take(ctx.row_count)
698 .map(|v| v == Some(tag))
699 .collect();
700 Ok(ColumnWithName::new(
701 fragment.clone(),
702 ColumnBuffer::bool(results),
703 ))
704 }
705 _ => Ok(ColumnWithName {
706 name: fragment.clone(),
707 data: ColumnBuffer::none_typed(Type::Boolean, ctx.row_count),
708 }),
709 }
710 } else {
711 Ok(ColumnWithName {
712 name: fragment.clone(),
713 data: ColumnBuffer::none_typed(Type::Boolean, ctx.row_count),
714 })
715 }
716 })
717 }
718
719 Expression::FieldAccess(e) => {
720 let field_name = e.field.text().to_string();
721 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)?;
814 Err(TypeError::Runtime {
815 kind: RuntimeErrorKind::FieldNotFound {
816 variable: "<expression>".to_string(),
817 field: field_name.to_string(),
818 available: vec![],
819 },
820 message: format!(
821 "Field '{}' not found on variable '<expression>'",
822 field_name
823 ),
824 }
825 .into())
826 }
827 })
828 }
829 })
830}
831
832fn compile_expressions(ctx: &CompileContext, exprs: &[Expression]) -> Result<Vec<CompiledExpr>> {
833 exprs.iter().map(|e| compile_expression(ctx, e)).collect()
834}
835
836fn combine_bool_columns(
837 left: ColumnWithName,
838 right: ColumnWithName,
839 fragment: Fragment,
840 combine_fn: fn(bool, bool) -> bool,
841) -> Result<ColumnWithName> {
842 binary_op_unwrap_option(&left, &right, fragment.clone(), |left, right| match (left.data(), right.data()) {
843 (ColumnBuffer::Bool(l), ColumnBuffer::Bool(r)) => {
844 let len = l.len();
845 let mut data = Vec::with_capacity(len);
846 let mut bitvec = Vec::with_capacity(len);
847
848 for i in 0..len {
849 let l_defined = l.is_defined(i);
850 let r_defined = r.is_defined(i);
851 let l_val = l.data().get(i);
852 let r_val = r.data().get(i);
853
854 if l_defined && r_defined {
855 data.push(combine_fn(l_val, r_val));
856 bitvec.push(true);
857 } else {
858 data.push(false);
859 bitvec.push(false);
860 }
861 }
862
863 Ok(ColumnWithName {
864 name: fragment.clone(),
865 data: ColumnBuffer::bool_with_bitvec(data, bitvec),
866 })
867 }
868 _ => {
869 unreachable!("combine_bool_columns should only be called with boolean columns")
870 }
871 })
872}
873
874fn list_items_contain(items: &[Value], element: &Value, fragment: &Fragment) -> bool {
875 if items.iter().any(|item| item == element) {
876 return true;
877 }
878 if items.is_empty() {
879 return false;
880 }
881
882 if let Some(items_buf) = build_homogeneous_buffer(items) {
883 let elems_buf = ColumnBuffer::from_many(element.clone(), items.len());
884 let items_col = ColumnWithName::new(fragment.clone(), items_buf);
885 let elems_col = ColumnWithName::new(fragment.clone(), elems_buf);
886 return compare_columns::<Equal>(&items_col, &elems_col, fragment.clone(), |f, l, r| {
887 TypeError::BinaryOperatorNotApplicable {
888 operator: BinaryOp::Equal,
889 left: l,
890 right: r,
891 fragment: f,
892 }
893 .into_diagnostic()
894 })
895 .map(|c| bool_column_has_true(&c))
896 .unwrap_or(false);
897 }
898
899 list_items_contain_per_item(items, element, fragment)
900}
901
902fn list_items_contain_per_item(items: &[Value], element: &Value, fragment: &Fragment) -> bool {
903 items.iter().any(|item| {
904 let item_col = ColumnWithName::new(fragment.clone(), ColumnBuffer::from(item.clone()));
905 let elem_col = ColumnWithName::new(fragment.clone(), ColumnBuffer::from(element.clone()));
906 compare_columns::<Equal>(&item_col, &elem_col, fragment.clone(), |f, l, r| {
907 TypeError::BinaryOperatorNotApplicable {
908 operator: BinaryOp::Equal,
909 left: l,
910 right: r,
911 fragment: f,
912 }
913 .into_diagnostic()
914 })
915 .ok()
916 .and_then(|c| match c.data() {
917 ColumnBuffer::Bool(b) => Some(b.data().get(0)),
918 _ => None,
919 })
920 .unwrap_or(false)
921 })
922}
923
924fn bool_column_has_true(col: &ColumnWithName) -> bool {
925 match col.data() {
926 ColumnBuffer::Bool(b) => b.data().any(),
927 ColumnBuffer::Option {
928 inner,
929 bitvec,
930 } => match inner.as_ref() {
931 ColumnBuffer::Bool(b) => {
932 let n = bitvec.len().min(b.len());
933 (0..n).any(|i| bitvec.get(i) && b.data().get(i))
934 }
935 _ => false,
936 },
937 _ => false,
938 }
939}
940
941fn build_homogeneous_buffer(items: &[Value]) -> Option<ColumnBuffer> {
942 let first = items.first()?;
943 let first_disc = discriminant(first);
944 if !items.iter().all(|v| discriminant(v) == first_disc) {
945 return None;
946 }
947
948 macro_rules! collect {
949 ($variant:ident, $constructor:ident, |$x:ident| $convert:expr) => {{
950 let data: Vec<_> = items
951 .iter()
952 .map(|v| match v {
953 Value::$variant($x) => $convert,
954 _ => unreachable!("homogeneous check guarantees variant"),
955 })
956 .collect();
957 Some(ColumnBuffer::$constructor(data))
958 }};
959 }
960
961 match first {
962 Value::Boolean(_) => collect!(Boolean, bool, |x| *x),
963 Value::Float4(_) => collect!(Float4, float4, |x| x.value()),
964 Value::Float8(_) => collect!(Float8, float8, |x| x.value()),
965 Value::Int1(_) => collect!(Int1, int1, |x| *x),
966 Value::Int2(_) => collect!(Int2, int2, |x| *x),
967 Value::Int4(_) => collect!(Int4, int4, |x| *x),
968 Value::Int8(_) => collect!(Int8, int8, |x| *x),
969 Value::Int16(_) => collect!(Int16, int16, |x| *x),
970 Value::Uint1(_) => collect!(Uint1, uint1, |x| *x),
971 Value::Uint2(_) => collect!(Uint2, uint2, |x| *x),
972 Value::Uint4(_) => collect!(Uint4, uint4, |x| *x),
973 Value::Uint8(_) => collect!(Uint8, uint8, |x| *x),
974 Value::Uint16(_) => collect!(Uint16, uint16, |x| *x),
975 Value::Utf8(_) => collect!(Utf8, utf8, |x| x.clone()),
976 Value::Date(_) => collect!(Date, date, |x| *x),
977 Value::DateTime(_) => collect!(DateTime, datetime, |x| *x),
978 Value::Time(_) => collect!(Time, time, |x| *x),
979 Value::Duration(_) => collect!(Duration, duration, |x| *x),
980 Value::Uuid4(_) => collect!(Uuid4, uuid4, |x| *x),
981 Value::Uuid7(_) => collect!(Uuid7, uuid7, |x| *x),
982 Value::IdentityId(_) => collect!(IdentityId, identity_id, |x| *x),
983 Value::Blob(_) => collect!(Blob, blob, |x| x.clone()),
984 Value::Int(_) => collect!(Int, int, |x| x.clone()),
985 Value::Uint(_) => collect!(Uint, uint, |x| x.clone()),
986 Value::Decimal(_) => collect!(Decimal, decimal, |x| x.clone()),
987 Value::DictionaryId(_) => collect!(DictionaryId, dictionary_id, |x| *x),
988 _ => None,
990 }
991}
992
993fn list_contains_element(
994 list_col: &ColumnWithName,
995 element_col: &ColumnWithName,
996 fragment: &Fragment,
997) -> Result<ColumnWithName> {
998 let len = list_col.data().len();
999 let mut data = Vec::with_capacity(len);
1000
1001 for i in 0..len {
1002 let list_value = list_col.data().get_value(i);
1003 let element_value = element_col.data().get_value(i);
1004
1005 let contained = match &list_value {
1006 Value::List(items) => list_items_contain(items, &element_value, fragment),
1007 Value::Tuple(items) => list_items_contain(items, &element_value, fragment),
1008 Value::Any(boxed) => match boxed.as_ref() {
1009 Value::List(items) => list_items_contain(items, &element_value, fragment),
1010 Value::Tuple(items) => list_items_contain(items, &element_value, fragment),
1011 _ => false,
1012 },
1013 _ => false,
1014 };
1015 data.push(contained);
1016 }
1017
1018 Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(data)))
1019}
1020
1021fn negate_column(col: ColumnWithName, fragment: Fragment) -> ColumnWithName {
1022 unary_op_unwrap_option(&col, |col| match col.data() {
1023 ColumnBuffer::Bool(container) => {
1024 let len = container.len();
1025 let mut data = Vec::with_capacity(len);
1026 let mut bitvec = Vec::with_capacity(len);
1027
1028 for i in 0..len {
1029 if container.is_defined(i) {
1030 data.push(!container.data().get(i));
1031 bitvec.push(true);
1032 } else {
1033 data.push(false);
1034 bitvec.push(false);
1035 }
1036 }
1037
1038 Ok(ColumnWithName {
1039 name: fragment.clone(),
1040 data: ColumnBuffer::bool_with_bitvec(data, bitvec),
1041 })
1042 }
1043 _ => unreachable!("negate_column should only be called with boolean columns"),
1044 })
1045 .unwrap()
1046}
1047
1048fn is_truthy(value: &Value) -> bool {
1049 match value {
1050 Value::Boolean(true) => true,
1051 Value::Boolean(false) => false,
1052 Value::None {
1053 ..
1054 } => false,
1055 Value::Int1(0) | Value::Int2(0) | Value::Int4(0) | Value::Int8(0) | Value::Int16(0) => false,
1056 Value::Uint1(0) | Value::Uint2(0) | Value::Uint4(0) | Value::Uint8(0) | Value::Uint16(0) => false,
1057 Value::Int1(_) | Value::Int2(_) | Value::Int4(_) | Value::Int8(_) | Value::Int16(_) => true,
1058 Value::Uint1(_) | Value::Uint2(_) | Value::Uint4(_) | Value::Uint8(_) | Value::Uint16(_) => true,
1059 Value::Utf8(s) => !s.is_empty(),
1060 _ => true,
1061 }
1062}
1063
1064fn execute_if_multi(
1065 ctx: &EvalContext,
1066 condition: &CompiledExpr,
1067 then_expr: &[CompiledExpr],
1068 else_ifs: &[(CompiledExpr, Vec<CompiledExpr>)],
1069 else_branch: &Option<Vec<CompiledExpr>>,
1070 _fragment: &Fragment,
1071) -> Result<Vec<ColumnWithName>> {
1072 let condition_column = condition.execute(ctx)?;
1073
1074 let mut result_data: Option<Vec<ColumnBuffer>> = None;
1075 let mut result_names: Vec<Fragment> = Vec::new();
1076
1077 for row_idx in 0..ctx.row_count {
1078 let condition_value = condition_column.data().get_value(row_idx);
1079
1080 let branch_results = if is_truthy(&condition_value) {
1081 execute_multi_exprs(ctx, then_expr)?
1082 } else {
1083 let mut found_branch = false;
1084 let mut branch_columns = None;
1085
1086 for (else_if_condition, else_if_then) in else_ifs {
1087 let else_if_col = else_if_condition.execute(ctx)?;
1088 let else_if_value = else_if_col.data().get_value(row_idx);
1089
1090 if is_truthy(&else_if_value) {
1091 branch_columns = Some(execute_multi_exprs(ctx, else_if_then)?);
1092 found_branch = true;
1093 break;
1094 }
1095 }
1096
1097 if found_branch {
1098 branch_columns.unwrap()
1099 } else if let Some(else_exprs) = else_branch {
1100 execute_multi_exprs(ctx, else_exprs)?
1101 } else {
1102 vec![]
1103 }
1104 };
1105
1106 let is_empty_result = branch_results.is_empty();
1108 if is_empty_result {
1109 if let Some(data) = result_data.as_mut() {
1110 for col_data in data.iter_mut() {
1111 col_data.push_value(Value::none());
1112 }
1113 }
1114 continue;
1115 }
1116
1117 if result_data.is_none() {
1119 let mut data: Vec<ColumnBuffer> = branch_results
1120 .iter()
1121 .map(|col| ColumnBuffer::with_capacity(col.data().get_type(), ctx.row_count))
1122 .collect();
1123 for _ in 0..row_idx {
1124 for col_data in data.iter_mut() {
1125 col_data.push_value(Value::none());
1126 }
1127 }
1128 result_data = Some(data);
1129 result_names = branch_results.iter().map(|col| col.name.clone()).collect();
1130 }
1131
1132 let data = result_data.as_mut().unwrap();
1133 for (i, branch_col) in branch_results.iter().enumerate() {
1134 if i < data.len() {
1135 let branch_value = branch_col.data().get_value(row_idx);
1136 data[i].push_value(branch_value);
1137 }
1138 }
1139 }
1140
1141 let result_data = result_data.unwrap_or_default();
1142 let result: Vec<ColumnWithName> = result_data
1143 .into_iter()
1144 .enumerate()
1145 .map(|(i, data)| ColumnWithName {
1146 name: result_names.get(i).cloned().unwrap_or_else(|| Fragment::internal("column")),
1147 data,
1148 })
1149 .collect();
1150
1151 if result.is_empty() {
1152 Ok(vec![ColumnWithName {
1153 name: Fragment::internal("none"),
1154 data: ColumnBuffer::none_typed(Type::Boolean, ctx.row_count),
1155 }])
1156 } else {
1157 Ok(result)
1158 }
1159}
1160
1161fn execute_multi_exprs(ctx: &EvalContext, exprs: &[CompiledExpr]) -> Result<Vec<ColumnWithName>> {
1162 let mut result = Vec::new();
1163 for expr in exprs {
1164 result.extend(expr.execute_multi(ctx)?);
1165 }
1166 Ok(result)
1167}
1168
1169fn execute_projection_multi(ctx: &EvalContext, expressions: &[CompiledExpr]) -> Result<Vec<ColumnWithName>> {
1170 let mut result = Vec::with_capacity(expressions.len());
1171
1172 for expr in expressions {
1173 let column = expr.execute(ctx)?;
1174 let name = column.name.text().to_string();
1175 result.push(ColumnWithName::new(Fragment::internal(name), column.data));
1176 }
1177
1178 Ok(result)
1179}