Skip to main content

reifydb_sdk/flow/operator/
context.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::ops::Bound;
5
6use reifydb_codec::{
7	key::encoded::EncodedKey,
8	row::operator::{EncodedOperatorRow, OperatorState},
9};
10use reifydb_core::{
11	interface::catalog::flow::OperatorId,
12	key::operator_state::{GroupId, GroupStateKey},
13	state::store::TimerKind,
14};
15use reifydb_flow::operator::state::reclaim::ReclaimOutcome;
16use reifydb_value::value::{
17	Value,
18	datetime::DateTime,
19	dictionary::{DictionaryEntryId, DictionaryId},
20	row_number::RowNumber,
21};
22
23use crate::{
24	error::Result,
25	flow::operator::column::{row::Row, sink::RowSink},
26};
27
28pub trait GuestEmit {
29	type Sink: RowSink;
30	fn sink(&mut self) -> &mut Self::Sink;
31	fn finish(self, row_numbers: &[RowNumber]) -> Result<()>;
32}
33
34pub trait GuestUpdateEmit {
35	type Sink: RowSink;
36	fn pre(&mut self) -> &mut Self::Sink;
37	fn post(&mut self) -> &mut Self::Sink;
38	fn finish(self, row_numbers: &[RowNumber]) -> Result<()>;
39}
40
41pub trait GuestState {
42	fn get<T: OperatorState>(&self, key: &GroupStateKey) -> Result<Option<T>>;
43	fn set<T: OperatorState>(&mut self, key: &GroupStateKey, value: &T) -> Result<()>;
44	fn remove(&mut self, key: &GroupStateKey) -> Result<()>;
45	fn contains(&self, key: &GroupStateKey) -> Result<bool>;
46	fn clear(&mut self) -> Result<()>;
47	fn scan_prefix<T: OperatorState>(&self, prefix: &GroupStateKey) -> Result<Vec<(GroupStateKey, T)>>;
48	fn get_many<T: OperatorState>(&self, keys: &[GroupStateKey]) -> Result<Vec<(GroupStateKey, T)>>;
49	fn keys_with_prefix(&self, prefix: &GroupStateKey) -> Result<Vec<GroupStateKey>>;
50	fn range<T: OperatorState>(
51		&self,
52		start: Bound<&GroupStateKey>,
53		end: Bound<&GroupStateKey>,
54	) -> Result<Vec<(GroupStateKey, T)>>;
55	fn get_bytes(&self, key: &GroupStateKey) -> Result<Option<EncodedOperatorRow>>;
56
57	fn set_bytes(&mut self, key: &GroupStateKey, payload: EncodedOperatorRow) -> Result<()>;
58
59	fn get_many_bytes_visit(
60		&self,
61		keys: &[GroupStateKey],
62		visit: &mut dyn FnMut(GroupStateKey, EncodedOperatorRow) -> Result<()>,
63	) -> Result<()>;
64
65	fn range_bytes_visit(
66		&self,
67		start: Bound<&GroupStateKey>,
68		end: Bound<&GroupStateKey>,
69		visit: &mut dyn FnMut(GroupStateKey, EncodedOperatorRow) -> Result<()>,
70	) -> Result<()>;
71}
72
73pub trait GuestDictionary {
74	fn id_by_name(&mut self, name: &str) -> Result<Option<DictionaryId>>;
75	fn find(&mut self, dictionary: DictionaryId, value: &Value) -> Result<Option<DictionaryEntryId>>;
76	fn get(&mut self, dictionary: DictionaryId, id: DictionaryEntryId) -> Result<Option<Value>>;
77}
78
79pub trait GuestContext {
80	type InsertEmit<'a>: GuestEmit
81	where
82		Self: 'a;
83	type UpdateEmit<'a>: GuestUpdateEmit
84	where
85		Self: 'a;
86	type RemoveEmit<'a>: GuestEmit
87	where
88		Self: 'a;
89
90	fn operator_id(&self) -> OperatorId;
91	fn written_at(&self) -> DateTime;
92	fn state(&mut self) -> impl GuestState + '_;
93	fn dictionary(&mut self) -> impl GuestDictionary + '_;
94	fn intern_groups(&mut self, groups: &[EncodedKey]) -> Result<Vec<(GroupId, bool)>>;
95	fn lookup_groups(&mut self, groups: &[EncodedKey]) -> Result<Vec<Option<GroupId>>>;
96	fn get_or_create_row_numbers(&mut self, group: GroupId, keys: &[EncodedKey]) -> Result<Vec<(RowNumber, bool)>>;
97	fn get_or_create_row_numbers_for_pairs(
98		&mut self,
99		pairs: &[(GroupId, EncodedKey)],
100	) -> Result<Vec<(RowNumber, bool)>>;
101	fn remove_row_number(&mut self, group: GroupId, key: &EncodedKey) -> Result<()>;
102	fn remove_row_numbers_below(&mut self, group: GroupId, upper: &EncodedKey) -> Result<Vec<RowNumber>>;
103	fn reclaim_group_identity(&mut self, group: GroupId, limit: usize) -> Result<ReclaimOutcome>;
104	fn arm_timer(&mut self, due: DateTime, kind: TimerKind, key: &EncodedKey) -> Result<()>;
105	fn disarm_timer(&mut self, due: DateTime, kind: TimerKind, key: &EncodedKey) -> Result<()>;
106	fn flow_watermark(&mut self) -> Result<Option<DateTime>>;
107
108	fn insert_emit<R: Row>(&mut self, row_capacity: usize) -> Result<Self::InsertEmit<'_>>;
109	fn update_emit<R: Row>(&mut self, row_capacity: usize) -> Result<Self::UpdateEmit<'_>>;
110	fn remove_emit<R: Row>(&mut self, row_capacity: usize) -> Result<Self::RemoveEmit<'_>>;
111
112	fn emit_insert<R: Row>(&mut self, rows: &[R], row_numbers: &[RowNumber]) -> Result<()> {
113		if rows.is_empty() {
114			return Ok(());
115		}
116		let mut emit = self.insert_emit::<R>(rows.len())?;
117		for bytes in rows {
118			bytes.encode_into(emit.sink())?;
119		}
120		emit.finish(row_numbers)
121	}
122
123	fn emit_update<R: Row>(&mut self, pre: &[R], post: &[R], row_numbers: &[RowNumber]) -> Result<()> {
124		if row_numbers.is_empty() {
125			return Ok(());
126		}
127		let mut emit = self.update_emit::<R>(row_numbers.len())?;
128		for row in pre {
129			row.encode_into(emit.pre())?;
130		}
131		for row in post {
132			row.encode_into(emit.post())?;
133		}
134		emit.finish(row_numbers)
135	}
136
137	fn emit_remove<R: Row>(&mut self, rows: &[R], row_numbers: &[RowNumber]) -> Result<()> {
138		if rows.is_empty() {
139			return Ok(());
140		}
141		let mut emit = self.remove_emit::<R>(rows.len())?;
142		for row in rows {
143			row.encode_into(emit.sink())?;
144		}
145		emit.finish(row_numbers)
146	}
147}