Skip to main content

reifydb_sub_flow/operator/window/
aggregate.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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}