reifydb_engine/vm/volcano/
assert.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::sync::Arc;
5
6use arrow_array::{ArrayRef, RecordBatch};
7use arrow_schema::FieldRef;
8use reifydb_core::{
9	expression::{Expression, name::display_label},
10	value::column::headers::ColumnHeaders,
11};
12use reifydb_evaluate::expression::{context::EvalContext, eval::evaluate};
13use reifydb_transaction::transaction::Transaction;
14use reifydb_value::{
15	reifydb_assertions,
16	value::column_view::{ColumnView, ViewData},
17};
18use tracing::instrument;
19
20use crate::{
21	Result,
22	error::EngineError,
23	vm::volcano::query::{QueryContext, QueryNode, eval_context_from_query},
24};
25
26pub(crate) struct AssertNode {
27	input: Box<dyn QueryNode>,
28	expressions: Vec<Expression>,
29	message: Option<String>,
30	context: Option<Arc<QueryContext>>,
31}
32
33impl AssertNode {
34	pub fn new(input: Box<dyn QueryNode>, expressions: Vec<Expression>, message: Option<String>) -> Self {
35		Self {
36			input,
37			expressions,
38			message,
39			context: None,
40		}
41	}
42
43	#[instrument(level = "trace", skip_all, name = "volcano::assert::eval")]
44	fn eval(
45		session: &EvalContext<'_>,
46		columns: &RecordBatch,
47		row_count: usize,
48		assert_expr: &Expression,
49	) -> Result<(FieldRef, ArrayRef)> {
50		let eval_ctx = session.with_eval(columns.clone(), row_count);
51		evaluate(&eval_ctx, assert_expr)
52	}
53
54	#[instrument(level = "trace", skip_all, name = "volcano::assert::verify")]
55	fn verify(&self, data: &ColumnView<'_>, row_count: usize, assert_expr: &Expression) -> Result<()> {
56		let frag = assert_expr.full_fragment_owned();
57		let label = display_label(assert_expr);
58		match &data.data {
59			ViewData::Bool(container) => {
60				for i in 0..row_count {
61					let valid = data.is_defined(i);
62					let value = match data.logical_nulls() {
63						Some(_) => valid && container.value(i),
64						None => container.value(i),
65					};
66					if !valid || !value {
67						return Err(EngineError::AssertionFailed {
68							fragment: frag.clone(),
69							message: self.message.clone().unwrap_or_default(),
70							expression: Some(label.text().to_string()),
71						}
72						.into());
73					}
74				}
75			}
76			_ => {
77				return Err(EngineError::AssertionFailed {
78					fragment: frag.clone(),
79					message: "assert expression must evaluate to a boolean".to_string(),
80					expression: Some(label.text().to_string()),
81				}
82				.into());
83			}
84		}
85		Ok(())
86	}
87}
88
89impl QueryNode for AssertNode {
90	#[instrument(level = "trace", skip_all, name = "volcano::assert::initialize")]
91	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
92		self.context = Some(Arc::new(ctx.clone()));
93		self.input.initialize(rx, ctx)?;
94		Ok(())
95	}
96
97	#[instrument(level = "trace", skip_all, name = "volcano::assert::next")]
98	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
99		reifydb_assertions! {
100			assert!(self.context.is_some(), "AssertNode::next() called before initialize()");
101		}
102		let stored_ctx = self.context.as_ref().unwrap();
103
104		if let Some(columns) = self.input.next(rx, ctx)? {
105			let row_count = columns.num_rows();
106			let session = eval_context_from_query(stored_ctx);
107
108			for assert_expr in &self.expressions {
109				let result = Self::eval(&session, &columns, row_count, assert_expr)?;
110				self.verify(&ColumnView::try_from(&result)?, row_count, assert_expr)?;
111			}
112
113			Ok(Some(columns))
114		} else {
115			Ok(None)
116		}
117	}
118
119	fn headers(&self) -> Option<ColumnHeaders> {
120		self.input.headers()
121	}
122}
123
124pub(crate) struct AssertWithoutInputNode {
125	expressions: Vec<Expression>,
126	message: Option<String>,
127	context: Option<Arc<QueryContext>>,
128	done: bool,
129}
130
131impl AssertWithoutInputNode {
132	pub fn new(expressions: Vec<Expression>, message: Option<String>) -> Self {
133		Self {
134			expressions,
135			message,
136			context: None,
137			done: false,
138		}
139	}
140
141	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::eval")]
142	fn eval(session: &EvalContext<'_>, assert_expr: &Expression) -> Result<(FieldRef, ArrayRef)> {
143		let eval_ctx = session.with_eval_empty();
144		evaluate(&eval_ctx, assert_expr)
145	}
146
147	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::verify")]
148	fn verify(&self, data: &ColumnView<'_>, assert_expr: &Expression) -> Result<()> {
149		let frag = assert_expr.full_fragment_owned();
150		let label = display_label(assert_expr);
151		match &data.data {
152			ViewData::Bool(container) => {
153				let valid = data.is_defined(0);
154				let value = match data.logical_nulls() {
155					Some(_) => valid && container.value(0),
156					None => container.value(0),
157				};
158				if !valid || !value {
159					return Err(EngineError::AssertionFailed {
160						fragment: frag.clone(),
161						message: self.message.clone().unwrap_or_default(),
162						expression: Some(label.text().to_string()),
163					}
164					.into());
165				}
166			}
167			_ => {
168				return Err(EngineError::AssertionFailed {
169					fragment: frag.clone(),
170					message: "assert expression must evaluate to a boolean".to_string(),
171					expression: Some(label.text().to_string()),
172				}
173				.into());
174			}
175		}
176		Ok(())
177	}
178}
179
180impl QueryNode for AssertWithoutInputNode {
181	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::initialize")]
182	fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
183		self.context = Some(Arc::new(ctx.clone()));
184		Ok(())
185	}
186
187	#[instrument(level = "trace", skip_all, name = "volcano::assert::noinput::next")]
188	fn next<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
189		if self.done {
190			return Ok(None);
191		}
192		self.done = true;
193
194		reifydb_assertions! {
195			assert!(self.context.is_some(), "AssertWithoutInputNode::next() called before initialize()");
196		}
197		let stored_ctx = self.context.as_ref().unwrap();
198		let session = eval_context_from_query(stored_ctx);
199
200		for assert_expr in &self.expressions {
201			let result = Self::eval(&session, assert_expr)?;
202			self.verify(&ColumnView::try_from(&result)?, assert_expr)?;
203		}
204
205		Ok(None)
206	}
207
208	fn headers(&self) -> Option<ColumnHeaders> {
209		None
210	}
211}
reifydb_engine/vm/volcano/assert.rs

reifydb_engine/vm/volcano/
assert.rs
