reifydb_sub_flow/operator/window/
aggregate.rs1use std::collections::{BTreeMap, HashMap};
5
6use reifydb_abi::operator::capabilities::OperatorCapability;
7use reifydb_core::{
8 interface::{
9 catalog::flow::FlowNodeId,
10 change::{Change, Diff},
11 },
12 value::column::columns::Columns,
13 window::{
14 engine::{LatePolicy, config::WindowEngineConfig, tumbling::TumblingBuckets},
15 span::WindowSpan,
16 },
17};
18use reifydb_engine::flow::aggregate::AggregateContext;
19use reifydb_routine::routine::registry::Routines;
20use reifydb_rql::expression::Expression;
21use reifydb_runtime::context::RuntimeContext;
22use reifydb_value::{Result, util::hash::Hash128, value::Value};
23
24use super::{
25 accumulator::WindowSlotKey,
26 aggregation::Aggregation,
27 tumbling::{finish_tumbling_engine, route_into_buckets},
28};
29use crate::{
30 operator::{Operator, OperatorCell},
31 transaction::FlowTransaction,
32};
33
34type EngineBuckets = TumblingBuckets<Hash128, u64, (WindowSlotKey, Vec<Option<Value>>)>;
35
36pub struct AggregateOperator {
37 core: Aggregation,
38}
39
40impl AggregateOperator {
41 pub fn new(
42 parent: OperatorCell,
43 node: FlowNodeId,
44 by: Vec<Expression>,
45 map: Vec<Expression>,
46 routines: Routines,
47 runtime_context: RuntimeContext,
48 ) -> Self {
49 Self {
50 core: Aggregation::new(
51 node,
52 parent,
53 by,
54 map,
55 routines,
56 runtime_context,
57 AggregateContext::Grouped,
58 ),
59 }
60 }
61
62 pub(crate) fn output_schema(&self) -> Option<Columns> {
63 self.core.parent.output_schema()
64 }
65}
66
67impl Operator for AggregateOperator {
68 fn id(&self) -> FlowNodeId {
69 self.core.node
70 }
71
72 fn capabilities(&self) -> &[OperatorCapability] {
73 OperatorCapability::STANDARD
74 }
75
76 fn apply(&self, txn: &mut FlowTransaction, change: Change) -> Result<Change> {
77 apply_aggregate_engine(&self.core, txn, change)
78 }
79}
80
81pub fn apply_aggregate_engine(core: &Aggregation, txn: &mut FlowTransaction, change: Change) -> Result<Change> {
82 let kinds = core.slot_kinds.clone().expect("aggregate requires representable slot kinds");
83
84 let mut buckets: EngineBuckets = BTreeMap::new();
85 let mut group_values: HashMap<Hash128, Vec<Value>> = HashMap::new();
86 let mut arrival: Vec<(Hash128, WindowSpan<u64>)> = Vec::new();
87 let mut window_max_ts: HashMap<(Hash128, WindowSpan<u64>), u64> = HashMap::new();
88
89 let degenerate_span = |_row_idx: usize| (WindowSpan::new(0u64, 1u64), 0u64);
90
91 for diff in change.diffs.iter() {
92 match diff {
93 Diff::Insert {
94 post,
95 ..
96 } => route_into_buckets(
97 core,
98 post,
99 true,
100 degenerate_span,
101 &mut buckets,
102 &mut group_values,
103 &mut arrival,
104 &mut window_max_ts,
105 )?,
106 Diff::Remove {
107 pre,
108 ..
109 } => route_into_buckets(
110 core,
111 pre,
112 false,
113 degenerate_span,
114 &mut buckets,
115 &mut group_values,
116 &mut arrival,
117 &mut window_max_ts,
118 )?,
119 Diff::Update {
120 pre,
121 post,
122 ..
123 } => {
124 route_into_buckets(
125 core,
126 pre,
127 false,
128 degenerate_span,
129 &mut buckets,
130 &mut group_values,
131 &mut arrival,
132 &mut window_max_ts,
133 )?;
134 route_into_buckets(
135 core,
136 post,
137 true,
138 degenerate_span,
139 &mut buckets,
140 &mut group_values,
141 &mut arrival,
142 &mut window_max_ts,
143 )?;
144 }
145 }
146 }
147
148 let diffs = finish_tumbling_engine(
149 core,
150 txn,
151 &change,
152 buckets,
153 &group_values,
154 arrival,
155 window_max_ts,
156 &kinds,
157 WindowEngineConfig::builder().late_policy(LatePolicy::Process).build(),
158 None,
159 false,
160 )?;
161 Ok(Change::from_flow(core.node, change.version, diffs, change.changed_at))
162}