laminar_sql/parser/aggregation_parser/
mod.rs

1//! Aggregate function detection and extraction
2//!
3//! This module analyzes SQL queries to extract aggregate functions like
4//! COUNT, SUM, MIN, MAX, AVG, STDDEV, VARIANCE, PERCENTILE, and more.
5//! It determines the aggregation strategy and maps to DataFusion names.
6
7use sqlparser::ast::{
8    Expr, Function, FunctionArg, FunctionArgExpr, GroupByExpr, OrderByExpr, Select, SelectItem,
9    SetExpr, Statement,
10};
11
12/// Types of aggregate functions supported.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum AggregateType {
15    // ── Core aggregates ─────────────────────────────────────────────
16    /// COUNT function
17    Count,
18    /// COUNT DISTINCT function
19    CountDistinct,
20    /// SUM function
21    Sum,
22    /// MIN function
23    Min,
24    /// MAX function
25    Max,
26    /// AVG function
27    Avg,
28    /// `FIRST_VALUE` function
29    FirstValue,
30    /// `LAST_VALUE` function
31    LastValue,
32
33    // ── Statistical aggregates ──────────────────────────────────────
34    /// Sample standard deviation (STDDEV / STDDEV_SAMP)
35    StdDev,
36    /// Population standard deviation (STDDEV_POP)
37    StdDevPop,
38    /// Sample variance (VARIANCE / VAR_SAMP)
39    Variance,
40    /// Population variance (VAR_POP / VARIANCE_POP)
41    VariancePop,
42    /// Median
43    Median,
44
45    // ── Percentile aggregates ───────────────────────────────────────
46    /// PERCENTILE_CONT (continuous interpolation)
47    PercentileCont,
48    /// PERCENTILE_DISC (discrete, nearest-rank)
49    PercentileDisc,
50
51    // ── Boolean aggregates ──────────────────────────────────────────
52    /// BOOL_AND / EVERY
53    BoolAnd,
54    /// BOOL_OR / ANY
55    BoolOr,
56
57    // ── Collection aggregates ───────────────────────────────────────
58    /// STRING_AGG / LISTAGG / GROUP_CONCAT
59    StringAgg,
60    /// ARRAY_AGG
61    ArrayAgg,
62
63    // ── Approximate aggregates ──────────────────────────────────────
64    /// APPROX_COUNT_DISTINCT
65    ApproxCountDistinct,
66    /// APPROX_PERCENTILE_CONT
67    ApproxPercentile,
68    /// APPROX_MEDIAN
69    ApproxMedian,
70
71    // ── Correlation / Regression ────────────────────────────────────
72    /// Covariance sample (COVAR_SAMP)
73    Covar,
74    /// Covariance population (COVAR_POP)
75    CovarPop,
76    /// Pearson correlation (CORR)
77    Corr,
78    /// Linear regression slope (REGR_SLOPE)
79    RegrSlope,
80    /// Linear regression intercept (REGR_INTERCEPT)
81    RegrIntercept,
82
83    // ── Bit aggregates ──────────────────────────────────────────────
84    /// BIT_AND
85    BitAnd,
86    /// BIT_OR
87    BitOr,
88    /// BIT_XOR
89    BitXor,
90
91    /// Custom / unrecognized aggregate function
92    Custom,
93}
94
95impl AggregateType {
96    /// Check if this aggregate is order-sensitive.
97    /// Order-sensitive aggregates require maintaining event order.
98    #[must_use]
99    pub fn is_order_sensitive(&self) -> bool {
100        matches!(
101            self,
102            AggregateType::FirstValue
103                | AggregateType::LastValue
104                | AggregateType::PercentileCont
105                | AggregateType::PercentileDisc
106                | AggregateType::StringAgg
107                | AggregateType::ArrayAgg
108        )
109    }
110
111    /// Check if this aggregate is decomposable (can be computed incrementally).
112    ///
113    /// Decomposable aggregates can be split into partial and final steps,
114    /// enabling parallel or distributed computation.
115    #[must_use]
116    pub fn is_decomposable(&self) -> bool {
117        matches!(
118            self,
119            AggregateType::Count
120                | AggregateType::Sum
121                | AggregateType::Min
122                | AggregateType::Max
123                | AggregateType::BoolAnd
124                | AggregateType::BoolOr
125                | AggregateType::BitAnd
126                | AggregateType::BitOr
127                | AggregateType::BitXor
128        )
129    }
130
131    /// Returns the DataFusion function registry name for this aggregate type,
132    /// or `None` if not directly mappable.
133    #[must_use]
134    pub fn datafusion_name(&self) -> Option<&'static str> {
135        match self {
136            AggregateType::Count | AggregateType::CountDistinct => Some("count"),
137            AggregateType::Sum => Some("sum"),
138            AggregateType::Min => Some("min"),
139            AggregateType::Max => Some("max"),
140            AggregateType::Avg => Some("avg"),
141            AggregateType::FirstValue => Some("first_value"),
142            AggregateType::LastValue => Some("last_value"),
143            AggregateType::StdDev => Some("stddev"),
144            AggregateType::StdDevPop => Some("stddev_pop"),
145            AggregateType::Variance => Some("variance"),
146            AggregateType::VariancePop => Some("variance_pop"),
147            AggregateType::Median => Some("median"),
148            AggregateType::PercentileCont => Some("percentile_cont"),
149            AggregateType::PercentileDisc => Some("percentile_disc"),
150            AggregateType::BoolAnd => Some("bool_and"),
151            AggregateType::BoolOr => Some("bool_or"),
152            AggregateType::StringAgg => Some("string_agg"),
153            AggregateType::ArrayAgg => Some("array_agg"),
154            AggregateType::ApproxCountDistinct => Some("approx_distinct"),
155            AggregateType::ApproxPercentile => Some("approx_percentile_cont"),
156            AggregateType::ApproxMedian => Some("approx_median"),
157            AggregateType::Covar => Some("covar_samp"),
158            AggregateType::CovarPop => Some("covar_pop"),
159            AggregateType::Corr => Some("corr"),
160            AggregateType::RegrSlope => Some("regr_slope"),
161            AggregateType::RegrIntercept => Some("regr_intercept"),
162            AggregateType::BitAnd => Some("bit_and"),
163            AggregateType::BitOr => Some("bit_or"),
164            AggregateType::BitXor => Some("bit_xor"),
165            AggregateType::Custom => None,
166        }
167    }
168
169    /// Returns the number of input columns required by this aggregate.
170    #[must_use]
171    pub fn arity(&self) -> usize {
172        match self {
173            AggregateType::Covar
174            | AggregateType::CovarPop
175            | AggregateType::Corr
176            | AggregateType::RegrSlope
177            | AggregateType::RegrIntercept => 2,
178            _ => 1,
179        }
180    }
181}
182
183/// Information about a detected aggregate function.
184#[derive(Debug, Clone)]
185pub struct AggregateInfo {
186    /// Type of aggregate
187    pub aggregate_type: AggregateType,
188    /// Column being aggregated (None for COUNT(*))
189    pub column: Option<String>,
190    /// Optional alias for the aggregate result
191    pub alias: Option<String>,
192    /// Whether DISTINCT is applied
193    pub distinct: bool,
194    /// FILTER clause expression (e.g. `COUNT(x) FILTER (WHERE x > 5)`)
195    pub filter: Option<Box<Expr>>,
196    /// WITHIN GROUP ORDER BY expressions
197    pub within_group: Vec<OrderByExpr>,
198}
199
200impl AggregateInfo {
201    /// Create a new aggregate info.
202    #[must_use]
203    pub fn new(aggregate_type: AggregateType, column: Option<String>) -> Self {
204        Self {
205            aggregate_type,
206            column,
207            alias: None,
208            distinct: false,
209            filter: None,
210            within_group: Vec::new(),
211        }
212    }
213
214    /// Set the alias.
215    #[must_use]
216    pub fn with_alias(mut self, alias: String) -> Self {
217        self.alias = Some(alias);
218        self
219    }
220
221    /// Set distinct flag.
222    #[must_use]
223    pub fn with_distinct(mut self, distinct: bool) -> Self {
224        self.distinct = distinct;
225        self
226    }
227
228    /// Check whether a FILTER clause is present.
229    #[must_use]
230    pub fn has_filter(&self) -> bool {
231        self.filter.is_some()
232    }
233
234    /// Check whether a WITHIN GROUP clause is present.
235    #[must_use]
236    pub fn has_within_group(&self) -> bool {
237        !self.within_group.is_empty()
238    }
239}
240
241/// Analysis result for aggregations in a query.
242#[derive(Debug, Clone, Default)]
243pub struct AggregationAnalysis {
244    /// List of aggregate functions found
245    pub aggregates: Vec<AggregateInfo>,
246    /// GROUP BY columns
247    pub group_by_columns: Vec<String>,
248    /// Whether the query has a HAVING clause
249    pub has_having: bool,
250}
251
252impl AggregationAnalysis {
253    /// Check if this analysis contains any aggregates.
254    #[must_use]
255    pub fn has_aggregates(&self) -> bool {
256        !self.aggregates.is_empty()
257    }
258
259    /// Check if any aggregate is order-sensitive.
260    #[must_use]
261    pub fn has_order_sensitive(&self) -> bool {
262        self.aggregates
263            .iter()
264            .any(|a| a.aggregate_type.is_order_sensitive())
265    }
266
267    /// Check if all aggregates are decomposable.
268    #[must_use]
269    pub fn all_decomposable(&self) -> bool {
270        self.aggregates
271            .iter()
272            .all(|a| a.aggregate_type.is_decomposable())
273    }
274
275    /// Get aggregates by type.
276    #[must_use]
277    pub fn get_by_type(&self, agg_type: AggregateType) -> Vec<&AggregateInfo> {
278        self.aggregates
279            .iter()
280            .filter(|a| a.aggregate_type == agg_type)
281            .collect()
282    }
283
284    /// Check if any aggregate has a FILTER clause.
285    #[must_use]
286    pub fn has_any_filter(&self) -> bool {
287        self.aggregates.iter().any(AggregateInfo::has_filter)
288    }
289
290    /// Check if any aggregate has a WITHIN GROUP clause.
291    #[must_use]
292    pub fn has_any_within_group(&self) -> bool {
293        self.aggregates.iter().any(AggregateInfo::has_within_group)
294    }
295}
296
297/// Analyze a SQL statement for aggregate functions.
298#[must_use]
299pub fn analyze_aggregates(stmt: &Statement) -> AggregationAnalysis {
300    let mut analysis = AggregationAnalysis::default();
301
302    if let Statement::Query(query) = stmt {
303        if let SetExpr::Select(select) = query.body.as_ref() {
304            analyze_select(&mut analysis, select);
305        }
306    }
307
308    analysis
309}
310
311/// Analyze a SELECT statement for aggregates.
312fn analyze_select(analysis: &mut AggregationAnalysis, select: &Select) {
313    // Check SELECT items for aggregate functions
314    for item in &select.projection {
315        match item {
316            SelectItem::UnnamedExpr(expr) => {
317                if let Some(agg) = extract_aggregate(expr, None) {
318                    analysis.aggregates.push(agg);
319                }
320            }
321            SelectItem::ExprWithAlias { expr, alias } => {
322                if let Some(agg) = extract_aggregate(expr, Some(alias.value.clone())) {
323                    analysis.aggregates.push(agg);
324                }
325            }
326            SelectItem::QualifiedWildcard(_, _) | SelectItem::Wildcard(_) => {}
327        }
328    }
329
330    // Extract GROUP BY columns
331    match &select.group_by {
332        GroupByExpr::Expressions(exprs, _modifiers) => {
333            for expr in exprs {
334                if let Some(col) = extract_column_name(expr) {
335                    analysis.group_by_columns.push(col);
336                }
337            }
338        }
339        GroupByExpr::All(_) => {}
340    }
341
342    analysis.has_having = select.having.is_some();
343}
344
345/// Resolve a SQL function name (upper-cased) to an [`AggregateType`], handling
346/// both canonical names and common aliases.
347fn resolve_aggregate_type(name: &str, func: &Function) -> Option<AggregateType> {
348    match name {
349        // ── Core ────────────────────────────────────────────────────
350        "COUNT" => {
351            if has_distinct_arg(func) {
352                Some(AggregateType::CountDistinct)
353            } else {
354                Some(AggregateType::Count)
355            }
356        }
357        "SUM" => Some(AggregateType::Sum),
358        "MIN" => Some(AggregateType::Min),
359        "MAX" => Some(AggregateType::Max),
360        "AVG" | "MEAN" => Some(AggregateType::Avg),
361        "FIRST_VALUE" | "FIRST" => Some(AggregateType::FirstValue),
362        "LAST_VALUE" | "LAST" => Some(AggregateType::LastValue),
363
364        // ── Statistical ────────────────────────────────────────────
365        "STDDEV" | "STDDEV_SAMP" => Some(AggregateType::StdDev),
366        "STDDEV_POP" => Some(AggregateType::StdDevPop),
367        "VARIANCE" | "VAR_SAMP" | "VAR" => Some(AggregateType::Variance),
368        "VAR_POP" | "VARIANCE_POP" => Some(AggregateType::VariancePop),
369        "MEDIAN" => Some(AggregateType::Median),
370
371        // ── Percentile ─────────────────────────────────────────────
372        "PERCENTILE_CONT" => Some(AggregateType::PercentileCont),
373        "PERCENTILE_DISC" => Some(AggregateType::PercentileDisc),
374
375        // ── Boolean ────────────────────────────────────────────────
376        "BOOL_AND" | "EVERY" => Some(AggregateType::BoolAnd),
377        "BOOL_OR" | "ANY" => Some(AggregateType::BoolOr),
378
379        // ── Collection ─────────────────────────────────────────────
380        "STRING_AGG" | "LISTAGG" | "GROUP_CONCAT" => Some(AggregateType::StringAgg),
381        "ARRAY_AGG" => Some(AggregateType::ArrayAgg),
382
383        // ── Approximate ────────────────────────────────────────────
384        "APPROX_COUNT_DISTINCT" | "APPROX_DISTINCT" => Some(AggregateType::ApproxCountDistinct),
385        "APPROX_PERCENTILE_CONT" | "APPROX_PERCENTILE" => Some(AggregateType::ApproxPercentile),
386        "APPROX_MEDIAN" => Some(AggregateType::ApproxMedian),
387
388        // ── Correlation / Regression ───────────────────────────────
389        "COVAR_SAMP" | "COVAR" => Some(AggregateType::Covar),
390        "COVAR_POP" => Some(AggregateType::CovarPop),
391        "CORR" => Some(AggregateType::Corr),
392        "REGR_SLOPE" => Some(AggregateType::RegrSlope),
393        "REGR_INTERCEPT" => Some(AggregateType::RegrIntercept),
394
395        // ── Bit ────────────────────────────────────────────────────
396        "BIT_AND" => Some(AggregateType::BitAnd),
397        "BIT_OR" => Some(AggregateType::BitOr),
398        "BIT_XOR" => Some(AggregateType::BitXor),
399
400        _ => None,
401    }
402}
403
404/// Extract aggregate function from an expression.
405fn extract_aggregate(expr: &Expr, alias: Option<String>) -> Option<AggregateInfo> {
406    match expr {
407        Expr::Function(func) => {
408            let func_name = func.name.to_string().to_uppercase();
409            let agg_type = resolve_aggregate_type(&func_name, func)?;
410
411            let column = extract_first_arg_column(func);
412            let distinct = has_distinct_arg(func);
413
414            let mut info = AggregateInfo::new(agg_type, column).with_distinct(distinct);
415
416            // Extract FILTER clause
417            if let Some(filter_expr) = &func.filter {
418                info.filter = Some(filter_expr.clone());
419            }
420
421            // Extract WITHIN GROUP clause
422            if !func.within_group.is_empty() {
423                info.within_group.clone_from(&func.within_group);
424            }
425
426            if let Some(a) = alias {
427                info = info.with_alias(a);
428            }
429            Some(info)
430        }
431        // Handle nested expressions (e.g., CAST(COUNT(*) AS INT))
432        Expr::Cast { expr, .. } | Expr::Nested(expr) => extract_aggregate(expr, alias),
433        _ => None,
434    }
435}
436
437/// Check if the function has a DISTINCT argument.
438fn has_distinct_arg(func: &Function) -> bool {
439    // In sqlparser 0.60, DISTINCT is part of FunctionArgumentList
440    match &func.args {
441        sqlparser::ast::FunctionArguments::List(list) => list.duplicate_treatment.is_some(),
442        _ => false,
443    }
444}
445
446/// Extract the column name from the first argument of a function.
447fn extract_first_arg_column(func: &Function) -> Option<String> {
448    // Handle FunctionArguments::List
449    match &func.args {
450        sqlparser::ast::FunctionArguments::List(list) => {
451            if list.args.is_empty() {
452                return None;
453            }
454            match &list.args[0] {
455                FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => extract_column_name(expr),
456                FunctionArg::Named { arg, .. } | FunctionArg::ExprNamed { arg, .. } => {
457                    if let FunctionArgExpr::Expr(expr) = arg {
458                        extract_column_name(expr)
459                    } else {
460                        None
461                    }
462                }
463                // COUNT(*), QualifiedWildcard, etc.
464                FunctionArg::Unnamed(_) => None,
465            }
466        }
467        sqlparser::ast::FunctionArguments::Subquery(_)
468        | sqlparser::ast::FunctionArguments::None => None,
469    }
470}
471
472/// Extract column name from an expression.
473fn extract_column_name(expr: &Expr) -> Option<String> {
474    match expr {
475        Expr::Identifier(ident) => Some(ident.value.clone()),
476        Expr::CompoundIdentifier(parts) => parts.last().map(|p| p.value.clone()),
477        _ => None,
478    }
479}
480
481/// Check if a SELECT statement contains any aggregate functions.
482#[must_use]
483pub fn has_aggregates(stmt: &Statement) -> bool {
484    analyze_aggregates(stmt).has_aggregates()
485}
486
487/// Count the number of aggregate functions in a statement.
488#[must_use]
489pub fn count_aggregates(stmt: &Statement) -> usize {
490    analyze_aggregates(stmt).aggregates.len()
491}
492
493#[cfg(test)]
494mod tests;
laminar_sql/parser/aggregation_parser/mod.rs

laminar_sql/parser/aggregation_parser/
mod.rs
