1use {
2 super::ExprNode,
3 crate::{
4 ast::{Aggregate, CountArgExpr},
5 parse_sql::parse_expr,
6 plan::{AggregateExprPlan, AggregateFunctionPlan, CountArgExprPlan},
7 result::Result,
8 translate::{NO_PARAMS, translate_expr},
9 },
10};
11
12#[derive(Clone, Debug)]
13pub enum AggregateNode<'a> {
14 Count(CountArgExprNode<'a>, bool), Sum(ExprNode<'a>, bool),
16 Min(ExprNode<'a>, bool),
17 Max(ExprNode<'a>, bool),
18 Avg(ExprNode<'a>, bool),
19 Variance(ExprNode<'a>, bool),
20 Stdev(ExprNode<'a>, bool),
21}
22
23#[derive(Clone, Debug)]
24pub enum CountArgExprNode<'a> {
25 Text(String),
26 Expr(ExprNode<'a>),
27}
28
29impl<'a> From<&'a str> for CountArgExprNode<'a> {
30 fn from(count_arg_str: &str) -> Self {
31 Self::Text(count_arg_str.to_owned())
32 }
33}
34
35impl<'a> From<ExprNode<'a>> for CountArgExprNode<'a> {
36 fn from(expr_node: ExprNode<'a>) -> Self {
37 Self::Expr(expr_node)
38 }
39}
40
41impl CountArgExprNode<'_> {
42 pub(super) fn build_count_arg_expr(self) -> Result<CountArgExpr> {
43 match self {
44 CountArgExprNode::Text(s) if &s == "*" => Ok(CountArgExpr::Wildcard),
45 CountArgExprNode::Text(s) => {
46 let expr = parse_expr(s).and_then(|expr| translate_expr(&expr, NO_PARAMS))?;
47
48 Ok(CountArgExpr::Expr(expr))
49 }
50 CountArgExprNode::Expr(expr_node) => expr_node.build_expr().map(CountArgExpr::Expr),
51 }
52 }
53
54 pub(super) fn build_count_arg_expr_plan(self) -> Result<CountArgExprPlan> {
55 match self {
56 CountArgExprNode::Text(s) if &s == "*" => Ok(CountArgExprPlan::Wildcard),
57 CountArgExprNode::Text(s) => {
58 let expr = parse_expr(s).and_then(|expr| translate_expr(&expr, NO_PARAMS))?;
59
60 Ok(CountArgExprPlan::Expr(expr.into()))
61 }
62 CountArgExprNode::Expr(expr_node) => {
63 expr_node.build_expr_plan().map(CountArgExprPlan::Expr)
64 }
65 }
66 }
67}
68
69impl AggregateNode<'_> {
70 pub(super) fn build_aggregate(self) -> Result<Aggregate> {
71 match self {
72 AggregateNode::Count(count_arg_expr_node, distinct) => count_arg_expr_node
73 .build_count_arg_expr()
74 .map(|expr| Aggregate::count(expr, distinct)),
75 AggregateNode::Sum(expr_node, distinct) => expr_node
76 .build_expr()
77 .map(|expr| Aggregate::sum(expr, distinct)),
78 AggregateNode::Min(expr_node, distinct) => expr_node
79 .build_expr()
80 .map(|expr| Aggregate::min(expr, distinct)),
81 AggregateNode::Max(expr_node, distinct) => expr_node
82 .build_expr()
83 .map(|expr| Aggregate::max(expr, distinct)),
84 AggregateNode::Avg(expr_node, distinct) => expr_node
85 .build_expr()
86 .map(|expr| Aggregate::avg(expr, distinct)),
87 AggregateNode::Variance(expr_node, distinct) => expr_node
88 .build_expr()
89 .map(|expr| Aggregate::variance(expr, distinct)),
90 AggregateNode::Stdev(expr_node, distinct) => expr_node
91 .build_expr()
92 .map(|expr| Aggregate::stdev(expr, distinct)),
93 }
94 }
95
96 pub(super) fn build_aggregate_expr_plan(self) -> Result<AggregateExprPlan> {
97 let (func, distinct) = match self {
98 AggregateNode::Count(count_arg_expr_node, distinct) => count_arg_expr_node
99 .build_count_arg_expr_plan()
100 .map(|expr| (AggregateFunctionPlan::Count(expr), distinct)),
101 AggregateNode::Sum(expr_node, distinct) => expr_node
102 .build_expr_plan()
103 .map(|expr| (AggregateFunctionPlan::Sum(expr), distinct)),
104 AggregateNode::Min(expr_node, distinct) => expr_node
105 .build_expr_plan()
106 .map(|expr| (AggregateFunctionPlan::Min(expr), distinct)),
107 AggregateNode::Max(expr_node, distinct) => expr_node
108 .build_expr_plan()
109 .map(|expr| (AggregateFunctionPlan::Max(expr), distinct)),
110 AggregateNode::Avg(expr_node, distinct) => expr_node
111 .build_expr_plan()
112 .map(|expr| (AggregateFunctionPlan::Avg(expr), distinct)),
113 AggregateNode::Variance(expr_node, distinct) => expr_node
114 .build_expr_plan()
115 .map(|expr| (AggregateFunctionPlan::Variance(expr), distinct)),
116 AggregateNode::Stdev(expr_node, distinct) => expr_node
117 .build_expr_plan()
118 .map(|expr| (AggregateFunctionPlan::Stdev(expr), distinct)),
119 }?;
120
121 Ok(AggregateExprPlan {
122 func,
123 distinct,
124 slot: None,
125 })
126 }
127}
128
129impl<'a> ExprNode<'a> {
130 #[must_use]
131 pub fn count(self) -> ExprNode<'a> {
132 ExprNode::Aggregate(Box::new(AggregateNode::Count(self.into(), false)))
133 }
134
135 #[must_use]
136 pub fn count_distinct(self) -> ExprNode<'a> {
137 ExprNode::Aggregate(Box::new(AggregateNode::Count(self.into(), true)))
138 }
139
140 #[must_use]
141 pub fn sum(self) -> ExprNode<'a> {
142 ExprNode::Aggregate(Box::new(AggregateNode::Sum(self, false)))
143 }
144
145 #[must_use]
146 pub fn sum_distinct(self) -> ExprNode<'a> {
147 ExprNode::Aggregate(Box::new(AggregateNode::Sum(self, true)))
148 }
149
150 #[must_use]
151 pub fn min(self) -> ExprNode<'a> {
152 ExprNode::Aggregate(Box::new(AggregateNode::Min(self, false)))
153 }
154
155 #[must_use]
156 pub fn min_distinct(self) -> ExprNode<'a> {
157 ExprNode::Aggregate(Box::new(AggregateNode::Min(self, true)))
158 }
159
160 #[must_use]
161 pub fn max(self) -> ExprNode<'a> {
162 ExprNode::Aggregate(Box::new(AggregateNode::Max(self, false)))
163 }
164
165 #[must_use]
166 pub fn max_distinct(self) -> ExprNode<'a> {
167 ExprNode::Aggregate(Box::new(AggregateNode::Max(self, true)))
168 }
169
170 #[must_use]
171 pub fn avg(self) -> ExprNode<'a> {
172 ExprNode::Aggregate(Box::new(AggregateNode::Avg(self, false)))
173 }
174
175 #[must_use]
176 pub fn avg_distinct(self) -> ExprNode<'a> {
177 ExprNode::Aggregate(Box::new(AggregateNode::Avg(self, true)))
178 }
179
180 #[must_use]
181 pub fn variance(self) -> ExprNode<'a> {
182 ExprNode::Aggregate(Box::new(AggregateNode::Variance(self, false)))
183 }
184
185 #[must_use]
186 pub fn variance_distinct(self) -> ExprNode<'a> {
187 ExprNode::Aggregate(Box::new(AggregateNode::Variance(self, true)))
188 }
189
190 #[must_use]
191 pub fn stdev(self) -> ExprNode<'a> {
192 ExprNode::Aggregate(Box::new(AggregateNode::Stdev(self, false)))
193 }
194
195 #[must_use]
196 pub fn stdev_distinct(self) -> ExprNode<'a> {
197 ExprNode::Aggregate(Box::new(AggregateNode::Stdev(self, true)))
198 }
199}
200
201pub fn count<'a, T: Into<CountArgExprNode<'a>>>(expr: T) -> ExprNode<'a> {
202 ExprNode::Aggregate(Box::new(AggregateNode::Count(expr.into(), false)))
203}
204
205pub fn count_distinct<'a, T: Into<CountArgExprNode<'a>>>(expr: T) -> ExprNode<'a> {
206 ExprNode::Aggregate(Box::new(AggregateNode::Count(expr.into(), true)))
207}
208
209pub fn sum<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
210 ExprNode::Aggregate(Box::new(AggregateNode::Sum(expr.into(), false)))
211}
212
213pub fn sum_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
214 ExprNode::Aggregate(Box::new(AggregateNode::Sum(expr.into(), true)))
215}
216
217pub fn min<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
218 ExprNode::Aggregate(Box::new(AggregateNode::Min(expr.into(), false)))
219}
220
221pub fn min_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
222 ExprNode::Aggregate(Box::new(AggregateNode::Min(expr.into(), true)))
223}
224
225pub fn max<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
226 ExprNode::Aggregate(Box::new(AggregateNode::Max(expr.into(), false)))
227}
228
229pub fn max_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
230 ExprNode::Aggregate(Box::new(AggregateNode::Max(expr.into(), true)))
231}
232
233pub fn avg<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
234 ExprNode::Aggregate(Box::new(AggregateNode::Avg(expr.into(), false)))
235}
236
237pub fn avg_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
238 ExprNode::Aggregate(Box::new(AggregateNode::Avg(expr.into(), true)))
239}
240
241pub fn variance<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
242 ExprNode::Aggregate(Box::new(AggregateNode::Variance(expr.into(), false)))
243}
244
245pub fn variance_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
246 ExprNode::Aggregate(Box::new(AggregateNode::Variance(expr.into(), true)))
247}
248
249pub fn stdev<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
250 ExprNode::Aggregate(Box::new(AggregateNode::Stdev(expr.into(), false)))
251}
252
253pub fn stdev_distinct<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
254 ExprNode::Aggregate(Box::new(AggregateNode::Stdev(expr.into(), true)))
255}
256
257#[cfg(test)]
258mod tests {
259 use crate::{
260 query_builder::{
261 avg, avg_distinct, col, count, count_distinct, expr, max, max_distinct, min,
262 min_distinct, stdev, stdev_distinct, sum, sum_distinct, test_expr, variance,
263 variance_distinct,
264 },
265 result::Error,
266 };
267
268 #[test]
269 fn aggregate() {
270 let actual = col("id").count();
271 let expected = "COUNT(id)";
272 test_expr(actual, expected);
273
274 let actual = count("id");
275 let expected = "COUNT(id)";
276 test_expr(actual, expected);
277
278 let actual = count("*");
279 let expected = "COUNT(*)";
280 test_expr(actual, expected);
281
282 let actual = count_distinct("*");
283 let expected = "COUNT(DISTINCT *)";
284 test_expr(actual, expected);
285
286 let actual = col("id").count_distinct();
287 let expected = "COUNT(DISTINCT id)";
288 test_expr(actual, expected);
289
290 let actual = count_distinct("id");
291 let expected = "COUNT(DISTINCT id)";
292 test_expr(actual, expected);
293
294 let actual = col("amount").sum();
295 let expected = "SUM(amount)";
296 test_expr(actual, expected);
297
298 let actual = sum("amount");
299 let expected = "SUM(amount)";
300 test_expr(actual, expected);
301
302 let actual = col("amount").sum_distinct();
303 let expected = "SUM(DISTINCT amount)";
304 test_expr(actual, expected);
305
306 let actual = sum_distinct("amount");
307 let expected = "SUM(DISTINCT amount)";
308 test_expr(actual, expected);
309
310 let actual = col("budget").min();
311 let expected = "MIN(budget)";
312 test_expr(actual, expected);
313
314 let actual = min("budget");
315 let expected = "MIN(budget)";
316 test_expr(actual, expected);
317
318 let actual = col("budget").min_distinct();
319 let expected = "MIN(DISTINCT budget)";
320 test_expr(actual, expected);
321
322 let actual = min_distinct("budget");
323 let expected = "MIN(DISTINCT budget)";
324 test_expr(actual, expected);
325
326 let actual = col("score").max();
327 let expected = "MAX(score)";
328 test_expr(actual, expected);
329
330 let actual = max("score");
331 let expected = "MAX(score)";
332 test_expr(actual, expected);
333
334 let actual = col("grade").max_distinct();
335 let expected = "MAX(DISTINCT grade)";
336 test_expr(actual, expected);
337
338 let actual = max_distinct("grade");
339 let expected = "MAX(DISTINCT grade)";
340 test_expr(actual, expected);
341
342 let actual = col("grade").avg();
343 let expected = "AVG(grade)";
344 test_expr(actual, expected);
345
346 let actual = avg("grade");
347 let expected = "AVG(grade)";
348 test_expr(actual, expected);
349
350 let actual = col("grade").avg_distinct();
351 let expected = "AVG(DISTINCT grade)";
352 test_expr(actual, expected);
353
354 let actual = avg_distinct("grade");
355 let expected = "AVG(DISTINCT grade)";
356 test_expr(actual, expected);
357
358 let actual = col("statistic").variance();
359 let expected = "VARIANCE(statistic)";
360 test_expr(actual, expected);
361
362 let actual = variance("statistic");
363 let expected = "VARIANCE(statistic)";
364 test_expr(actual, expected);
365
366 let actual = col("statistic").variance_distinct();
367 let expected = "VARIANCE(DISTINCT statistic)";
368 test_expr(actual, expected);
369
370 let actual = variance_distinct("statistic");
371 let expected = "VARIANCE(DISTINCT statistic)";
372 test_expr(actual, expected);
373
374 let actual = col("scatterplot").stdev();
375 let expected = "STDEV(scatterplot)";
376 test_expr(actual, expected);
377
378 let actual = stdev("scatterplot");
379 let expected = "STDEV(scatterplot)";
380 test_expr(actual, expected);
381
382 let actual = col("scatterplot").stdev_distinct();
383 let expected = "STDEV(DISTINCT scatterplot)";
384 test_expr(actual, expected);
385
386 let actual = stdev_distinct("scatterplot");
387 let expected = "STDEV(DISTINCT scatterplot)";
388 test_expr(actual, expected);
389 }
390
391 #[test]
392 fn aggregate_expr_plan_propagates_expr_error() {
393 let actual = sum(expr(")")).build_expr_plan();
394
395 assert!(matches!(actual, Err(Error::Parser(_))));
396 }
397}