datafusion_python/expr/
aggregate_expr.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::fmt::{Display, Formatter};
19
20use datafusion::logical_expr::expr::AggregateFunction;
21use pyo3::prelude::*;
22
23use crate::expr::PyExpr;
24
25#[pyclass(
26    frozen,
27    name = "AggregateFunction",
28    module = "datafusion.expr",
29    subclass
30)]
31#[derive(Clone)]
32pub struct PyAggregateFunction {
33    aggr: AggregateFunction,
34}
35
36impl From<PyAggregateFunction> for AggregateFunction {
37    fn from(aggr: PyAggregateFunction) -> Self {
38        aggr.aggr
39    }
40}
41
42impl From<AggregateFunction> for PyAggregateFunction {
43    fn from(aggr: AggregateFunction) -> PyAggregateFunction {
44        PyAggregateFunction { aggr }
45    }
46}
47
48impl Display for PyAggregateFunction {
49    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
50        let args: Vec<String> = self
51            .aggr
52            .params
53            .args
54            .iter()
55            .map(|expr| expr.to_string())
56            .collect();
57        write!(f, "{}({})", self.aggr.func.name(), args.join(", "))
58    }
59}
60
61#[pymethods]
62impl PyAggregateFunction {
63    /// Get the aggregate type, such as "MIN", or "MAX"
64    fn aggregate_type(&self) -> String {
65        self.aggr.func.name().to_string()
66    }
67
68    /// is this a distinct aggregate such as `COUNT(DISTINCT expr)`
69    fn is_distinct(&self) -> bool {
70        self.aggr.params.distinct
71    }
72
73    /// Get the arguments to the aggregate function
74    fn args(&self) -> Vec<PyExpr> {
75        self.aggr
76            .params
77            .args
78            .iter()
79            .map(|expr| PyExpr::from(expr.clone()))
80            .collect()
81    }
82
83    /// Get a String representation of this column
84    fn __repr__(&self) -> String {
85        format!("{self}")
86    }
87}
datafusion_python/expr/aggregate_expr.rs

datafusion_python/expr/
aggregate_expr.rs
