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