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