reifydb_sub_flow/operator/
extend.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::sync::LazyLock;
5
6use reifydb_abi::operator::capabilities::OperatorCapability;
7use reifydb_core::{
8	interface::{
9		catalog::flow::FlowNodeId,
10		change::{Change, Diff},
11	},
12	value::column::{ColumnWithName, columns::Columns},
13};
14use reifydb_engine::{
15	expression::{
16		compile::{CompiledExpr, compile_expression},
17		context::{CompileContext, EvalContext},
18	},
19	vm::stack::SymbolTable,
20};
21use reifydb_routine::routine::registry::Routines;
22use reifydb_rql::expression::{Expression, name::display_label};
23use reifydb_runtime::context::RuntimeContext;
24use reifydb_value::{Result, fragment::Fragment, params::Params, value::identity::IdentityId};
25
26use crate::{Operator, operator::OperatorCell, transaction::FlowTransaction};
27
28static EMPTY_PARAMS: Params = Params::None;
29static EMPTY_SYMBOL_TABLE: LazyLock<SymbolTable> = LazyLock::new(SymbolTable::new);
30
31pub struct ExtendOperator {
32	parent: OperatorCell,
33	node: FlowNodeId,
34	expressions: Vec<Expression>,
35	compiled_expressions: Vec<CompiledExpr>,
36	routines: Routines,
37	runtime_context: RuntimeContext,
38}
39
40impl ExtendOperator {
41	pub fn new(
42		parent: OperatorCell,
43		node: FlowNodeId,
44		expressions: Vec<Expression>,
45		routines: Routines,
46		runtime_context: RuntimeContext,
47	) -> Self {
48		let compile_ctx = CompileContext {
49			symbols: &EMPTY_SYMBOL_TABLE,
50		};
51		let compiled_expressions: Vec<CompiledExpr> = expressions
52			.iter()
53			.map(|e| compile_expression(&compile_ctx, e))
54			.collect::<Result<Vec<_>>>()
55			.expect("Failed to compile expressions");
56
57		Self {
58			parent,
59			node,
60			expressions,
61			compiled_expressions,
62			routines,
63			runtime_context,
64		}
65	}
66
67	pub(crate) fn output_schema(&self) -> Option<Columns> {
68		self.parent.output_schema()
69	}
70
71	fn extend(&self, columns: &Columns) -> Result<Columns> {
72		let row_count = columns.row_count();
73		if row_count == 0 {
74			return Ok(Columns::empty());
75		}
76
77		let session = EvalContext {
78			params: &EMPTY_PARAMS,
79			symbols: &EMPTY_SYMBOL_TABLE,
80			routines: &self.routines,
81			runtime_context: &self.runtime_context,
82			arena: None,
83			identity: IdentityId::root(),
84			is_aggregate_context: false,
85			columns: Columns::empty(),
86			row_count: 1,
87			target: None,
88			take: None,
89		};
90		let exec_ctx = session.with_eval(columns.clone(), row_count);
91
92		let mut result_columns: Vec<ColumnWithName> =
93			columns.iter().map(|col| ColumnWithName::new(col.name().clone(), col.data().clone())).collect();
94
95		for (i, compiled_expr) in self.compiled_expressions.iter().enumerate() {
96			let evaluated_col = compiled_expr.execute(&exec_ctx)?;
97
98			let expr = &self.expressions[i];
99			let field_name = display_label(expr).text().to_string();
100
101			result_columns.push(ColumnWithName::new(
102				Fragment::internal(field_name),
103				evaluated_col.data().clone(),
104			));
105		}
106
107		let row_numbers = if columns.row_numbers.is_empty() {
108			Vec::new()
109		} else {
110			columns.row_numbers.iter().cloned().collect()
111		};
112
113		Ok(Columns::with_system_columns(
114			result_columns,
115			row_numbers,
116			columns.created_at.to_vec(),
117			columns.updated_at.to_vec(),
118		))
119	}
120}
121
122impl Operator for ExtendOperator {
123	fn id(&self) -> FlowNodeId {
124		self.node
125	}
126
127	fn capabilities(&self) -> &[OperatorCapability] {
128		OperatorCapability::STANDARD
129	}
130
131	fn apply(&self, _txn: &mut FlowTransaction, change: Change) -> Result<Change> {
132		let mut result = Vec::new();
133
134		for diff in change.diffs.into_iter() {
135			match diff {
136				Diff::Insert {
137					post,
138					..
139				} => {
140					let extended = match self.extend(&post) {
141						Ok(extended) => extended,
142						Err(err) => {
143							panic!("{:#?}", err)
144						}
145					};
146
147					if !extended.is_empty() {
148						result.push(Diff::insert(extended));
149					}
150				}
151				Diff::Update {
152					pre,
153					post,
154					..
155				} => {
156					let extended_post = self.extend(&post)?;
157					let extended_pre = self.extend(&pre)?;
158
159					if !extended_post.is_empty() {
160						result.push(Diff::update(extended_pre, extended_post));
161					}
162				}
163				Diff::Remove {
164					pre,
165					..
166				} => {
167					let extended_pre = self.extend(&pre)?;
168					if !extended_pre.is_empty() {
169						result.push(Diff::remove(extended_pre));
170					}
171				}
172			}
173		}
174
175		Ok(Change::from_flow(self.node, change.version, result, change.changed_at))
176	}
177}
reifydb_sub_flow/operator/extend.rs

reifydb_sub_flow/operator/
extend.rs
