reifydb-engine 0.4.12

Query execution and processing engine for ReifyDB
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2025 ReifyDB

use std::sync::Arc;

use reifydb_core::value::column::{columns::Columns, headers::ColumnHeaders};
use reifydb_routine::{
	function::{Function, FunctionContext},
	procedure::{Procedure, context::ProcedureContext},
};
use reifydb_rql::expression::Expression;
use reifydb_transaction::transaction::Transaction;
use reifydb_type::{fragment::Fragment, params::Params, value::Value};

use crate::{
	Result,
	error::EngineError,
	expression::{context::EvalContext, eval::evaluate},
	vm::volcano::query::{QueryContext, QueryNode},
};

enum GeneratorImpl {
	Function(Arc<dyn Function>),
	Procedure(Box<dyn Procedure>),
}

pub(crate) struct GeneratorNode {
	function_name: Fragment,
	expressions: Vec<Expression>,
	context: Option<Arc<QueryContext>>,
	exhausted: bool,
	generator: Option<GeneratorImpl>,
}

impl GeneratorNode {
	pub fn new(function_name: Fragment, parameter_expressions: Vec<Expression>) -> Self {
		Self {
			function_name,
			expressions: parameter_expressions,
			context: None,
			exhausted: false,
			generator: None,
		}
	}
}

impl QueryNode for GeneratorNode {
	fn initialize<'a>(&mut self, _txn: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
		self.context = Some(Arc::new(ctx.clone()));

		let name = self.function_name.text();
		if let Some(func) = ctx.services.functions.get_generator(name) {
			self.generator = Some(GeneratorImpl::Function(func));
		} else if let Some(proc) = ctx.services.procedures.get_procedure(name) {
			self.generator = Some(GeneratorImpl::Procedure(proc));
		} else {
			return Err(EngineError::GeneratorNotFound {
				name: name.to_string(),
				fragment: self.function_name.clone(),
			}
			.into());
		}

		self.exhausted = false;
		Ok(())
	}

	fn next<'a>(&mut self, txn: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<Columns>> {
		if self.exhausted {
			return Ok(None);
		}

		let stored_ctx = self.context.as_ref().unwrap();

		let session = EvalContext::from_query(stored_ctx);
		let evaluation_ctx = session.with_eval_empty();

		// Evaluate all parameter expressions into columns
		let mut evaluated_columns = Vec::new();
		for expr in &self.expressions {
			let column = evaluate(&evaluation_ctx, expr)?;
			evaluated_columns.push(column);
		}

		let columns = match self.generator.as_ref().unwrap() {
			GeneratorImpl::Function(generator) => {
				let evaluated_params = Columns::new(evaluated_columns);
				let fn_ctx = FunctionContext::new(
					self.function_name.clone(),
					&stored_ctx.services.runtime_context,
					stored_ctx.identity,
					evaluated_params.row_count(),
				);
				generator.call(&fn_ctx, &evaluated_params)?
			}
			GeneratorImpl::Procedure(procedure) => {
				let values: Vec<Value> =
					evaluated_columns.iter().map(|col| col.data().get_value(0)).collect();
				let params = Params::Positional(Arc::new(values));
				let proc_ctx = ProcedureContext {
					params: &params,
					catalog: &stored_ctx.services.catalog,
					functions: &stored_ctx.services.functions,
					runtime_context: &stored_ctx.services.runtime_context,
					ioc: &stored_ctx.services.ioc,
				};
				procedure.call(&proc_ctx, txn)?
			}
		};

		self.exhausted = true;

		Ok(Some(columns))
	}

	fn headers(&self) -> Option<ColumnHeaders> {
		None
	}
}