reifydb_engine/flow/
aggregate.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use reifydb_core::interface::identifier::{ColumnIdentifier, ColumnShape};
5use reifydb_routine::routine::registry::Routines;
6use reifydb_rql::expression::{ColumnExpression, Expression};
7use reifydb_value::fragment::Fragment;
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum SlotKind {
11	Count {
12		count_star: bool,
13	},
14	Sum,
15	Avg,
16	Min,
17	Max,
18	First,
19	Last,
20}
21
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
23pub enum AggregateContext {
24	Windowed,
25	Grouped,
26}
27
28pub enum SlotArg {
29	Star,
30	Column(String),
31	Expr(Expression),
32}
33
34pub fn synthetic_aggregate_column_name(idx: usize) -> String {
35	format!("__aggregate{idx}")
36}
37
38pub fn synthetic_aggregate_column(idx: usize) -> Expression {
39	let name = synthetic_aggregate_column_name(idx);
40	Expression::Column(ColumnExpression(ColumnIdentifier {
41		shape: ColumnShape::Alias(Fragment::internal(name.clone())),
42		name: Fragment::internal(name),
43	}))
44}
45
46pub fn classify_slot(routines: &Routines, expr: &Expression, context: AggregateContext) -> Option<(SlotKind, SlotArg)> {
47	let inner = match expr {
48		Expression::Alias(alias) => alias.expression.as_ref(),
49		other => other,
50	};
51	let call = match inner {
52		Expression::Call(c) => c,
53		_ => return None,
54	};
55	let name = call.func.0.text().to_string();
56	let short = name.rsplit("::").next().unwrap_or(&name);
57	let is_first_or_last = matches!(short, "first" | "last");
58	if is_first_or_last {
59		if context == AggregateContext::Grouped {
60			return None;
61		}
62	} else {
63		routines.get_aggregate_function(&name)?;
64	}
65	let arg = match call.args.as_slice() {
66		[] => SlotArg::Star,
67		[Expression::Column(col)] => SlotArg::Column(col.0.name.text().to_string()),
68		[single] => SlotArg::Expr(single.clone()),
69		_ => return None,
70	};
71	let is_star = matches!(arg, SlotArg::Star);
72	let kind = match short {
73		"count" => SlotKind::Count {
74			count_star: is_star,
75		},
76		"sum" if !is_star => SlotKind::Sum,
77		"avg" if !is_star => SlotKind::Avg,
78		"min" if !is_star => SlotKind::Min,
79		"max" if !is_star => SlotKind::Max,
80		"first" if !is_star => SlotKind::First,
81		"last" if !is_star => SlotKind::Last,
82		_ => return None,
83	};
84	Some((kind, arg))
85}
86
87pub fn rewrite_aggregates(
88	routines: &Routines,
89	expr: &mut Expression,
90	slots: &mut Vec<(SlotKind, SlotArg)>,
91	context: AggregateContext,
92) -> bool {
93	if let Some((kind, arg)) = classify_slot(routines, expr, context) {
94		let idx = slots.len();
95		slots.push((kind, arg));
96		*expr = synthetic_aggregate_column(idx);
97		return true;
98	}
99	match expr {
100		Expression::Alias(a) => rewrite_aggregates(routines, a.expression.as_mut(), slots, context),
101		Expression::Cast(c) => rewrite_aggregates(routines, c.expression.as_mut(), slots, context),
102		Expression::Prefix(p) => rewrite_aggregates(routines, p.expression.as_mut(), slots, context),
103		Expression::Add(e) => {
104			let l = rewrite_aggregates(routines, e.left.as_mut(), slots, context);
105			let r = rewrite_aggregates(routines, e.right.as_mut(), slots, context);
106			l && r
107		}
108		Expression::Sub(e) => {
109			let l = rewrite_aggregates(routines, e.left.as_mut(), slots, context);
110			let r = rewrite_aggregates(routines, e.right.as_mut(), slots, context);
111			l && r
112		}
113		Expression::Mul(e) => {
114			let l = rewrite_aggregates(routines, e.left.as_mut(), slots, context);
115			let r = rewrite_aggregates(routines, e.right.as_mut(), slots, context);
116			l && r
117		}
118		Expression::Div(e) => {
119			let l = rewrite_aggregates(routines, e.left.as_mut(), slots, context);
120			let r = rewrite_aggregates(routines, e.right.as_mut(), slots, context);
121			l && r
122		}
123		Expression::Rem(e) => {
124			let l = rewrite_aggregates(routines, e.left.as_mut(), slots, context);
125			let r = rewrite_aggregates(routines, e.right.as_mut(), slots, context);
126			l && r
127		}
128		Expression::Constant(_) => true,
129		_ => false,
130	}
131}
132
133pub fn is_representable(routines: &Routines, expr: &Expression, context: AggregateContext) -> bool {
134	let mut cloned = expr.clone();
135	let mut slots: Vec<(SlotKind, SlotArg)> = Vec::new();
136	rewrite_aggregates(routines, &mut cloned, &mut slots, context)
137}
reifydb_engine/flow/aggregate.rs

reifydb_engine/flow/
aggregate.rs
