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