reifydb_engine/vm/volcano/
query.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	error::diagnostic::{operation, query},
9	interface::{
10		catalog::config::{ConfigKey, GetConfig},
11		resolved::ResolvedObject,
12	},
13	sort::SortKey,
14	util::budget::MemoryBudget,
15	value::{batch::empty_batch, column::headers::ColumnHeaders},
16};
17use reifydb_evaluate::{expression::context::EvalContext, stack::SymbolTable};
18use reifydb_extension::transform::context::TransformContext;
19use reifydb_transaction::transaction::Transaction;
20use reifydb_value::{
21	byte_size::ByteSize,
22	error,
23	params::Params,
24	reifydb_assertions,
25	value::{column_view::ColumnView, identity::IdentityId, system_columns::check_user_columns},
26};
27
28use crate::{Result, vm::services::Services};
29
30pub fn query_budget(services: &Services) -> Arc<MemoryBudget> {
31	let limit = services.catalog.get_config_uint8(ConfigKey::QueryMemoryLimit);
32	Arc::new(MemoryBudget::new(ByteSize::from_bytes(limit)))
33}
34
35pub(crate) fn ensure_sort_key_orderable(key: &SortKey, data: &ColumnView<'_>) -> Result<()> {
36	let ty = data.get_type();
37	if ty.is_scalar() {
38		Ok(())
39	} else {
40		Err(error!(operation::sort_key_not_orderable(key.column.clone(), ty)))
41	}
42}
43
44pub(crate) fn charge_query_memory_bytes(budget: &MemoryBudget, charged: &mut usize, total: usize) -> Result<()> {
45	if total > *charged {
46		let delta = (total - *charged) as u64;
47		if !budget.try_charge(ByteSize::from_bytes(delta)) {
48			return Err(error!(query::memory_limit_exceeded(budget.used(), budget.limit())));
49		}
50		*charged = total;
51	}
52	Ok(())
53}
54
55pub trait QueryNode: Send + Sync {
56	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()>;
57
58	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>>;
59
60	fn headers(&self) -> Option<ColumnHeaders>;
61}
62
63#[derive(Clone)]
64pub struct QueryContext {
65	pub services: Arc<Services>,
66	pub source: Option<ResolvedObject>,
67	pub batch_size: u64,
68	pub params: Params,
69	pub symbols: SymbolTable,
70	pub identity: IdentityId,
71	pub memory: Arc<MemoryBudget>,
72}
73
74impl QueryNode for Box<dyn QueryNode> {
75	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
76		(**self).initialize(rx, ctx)
77	}
78
79	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
80		let result = (**self).next(rx, ctx)?;
81		if let Some(ref batch) = result {
82			check_user_columns(batch)?;
83			reifydb_assertions! {
84				reifydb_value::value::canonical::assert_canonical_floats(batch, "volcano next");
85			}
86		}
87		Ok(result)
88	}
89
90	fn headers(&self) -> Option<ColumnHeaders> {
91		(**self).headers()
92	}
93}
94
95pub fn eval_context_from_query<'a>(ctx: &'a QueryContext) -> EvalContext<'a> {
96	EvalContext {
97		target: None,
98		batch: empty_batch(),
99		row_count: 1,
100		take: None,
101		params: &ctx.params,
102		symbols: &ctx.symbols,
103		routines: &ctx.services.routines,
104		runtime_context: &ctx.services.runtime_context,
105		identity: ctx.identity,
106	}
107}
108
109pub fn eval_context_from_transform<'a>(ctx: &'a TransformContext<'a>, stored: &'a QueryContext) -> EvalContext<'a> {
110	EvalContext {
111		target: None,
112		batch: empty_batch(),
113		row_count: 1,
114		take: None,
115		params: ctx.params,
116		symbols: &stored.symbols,
117		routines: &stored.services.routines,
118		runtime_context: ctx.runtime_context,
119		identity: stored.identity,
120	}
121}
122
123#[cfg(test)]
124mod tests {
125	#[cfg(reifydb_assertions)]
126	use std::sync::Arc;
127
128	#[cfg(reifydb_assertions)]
129	use arrow_array::{ArrayRef, Float64Array, RecordBatch};
130	#[cfg(reifydb_assertions)]
131	use reifydb_core::value::column::headers::ColumnHeaders;
132	use reifydb_core::{
133		util::budget::MemoryBudget,
134		value::{
135			batch::{batch, data_size},
136			column::factory::int4,
137		},
138	};
139	#[cfg(reifydb_assertions)]
140	use reifydb_evaluate::stack::SymbolTable;
141	#[cfg(reifydb_assertions)]
142	use reifydb_test_harness::engine::create_test_admin_transaction;
143	#[cfg(reifydb_assertions)]
144	use reifydb_transaction::transaction::Transaction;
145	use reifydb_value::byte_size::ByteSize;
146	#[cfg(reifydb_assertions)]
147	use reifydb_value::{params::Params, value::identity::IdentityId};
148
149	use super::charge_query_memory_bytes;
150	#[cfg(reifydb_assertions)]
151	use super::{QueryContext, QueryNode, query_budget};
152	#[cfg(reifydb_assertions)]
153	use crate::{Result, vm::services::Services};
154
155	#[test]
156	fn charge_query_memory_delta_charges_and_rejects_over_budget() {
157		let budget = MemoryBudget::new(ByteSize::from_kib(1));
158		let mut charged = 0usize;
159
160		let small = batch(vec![int4("c", [1i32, 2, 3, 4])]).expect("one int4 column");
161		let small_size = data_size(&small).expect("an int4 batch has a data size");
162		charge_query_memory_bytes(&budget, &mut charged, small_size).expect("small buffer fits under 1 KiB");
163		let after_first = budget.used().as_bytes();
164		assert!(after_first > 0, "charging a non-empty buffer must consume budget");
165		assert_eq!(charged as u64, after_first, "charged must track exactly what the budget recorded");
166
167		charge_query_memory_bytes(&budget, &mut charged, small_size)
168			.expect("re-charge of the same buffer is free");
169		assert_eq!(
170			budget.used().as_bytes(),
171			after_first,
172			"delta charging must not double count an unchanged buffer"
173		);
174
175		let big = batch(vec![int4("c", 0..4000i32)]).expect("one int4 column");
176		let mut big_charged = 0usize;
177		let big_size = data_size(&big).expect("an int4 batch has a data size");
178		let err = charge_query_memory_bytes(&budget, &mut big_charged, big_size).unwrap_err();
179		assert_eq!(err.0.code, "QUERY_006", "an over-budget charge must raise the memory-limit diagnostic");
180	}
181
182	#[cfg(reifydb_assertions)]
183	struct NegativeZeroNode;
184
185	#[cfg(reifydb_assertions)]
186	impl QueryNode for NegativeZeroNode {
187		fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
188			Ok(())
189		}
190
191		fn next<'a>(
192			&mut self,
193			_rx: &mut Transaction<'a>,
194			_ctx: &mut QueryContext,
195		) -> Result<Option<RecordBatch>> {
196			let column: ArrayRef = Arc::new(Float64Array::from(vec![-0.0f64]));
197			Ok(Some(RecordBatch::try_from_iter([("c", column)]).unwrap()))
198		}
199
200		fn headers(&self) -> Option<ColumnHeaders> {
201			None
202		}
203	}
204
205	#[test]
206	#[cfg(reifydb_assertions)]
207	#[should_panic(expected = "is not canonical")]
208	fn test_box_next_with_negative_zero_panics() {
209		// Every operator's output passes through this Box, so a raw -0.0 must stop here or compare splits zero.
210		let services = Services::testing();
211		let mut txn = create_test_admin_transaction();
212		let mut ctx = QueryContext {
213			services: services.clone(),
214			source: None,
215			batch_size: 1,
216			params: Params::default(),
217			symbols: SymbolTable::new(),
218			identity: IdentityId::system(),
219			memory: query_budget(&services),
220		};
221		let mut node: Box<dyn QueryNode> = Box::new(NegativeZeroNode);
222		let _ = <Box<dyn QueryNode> as QueryNode>::next(&mut node, &mut Transaction::Admin(&mut txn), &mut ctx);
223	}
224}
reifydb_engine/vm/volcano/query.rs

reifydb_engine/vm/volcano/
query.rs
