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