lance_datafusion/
aggregate.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Aggregate specification for DataFusion aggregates.
5
6use datafusion::logical_expr::Expr;
7
8use crate::planner::Planner;
9
10/// Aggregate specification with group by and aggregate expressions.
11#[derive(Debug, Clone)]
12pub struct Aggregate {
13    /// Expressions to group by (e.g., column references).
14    pub group_by: Vec<Expr>,
15    /// Aggregate function expressions (e.g., SUM, COUNT, AVG).
16    /// Use `.alias()` on the expression to set output column names.
17    pub aggregates: Vec<Expr>,
18}
19
20impl Aggregate {
21    /// Create a new Aggregate.
22    pub fn new(group_by: Vec<Expr>, aggregates: Vec<Expr>) -> Self {
23        Self {
24            group_by,
25            aggregates,
26        }
27    }
28
29    /// Compute column names required by this aggregate.
30    ///
31    /// For COUNT(*), this returns empty. For SUM(x), GROUP BY y, this returns [x, y].
32    pub fn required_columns(&self) -> Vec<String> {
33        let mut required_columns = Vec::new();
34        for expr in self.group_by.iter().chain(self.aggregates.iter()) {
35            required_columns.extend(Planner::column_names_in_expr(expr));
36        }
37        required_columns.sort();
38        required_columns.dedup();
39        required_columns
40    }
41}
lance_datafusion/aggregate.rs

lance_datafusion/
aggregate.rs
