reifydb_sdk/flow/operator/
context.rs1use 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}