reifydb_engine/vm/volcano/
variable.rs1use std::sync::Arc;
5
6use arrow_array::RecordBatch;
7use reifydb_core::{
8 expression::VariableExpression,
9 value::{
10 batch::{is_scalar, scalar_value, try_from_records},
11 column::headers::ColumnHeaders,
12 },
13};
14use reifydb_evaluate::{error::EvaluateError, stack::Variable};
15use reifydb_transaction::transaction::Transaction;
16use reifydb_value::{
17 error::{RuntimeErrorKind, TypeError},
18 reifydb_assertions,
19 value::Value,
20};
21use tracing::instrument;
22
23use crate::{
24 Result,
25 vm::volcano::query::{QueryContext, QueryNode},
26};
27
28pub(crate) struct VariableNode {
29 variable_expr: VariableExpression,
30 context: Option<Arc<QueryContext>>,
31 executed: bool,
32}
33
34impl VariableNode {
35 pub fn new(variable_expr: VariableExpression) -> Self {
36 Self {
37 variable_expr,
38 context: None,
39 executed: false,
40 }
41 }
42}
43
44impl QueryNode for VariableNode {
45 #[instrument(level = "trace", skip_all, name = "volcano::variable::initialize")]
46 fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
47 self.context = Some(Arc::new(ctx.clone()));
48 Ok(())
49 }
50
51 #[instrument(level = "trace", skip_all, name = "volcano::variable::next")]
52 fn next<'a>(&mut self, _rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
53 reifydb_assertions! {
54 assert!(self.context.is_some(), "VariableNode::next() called before initialize()");
55 }
56
57 if self.executed {
58 return Ok(None);
59 }
60
61 let variable_name = self.variable_expr.name();
62
63 match ctx.symbols.get(variable_name) {
64 Some(Variable::Columns {
65 batch,
66 }) if is_scalar(batch) => {
67 self.executed = true;
68 let mut value = scalar_value(batch)?;
69 while let Value::Any(inner) = value {
70 value = *inner;
71 }
72 match value {
73 Value::List(items) if items.iter().all(|v| matches!(v, Value::Record(_))) => {
74 Ok(Some(try_from_records(variable_name, &items)?))
75 }
76 _ => Ok(Some(batch.clone())),
77 }
78 }
79 Some(Variable::Columns {
80 batch,
81 }) => {
82 self.executed = true;
83 Ok(Some(batch.clone()))
84 }
85 Some(Variable::ForIterator {
86 batch,
87 ..
88 }) => {
89 self.executed = true;
90
91 Ok(Some(batch.clone()))
92 }
93 Some(Variable::Closure(_)) => Err(TypeError::Runtime {
94 kind: RuntimeErrorKind::VariableIsClosure {
95 fragment: self.variable_expr.fragment.clone(),
96 },
97 message: format!(
98 "Variable '{}' holds a closure and cannot be read as rows",
99 variable_name
100 ),
101 }
102 .into()),
103 None => Err(EvaluateError::VariableNotFound {
104 name: variable_name.to_string(),
105 }
106 .into()),
107 }
108 }
109
110 fn headers(&self) -> Option<ColumnHeaders> {
111 None
112 }
113}