Skip to main content

reifydb_engine/vm/volcano/
variable.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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}