Skip to main content

reifydb_engine/vm/volcano/
apply_transform.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::{error::diagnostic::query::unknown_apply_operator, value::column::headers::ColumnHeaders};
8use reifydb_extension::transform::{Transform, context::TransformContext};
9use reifydb_transaction::transaction::Transaction;
10use reifydb_value::{error::Error, fragment::Fragment, reifydb_assertions};
11use tracing::instrument;
12
13use crate::{
14	Result,
15	vm::volcano::query::{QueryContext, QueryNode},
16};
17
18pub(crate) struct ApplyTransformNode {
19	input: Box<dyn QueryNode>,
20	transform: Box<dyn Transform>,
21	context: Option<Arc<QueryContext>>,
22}
23
24impl ApplyTransformNode {
25	pub fn new(input: Box<dyn QueryNode>, transform: Box<dyn Transform>) -> Self {
26		Self {
27			input,
28			transform,
29			context: None,
30		}
31	}
32}
33
34impl QueryNode for ApplyTransformNode {
35	#[instrument(level = "trace", skip_all, name = "volcano::apply_transform::initialize")]
36	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
37		self.context = Some(Arc::new(ctx.clone()));
38		self.input.initialize(rx, ctx)?;
39		Ok(())
40	}
41
42	#[instrument(level = "trace", skip_all, name = "volcano::apply_transform::next")]
43	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
44		reifydb_assertions! {
45			assert!(self.context.is_some(), "ApplyTransformNode::next() called before initialize()");
46		}
47		let stored_ctx = self.context.as_ref().unwrap();
48
49		if let Some(columns) = self.input.next(rx, ctx)? {
50			let transform_ctx = TransformContext {
51				routines: &ctx.services.routines,
52				runtime_context: &stored_ctx.services.runtime_context,
53				params: &stored_ctx.params,
54			};
55			let result = self.transform.apply(&transform_ctx, columns)?;
56			Ok(Some(result))
57		} else {
58			Ok(None)
59		}
60	}
61
62	fn headers(&self) -> Option<ColumnHeaders> {
63		self.input.headers()
64	}
65}
66
67pub(crate) struct UnknownTransformNode {
68	operator: Fragment,
69}
70
71impl UnknownTransformNode {
72	pub fn new(operator: Fragment) -> Self {
73		Self {
74			operator,
75		}
76	}
77}
78
79impl QueryNode for UnknownTransformNode {
80	fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
81		Err(Error(Box::new(unknown_apply_operator(self.operator.clone()))))
82	}
83
84	fn next<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
85		Err(Error(Box::new(unknown_apply_operator(self.operator.clone()))))
86	}
87
88	fn headers(&self) -> Option<ColumnHeaders> {
89		None
90	}
91}