reifydb_engine/vm/volcano/
apply_transform.rs1use 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}