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 vm::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<Box<Value>> =
339 (0..row_count).map(|_| Box::new(Value::Type(ty.clone()))).collect();
340 Ok(ColumnWithName::new(fragment.text(), ColumnBuffer::any(values)))
341 })
342 }
343
344 Expression::AccessSource(e) => {
345 let col_name = e.column.name.text().to_string();
346 let expr = e.clone();
347 CompiledExpr::new_access(col_name, move |ctx| access_lookup(ctx, &expr))
348 }
349
350 Expression::Tuple(e) => {
351 if e.expressions.len() == 1 {
352 let inner = compile_expression(_ctx, &e.expressions[0])?;
353 CompiledExpr::new(move |ctx| inner.execute(ctx))
354 } else {
355 let compiled: Vec<CompiledExpr> = e
356 .expressions
357 .iter()
358 .map(|expr| compile_expression(_ctx, expr))
359 .collect::<Result<Vec<_>>>()?;
360 let fragment = e.fragment.clone();
361 CompiledExpr::new(move |ctx| {
362 let columns: Vec<ColumnWithName> = compiled
363 .iter()
364 .map(|expr| expr.execute(ctx))
365 .collect::<Result<Vec<_>>>()?;
366
367 let len = columns.first().map_or(1, |c| c.data().len());
368 let mut data: Vec<Box<Value>> = Vec::with_capacity(len);
369
370 for i in 0..len {
371 let items: Vec<Value> =
372 columns.iter().map(|col| col.data().get_value(i)).collect();
373 data.push(Box::new(Value::Tuple(items)));
374 }
375
376 Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::any(data)))
377 })
378 }
379 }
380
381 Expression::List(e) => {
382 let compiled: Vec<CompiledExpr> = e
383 .expressions
384 .iter()
385 .map(|expr| compile_expression(_ctx, expr))
386 .collect::<Result<Vec<_>>>()?;
387 let fragment = e.fragment.clone();
388 CompiledExpr::new(move |ctx| {
389 let columns: Vec<ColumnWithName> =
390 compiled.iter().map(|expr| expr.execute(ctx)).collect::<Result<Vec<_>>>()?;
391
392 let len = columns.first().map_or(1, |c| c.data().len());
393 let mut data: Vec<Box<Value>> = Vec::with_capacity(len);
394
395 for i in 0..len {
396 let items: Vec<Value> =
397 columns.iter().map(|col| col.data().get_value(i)).collect();
398 data.push(Box::new(Value::List(items)));
399 }
400
401 Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::any(data)))
402 })
403 }
404
405 Expression::Between(e) => {
406 let value = compile_expression(_ctx, &e.value)?;
407 let lower = compile_expression(_ctx, &e.lower)?;
408 let upper = compile_expression(_ctx, &e.upper)?;
409 let fragment = e.fragment.clone();
410 CompiledExpr::new(move |ctx| {
411 let value_col = value.execute(ctx)?;
412 let lower_col = lower.execute(ctx)?;
413 let upper_col = upper.execute(ctx)?;
414
415 let ge_result = compare_columns::<GreaterThanEqual>(
416 &value_col,
417 &lower_col,
418 fragment.clone(),
419 |f, l, r| {
420 TypeError::BinaryOperatorNotApplicable {
421 operator: BinaryOp::Between,
422 left: l,
423 right: r,
424 fragment: f,
425 }
426 .into_diagnostic()
427 },
428 )?;
429 let le_result = compare_columns::<LessThanEqual>(
430 &value_col,
431 &upper_col,
432 fragment.clone(),
433 |f, l, r| {
434 TypeError::BinaryOperatorNotApplicable {
435 operator: BinaryOp::Between,
436 left: l,
437 right: r,
438 fragment: f,
439 }
440 .into_diagnostic()
441 },
442 )?;
443
444 if !matches!(ge_result.data(), ColumnBuffer::Bool(_))
445 || !matches!(le_result.data(), ColumnBuffer::Bool(_))
446 {
447 return Err(TypeError::BinaryOperatorNotApplicable {
448 operator: BinaryOp::Between,
449 left: value_col.get_type(),
450 right: lower_col.get_type(),
451 fragment: fragment.clone(),
452 }
453 .into());
454 }
455
456 match (ge_result.data(), le_result.data()) {
457 (ColumnBuffer::Bool(ge_container), ColumnBuffer::Bool(le_container)) => {
458 let mut data = Vec::with_capacity(ge_container.len());
459 let mut bitvec = Vec::with_capacity(ge_container.len());
460
461 for i in 0..ge_container.len() {
462 if ge_container.is_defined(i) && le_container.is_defined(i) {
463 data.push(ge_container.data().get(i)
464 && le_container.data().get(i));
465 bitvec.push(true);
466 } else {
467 data.push(false);
468 bitvec.push(false);
469 }
470 }
471
472 Ok(ColumnWithName {
473 name: fragment.clone(),
474 data: ColumnBuffer::bool_with_bitvec(data, bitvec),
475 })
476 }
477 _ => unreachable!(
478 "Both comparison results should be boolean after the check above"
479 ),
480 }
481 })
482 }
483
484 Expression::In(e) => {
485 let list_expressions = match e.list.as_ref() {
486 Expression::Tuple(tuple) => &tuple.expressions,
487 Expression::List(list) => &list.expressions,
488 _ => from_ref(e.list.as_ref()),
489 };
490 let value = compile_expression(_ctx, &e.value)?;
491 let list: Vec<CompiledExpr> = list_expressions
492 .iter()
493 .map(|expr| compile_expression(_ctx, expr))
494 .collect::<Result<Vec<_>>>()?;
495 let negated = e.negated;
496 let fragment = e.fragment.clone();
497 CompiledExpr::new(move |ctx| {
498 if list.is_empty() {
499 let value_col = value.execute(ctx)?;
500 let len = value_col.data().len();
501 let result = vec![negated; len];
502 return Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(result)));
503 }
504
505 let value_col = value.execute(ctx)?;
506
507 let first_col = list[0].execute(ctx)?;
508 let mut result = compare_columns::<Equal>(
509 &value_col,
510 &first_col,
511 fragment.clone(),
512 |f, l, r| {
513 TypeError::BinaryOperatorNotApplicable {
514 operator: BinaryOp::Equal,
515 left: l,
516 right: r,
517 fragment: f,
518 }
519 .into_diagnostic()
520 },
521 )?;
522
523 for list_expr in list.iter().skip(1) {
524 let list_col = list_expr.execute(ctx)?;
525 let eq_result = compare_columns::<Equal>(
526 &value_col,
527 &list_col,
528 fragment.clone(),
529 |f, l, r| {
530 TypeError::BinaryOperatorNotApplicable {
531 operator: BinaryOp::Equal,
532 left: l,
533 right: r,
534 fragment: f,
535 }
536 .into_diagnostic()
537 },
538 )?;
539 result = combine_bool_columns(result, eq_result, fragment.clone(), |l, r| {
540 l || r
541 })?;
542 }
543
544 if negated {
545 result = negate_column(result, fragment.clone());
546 }
547
548 Ok(result)
549 })
550 }
551
552 Expression::Contains(e) => {
553 let list_expressions = match e.list.as_ref() {
554 Expression::Tuple(tuple) => &tuple.expressions,
555 Expression::List(list) => &list.expressions,
556 _ => from_ref(e.list.as_ref()),
557 };
558 let value = compile_expression(_ctx, &e.value)?;
559 let list: Vec<CompiledExpr> = list_expressions
560 .iter()
561 .map(|expr| compile_expression(_ctx, expr))
562 .collect::<Result<Vec<_>>>()?;
563 let fragment = e.fragment.clone();
564 CompiledExpr::new(move |ctx| {
565 let value_col = value.execute(ctx)?;
566
567 if list.is_empty() {
568 let len = value_col.data().len();
569 let result = vec![true; len];
570 return Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(result)));
571 }
572
573 let first_col = list[0].execute(ctx)?;
574 let mut result = list_contains_element(&value_col, &first_col, &fragment)?;
575
576 for list_expr in list.iter().skip(1) {
577 let list_col = list_expr.execute(ctx)?;
578 let element_result = list_contains_element(&value_col, &list_col, &fragment)?;
579 result = combine_bool_columns(
580 result,
581 element_result,
582 fragment.clone(),
583 |l, r| l && r,
584 )?;
585 }
586
587 Ok(result)
588 })
589 }
590
591 Expression::Cast(e) => {
592 let label = display_label(expr);
593 if let Expression::Constant(const_expr) = e.expression.as_ref() {
594 let const_expr = const_expr.clone();
595 let target_type = e.to.ty.clone();
596 let inner_fragment = e.expression.full_fragment_owned();
597 CompiledExpr::new(move |ctx| {
598 let row_count = ctx.take.unwrap_or(ctx.row_count);
599 let data = constant_value(&const_expr, row_count)?;
600 let casted = if data.get_type() == target_type {
601 data
602 } else {
603 apply_cast(ctx, &data, &target_type, &inner_fragment)?
604 };
605 Ok(ColumnWithName::new(label.clone(), casted))
606 })
607 } else {
608 let inner = compile_expression(_ctx, &e.expression)?;
609 let target_type = e.to.ty.clone();
610 let inner_fragment = e.expression.full_fragment_owned();
611 CompiledExpr::new(move |ctx| {
612 let column = inner.execute(ctx)?;
613 let casted = apply_cast(ctx, column.data(), &target_type, &inner_fragment)?;
614 Ok(ColumnWithName::new(label.clone(), casted))
615 })
616 }
617 }
618
619 Expression::If(e) => {
620 let condition = compile_expression(_ctx, &e.condition)?;
621 let then_expr = compile_expressions(_ctx, from_ref(e.then_expr.as_ref()))?;
622 let else_ifs: Vec<(CompiledExpr, Vec<CompiledExpr>)> = e
623 .else_ifs
624 .iter()
625 .map(|ei| {
626 Ok((
627 compile_expression(_ctx, &ei.condition)?,
628 compile_expressions(_ctx, from_ref(ei.then_expr.as_ref()))?,
629 ))
630 })
631 .collect::<Result<Vec<_>>>()?;
632 let else_branch: Option<Vec<CompiledExpr>> = match &e.else_expr {
633 Some(expr) => Some(compile_expressions(_ctx, from_ref(expr.as_ref()))?),
634 None => None,
635 };
636 let fragment = e.fragment.clone();
637 CompiledExpr::new_multi(move |ctx| {
638 execute_if_multi(ctx, &condition, &then_expr, &else_ifs, &else_branch, &fragment)
639 })
640 }
641
642 Expression::Map(e) => {
643 let expressions = compile_expressions(_ctx, &e.expressions)?;
644 CompiledExpr::new_multi(move |ctx| execute_projection_multi(ctx, &expressions))
645 }
646
647 Expression::Extend(e) => {
648 let expressions = compile_expressions(_ctx, &e.expressions)?;
649 CompiledExpr::new_multi(move |ctx| execute_projection_multi(ctx, &expressions))
650 }
651
652 Expression::Call(e) => {
653 let compiled_args: Vec<CompiledExpr> =
654 e.args.iter().map(|arg| compile_expression(_ctx, arg)).collect::<Result<Vec<_>>>()?;
655 let expr = e.clone();
656 CompiledExpr::new(move |ctx| {
657 let mut arg_columns = Vec::with_capacity(compiled_args.len());
658 for compiled_arg in &compiled_args {
659 arg_columns.push(compiled_arg.execute(ctx)?);
660 }
661 let arguments = Columns::new(arg_columns);
662 call_builtin(ctx, &expr, arguments)
663 })
664 }
665
666 Expression::SumTypeConstructor(_) => {
667 panic!(
668 "SumTypeConstructor in expression context - constructors should be expanded by InlineDataNode before expression compilation"
669 );
670 }
671
672 Expression::IsVariant(e) => {
673 let col_name = match e.expression.as_ref() {
674 Expression::Column(c) => c.0.name.text().to_string(),
675 other => display_label(other).text().to_string(),
676 };
677 let tag_col_name = format!("{}_tag", col_name);
678 let tag = e.tag.expect("IS variant tag must be resolved before compilation");
679 let fragment = e.fragment.clone();
680 CompiledExpr::new(move |ctx| {
681 if let Some(tag_col) =
682 ctx.columns.iter().find(|c| c.name().text() == tag_col_name.as_str())
683 {
684 match tag_col.data() {
685 ColumnBuffer::Uint1(container) => {
686 let results: Vec<bool> = container
687 .iter()
688 .take(ctx.row_count)
689 .map(|v| v == Some(tag))
690 .collect();
691 Ok(ColumnWithName::new(
692 fragment.clone(),
693 ColumnBuffer::bool(results),
694 ))
695 }
696 _ => Ok(ColumnWithName {
697 name: fragment.clone(),
698 data: ColumnBuffer::none_typed(
699 ValueType::Boolean,
700 ctx.row_count,
701 ),
702 }),
703 }
704 } else {
705 Ok(ColumnWithName {
706 name: fragment.clone(),
707 data: ColumnBuffer::none_typed(ValueType::Boolean, ctx.row_count),
708 })
709 }
710 })
711 }
712
713 Expression::FieldAccess(e) => {
714 let field_name = e.field.text().to_string();
715
716 let var_name = match e.object.as_ref() {
717 Expression::Variable(var_expr) => Some(var_expr.name().to_string()),
718 _ => None,
719 };
720 let object = compile_expression(_ctx, &e.object)?;
721 CompiledExpr::new(move |ctx| {
722 if let Some(ref variable_name) = var_name {
723 match ctx.symbols.get(variable_name) {
724 Some(Variable::Columns {
725 columns,
726 }) if !columns.is_scalar() => {
727 let col_pos = columns
728 .names
729 .iter()
730 .position(|n| n.text() == field_name);
731 match col_pos {
732 Some(pos) => {
733 let value = columns.columns[pos].get_value(0);
734 let row_count =
735 ctx.take.unwrap_or(ctx.row_count);
736 let mut data = ColumnBuffer::with_capacity(
737 value.get_type(),
738 row_count,
739 );
740 for _ in 0..row_count {
741 data.push_value(value.clone());
742 }
743 Ok(ColumnWithName {
744 name: Fragment::internal(&field_name),
745 data,
746 })
747 }
748 None => {
749 let available: Vec<String> = columns
750 .names
751 .iter()
752 .map(|n| n.text().to_string())
753 .collect();
754 Err(TypeError::Runtime {
755 kind: RuntimeErrorKind::FieldNotFound {
756 variable: variable_name
757 .to_string(),
758 field: field_name.to_string(),
759 available,
760 },
761 message: format!(
762 "Field '{}' not found on variable '{}'",
763 field_name, variable_name
764 ),
765 }
766 .into())
767 }
768 }
769 }
770 Some(Variable::Columns {
771 ..
772 })
773 | Some(Variable::Closure(_)) => Err(TypeError::Runtime {
774 kind: RuntimeErrorKind::FieldNotFound {
775 variable: variable_name.to_string(),
776 field: field_name.to_string(),
777 available: vec![],
778 },
779 message: format!(
780 "Field '{}' not found on variable '{}'",
781 field_name, variable_name
782 ),
783 }
784 .into()),
785 Some(Variable::ForIterator {
786 ..
787 }) => Err(TypeError::Runtime {
788 kind: RuntimeErrorKind::VariableIsDataframe {
789 name: variable_name.to_string(),
790 },
791 message: format!(
792 "Variable '{}' contains a dataframe and cannot be used directly in scalar expressions",
793 variable_name
794 ),
795 }
796 .into()),
797 None => Err(TypeError::Runtime {
798 kind: RuntimeErrorKind::VariableNotFound {
799 name: variable_name.to_string(),
800 },
801 message: format!("Variable '{}' is not defined", variable_name),
802 }
803 .into()),
804 }
805 } else {
806 let _obj_col = object.execute(ctx)?;
807 Err(TypeError::Runtime {
808 kind: RuntimeErrorKind::FieldNotFound {
809 variable: "<expression>".to_string(),
810 field: field_name.to_string(),
811 available: vec![],
812 },
813 message: format!(
814 "Field '{}' not found on variable '<expression>'",
815 field_name
816 ),
817 }
818 .into())
819 }
820 })
821 }
822 })
823}
824
825fn compile_expressions(ctx: &CompileContext, exprs: &[Expression]) -> Result<Vec<CompiledExpr>> {
826 exprs.iter().map(|e| compile_expression(ctx, e)).collect()
827}
828
829fn combine_bool_columns(
830 left: ColumnWithName,
831 right: ColumnWithName,
832 fragment: Fragment,
833 combine_fn: fn(bool, bool) -> bool,
834) -> Result<ColumnWithName> {
835 binary_op_unwrap_option(&left, &right, fragment.clone(), |left, right| match (left.data(), right.data()) {
836 (ColumnBuffer::Bool(l), ColumnBuffer::Bool(r)) => {
837 let len = l.len();
838 let mut data = Vec::with_capacity(len);
839 let mut bitvec = Vec::with_capacity(len);
840
841 for i in 0..len {
842 let l_defined = l.is_defined(i);
843 let r_defined = r.is_defined(i);
844 let l_val = l.data().get(i);
845 let r_val = r.data().get(i);
846
847 if l_defined && r_defined {
848 data.push(combine_fn(l_val, r_val));
849 bitvec.push(true);
850 } else {
851 data.push(false);
852 bitvec.push(false);
853 }
854 }
855
856 Ok(ColumnWithName {
857 name: fragment.clone(),
858 data: ColumnBuffer::bool_with_bitvec(data, bitvec),
859 })
860 }
861 _ => {
862 unreachable!("combine_bool_columns should only be called with boolean columns")
863 }
864 })
865}
866
867fn list_items_contain(items: &[Value], element: &Value, fragment: &Fragment) -> bool {
868 if items.iter().any(|item| item == element) {
869 return true;
870 }
871 if items.is_empty() {
872 return false;
873 }
874
875 if let Some(items_buf) = build_homogeneous_buffer(items) {
876 let elems_buf = ColumnBuffer::from_many(element.clone(), items.len());
877 let items_col = ColumnWithName::new(fragment.clone(), items_buf);
878 let elems_col = ColumnWithName::new(fragment.clone(), elems_buf);
879 return compare_columns::<Equal>(&items_col, &elems_col, fragment.clone(), |f, l, r| {
880 TypeError::BinaryOperatorNotApplicable {
881 operator: BinaryOp::Equal,
882 left: l,
883 right: r,
884 fragment: f,
885 }
886 .into_diagnostic()
887 })
888 .map(|c| bool_column_has_true(&c))
889 .unwrap_or(false);
890 }
891
892 list_items_contain_per_item(items, element, fragment)
893}
894
895fn list_items_contain_per_item(items: &[Value], element: &Value, fragment: &Fragment) -> bool {
896 items.iter().any(|item| {
897 let item_col = ColumnWithName::new(fragment.clone(), ColumnBuffer::from(item.clone()));
898 let elem_col = ColumnWithName::new(fragment.clone(), ColumnBuffer::from(element.clone()));
899 compare_columns::<Equal>(&item_col, &elem_col, fragment.clone(), |f, l, r| {
900 TypeError::BinaryOperatorNotApplicable {
901 operator: BinaryOp::Equal,
902 left: l,
903 right: r,
904 fragment: f,
905 }
906 .into_diagnostic()
907 })
908 .ok()
909 .and_then(|c| match c.data() {
910 ColumnBuffer::Bool(b) => Some(b.data().get(0)),
911 _ => None,
912 })
913 .unwrap_or(false)
914 })
915}
916
917fn bool_column_has_true(col: &ColumnWithName) -> bool {
918 match col.data() {
919 ColumnBuffer::Bool(b) => b.data().any(),
920 ColumnBuffer::Option {
921 inner,
922 bitvec,
923 } => match inner.as_ref() {
924 ColumnBuffer::Bool(b) => {
925 let n = bitvec.len().min(b.len());
926 (0..n).any(|i| bitvec.get(i) && b.data().get(i))
927 }
928 _ => false,
929 },
930 _ => false,
931 }
932}
933
934fn build_homogeneous_buffer(items: &[Value]) -> Option<ColumnBuffer> {
935 let first = items.first()?;
936 let first_disc = discriminant(first);
937 if !items.iter().all(|v| discriminant(v) == first_disc) {
938 return None;
939 }
940
941 macro_rules! collect {
942 ($variant:ident, $constructor:ident, |$x:ident| $convert:expr) => {{
943 let data: Vec<_> = items
944 .iter()
945 .map(|v| match v {
946 Value::$variant($x) => $convert,
947 _ => unreachable!("homogeneous check guarantees variant"),
948 })
949 .collect();
950 Some(ColumnBuffer::$constructor(data))
951 }};
952 }
953
954 match first {
955 Value::Boolean(_) => collect!(Boolean, bool, |x| *x),
956 Value::Float4(_) => collect!(Float4, float4, |x| x.value()),
957 Value::Float8(_) => collect!(Float8, float8, |x| x.value()),
958 Value::Int1(_) => collect!(Int1, int1, |x| *x),
959 Value::Int2(_) => collect!(Int2, int2, |x| *x),
960 Value::Int4(_) => collect!(Int4, int4, |x| *x),
961 Value::Int8(_) => collect!(Int8, int8, |x| *x),
962 Value::Int16(_) => collect!(Int16, int16, |x| *x),
963 Value::Uint1(_) => collect!(Uint1, uint1, |x| *x),
964 Value::Uint2(_) => collect!(Uint2, uint2, |x| *x),
965 Value::Uint4(_) => collect!(Uint4, uint4, |x| *x),
966 Value::Uint8(_) => collect!(Uint8, uint8, |x| *x),
967 Value::Uint16(_) => collect!(Uint16, uint16, |x| *x),
968 Value::Utf8(_) => collect!(Utf8, utf8, |x| x.clone()),
969 Value::Date(_) => collect!(Date, date, |x| *x),
970 Value::DateTime(_) => collect!(DateTime, datetime, |x| *x),
971 Value::Time(_) => collect!(Time, time, |x| *x),
972 Value::Duration(_) => collect!(Duration, duration, |x| *x),
973 Value::Uuid4(_) => collect!(Uuid4, uuid4, |x| *x),
974 Value::Uuid7(_) => collect!(Uuid7, uuid7, |x| *x),
975 Value::IdentityId(_) => collect!(IdentityId, identity_id, |x| *x),
976 Value::Blob(_) => collect!(Blob, blob, |x| x.clone()),
977 Value::Int(_) => collect!(Int, int, |x| x.clone()),
978 Value::Uint(_) => collect!(Uint, uint, |x| x.clone()),
979 Value::Decimal(_) => collect!(Decimal, decimal, |x| x.clone()),
980 Value::DictionaryId(_) => collect!(DictionaryId, dictionary_id, |x| *x),
981
982 _ => None,
983 }
984}
985
986fn list_contains_element(
987 list_col: &ColumnWithName,
988 element_col: &ColumnWithName,
989 fragment: &Fragment,
990) -> Result<ColumnWithName> {
991 let len = list_col.data().len();
992 let mut data = Vec::with_capacity(len);
993
994 for i in 0..len {
995 let list_value = list_col.data().get_value(i);
996 let element_value = element_col.data().get_value(i);
997
998 let contained = match &list_value {
999 Value::List(items) => list_items_contain(items, &element_value, fragment),
1000 Value::Tuple(items) => list_items_contain(items, &element_value, fragment),
1001 Value::Any(boxed) => match boxed.as_ref() {
1002 Value::List(items) => list_items_contain(items, &element_value, fragment),
1003 Value::Tuple(items) => list_items_contain(items, &element_value, fragment),
1004 _ => false,
1005 },
1006 _ => false,
1007 };
1008 data.push(contained);
1009 }
1010
1011 Ok(ColumnWithName::new(fragment.clone(), ColumnBuffer::bool(data)))
1012}
1013
1014fn negate_column(col: ColumnWithName, fragment: Fragment) -> ColumnWithName {
1015 unary_op_unwrap_option(&col, |col| match col.data() {
1016 ColumnBuffer::Bool(container) => {
1017 let len = container.len();
1018 let mut data = Vec::with_capacity(len);
1019 let mut bitvec = Vec::with_capacity(len);
1020
1021 for i in 0..len {
1022 if container.is_defined(i) {
1023 data.push(!container.data().get(i));
1024 bitvec.push(true);
1025 } else {
1026 data.push(false);
1027 bitvec.push(false);
1028 }
1029 }
1030
1031 Ok(ColumnWithName {
1032 name: fragment.clone(),
1033 data: ColumnBuffer::bool_with_bitvec(data, bitvec),
1034 })
1035 }
1036 _ => unreachable!("negate_column should only be called with boolean columns"),
1037 })
1038 .unwrap()
1039}
1040
1041fn is_truthy(value: &Value) -> bool {
1042 match value {
1043 Value::Boolean(true) => true,
1044 Value::Boolean(false) => false,
1045 Value::None {
1046 ..
1047 } => false,
1048 Value::Int1(0) | Value::Int2(0) | Value::Int4(0) | Value::Int8(0) | Value::Int16(0) => false,
1049 Value::Uint1(0) | Value::Uint2(0) | Value::Uint4(0) | Value::Uint8(0) | Value::Uint16(0) => false,
1050 Value::Int1(_) | Value::Int2(_) | Value::Int4(_) | Value::Int8(_) | Value::Int16(_) => true,
1051 Value::Uint1(_) | Value::Uint2(_) | Value::Uint4(_) | Value::Uint8(_) | Value::Uint16(_) => true,
1052 Value::Utf8(s) => !s.is_empty(),
1053 _ => true,
1054 }
1055}
1056
1057fn execute_if_multi(
1058 ctx: &EvalContext,
1059 condition: &CompiledExpr,
1060 then_expr: &[CompiledExpr],
1061 else_ifs: &[(CompiledExpr, Vec<CompiledExpr>)],
1062 else_branch: &Option<Vec<CompiledExpr>>,
1063 _fragment: &Fragment,
1064) -> Result<Vec<ColumnWithName>> {
1065 let condition_column = condition.execute(ctx)?;
1066
1067 let mut result_data: Option<Vec<ColumnBuffer>> = None;
1068 let mut result_names: Vec<Fragment> = Vec::new();
1069
1070 for row_idx in 0..ctx.row_count {
1071 let condition_value = condition_column.data().get_value(row_idx);
1072
1073 let branch_results = if is_truthy(&condition_value) {
1074 execute_multi_exprs(ctx, then_expr)?
1075 } else {
1076 let mut found_branch = false;
1077 let mut branch_columns = None;
1078
1079 for (else_if_condition, else_if_then) in else_ifs {
1080 let else_if_col = else_if_condition.execute(ctx)?;
1081 let else_if_value = else_if_col.data().get_value(row_idx);
1082
1083 if is_truthy(&else_if_value) {
1084 branch_columns = Some(execute_multi_exprs(ctx, else_if_then)?);
1085 found_branch = true;
1086 break;
1087 }
1088 }
1089
1090 if found_branch {
1091 branch_columns.unwrap()
1092 } else if let Some(else_exprs) = else_branch {
1093 execute_multi_exprs(ctx, else_exprs)?
1094 } else {
1095 vec![]
1096 }
1097 };
1098
1099 let is_empty_result = branch_results.is_empty();
1100 if is_empty_result {
1101 if let Some(data) = result_data.as_mut() {
1102 for col_data in data.iter_mut() {
1103 col_data.push_value(Value::none());
1104 }
1105 }
1106 continue;
1107 }
1108
1109 if result_data.is_none() {
1110 let mut data: Vec<ColumnBuffer> = branch_results
1111 .iter()
1112 .map(|col| ColumnBuffer::with_capacity(col.data().get_type(), ctx.row_count))
1113 .collect();
1114 for _ in 0..row_idx {
1115 for col_data in data.iter_mut() {
1116 col_data.push_value(Value::none());
1117 }
1118 }
1119 result_data = Some(data);
1120 result_names = branch_results.iter().map(|col| col.name.clone()).collect();
1121 }
1122
1123 let data = result_data.as_mut().unwrap();
1124 for (i, branch_col) in branch_results.iter().enumerate() {
1125 if i < data.len() {
1126 let branch_value = branch_col.data().get_value(row_idx);
1127 data[i].push_value(branch_value);
1128 }
1129 }
1130 }
1131
1132 let result_data = result_data.unwrap_or_default();
1133 let result: Vec<ColumnWithName> = result_data
1134 .into_iter()
1135 .enumerate()
1136 .map(|(i, data)| ColumnWithName {
1137 name: result_names.get(i).cloned().unwrap_or_else(|| Fragment::internal("column")),
1138 data,
1139 })
1140 .collect();
1141
1142 if result.is_empty() {
1143 Ok(vec![ColumnWithName {
1144 name: Fragment::internal("none"),
1145 data: ColumnBuffer::none_typed(ValueType::Boolean, ctx.row_count),
1146 }])
1147 } else {
1148 Ok(result)
1149 }
1150}
1151
1152fn execute_multi_exprs(ctx: &EvalContext, exprs: &[CompiledExpr]) -> Result<Vec<ColumnWithName>> {
1153 let mut result = Vec::new();
1154 for expr in exprs {
1155 result.extend(expr.execute_multi(ctx)?);
1156 }
1157 Ok(result)
1158}
1159
1160fn execute_projection_multi(ctx: &EvalContext, expressions: &[CompiledExpr]) -> Result<Vec<ColumnWithName>> {
1161 let mut result = Vec::with_capacity(expressions.len());
1162
1163 for expr in expressions {
1164 let column = expr.execute(ctx)?;
1165 let name = column.name.text().to_string();
1166 result.push(ColumnWithName::new(Fragment::internal(name), column.data));
1167 }
1168
1169 Ok(result)
1170}
1171
1172fn apply_cast(ctx: &EvalContext, data: &ColumnBuffer, target: &ValueType, fragment: &Fragment) -> Result<ColumnBuffer> {
1173 cast_column_data(ctx, data, target.clone(), &|| fragment.clone())
1174 .map_err(|e| wrap_cast_error(e, fragment.clone(), target))
1175}
1176
1177fn wrap_cast_error(err: Error, fragment: Fragment, target: &ValueType) -> Error {
1178 if err.0.code.starts_with("CAST_") {
1179 return err;
1180 }
1181 let cause = err.diagnostic();
1182 let wrapped = if target.is_bool() {
1183 CastError::InvalidBoolean {
1184 fragment,
1185 cause,
1186 }
1187 } else if target.is_temporal() {
1188 CastError::InvalidTemporal {
1189 fragment,
1190 target: target.clone(),
1191 cause,
1192 }
1193 } else if target.is_uuid() || *target == ValueType::IdentityId {
1194 CastError::InvalidUuid {
1195 fragment,
1196 target: target.clone(),
1197 cause,
1198 }
1199 } else {
1200 CastError::InvalidNumber {
1201 fragment,
1202 target: target.clone(),
1203 cause,
1204 }
1205 };
1206 Error::from(wrapped)
1207}