datafusion_python/expr/
grouping_set.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 datafusion::logical_expr::{Expr, GroupingSet};
19use pyo3::prelude::*;
20
21use crate::expr::PyExpr;
22
23#[pyclass(
24    from_py_object,
25    frozen,
26    name = "GroupingSet",
27    module = "datafusion.expr",
28    subclass
29)]
30#[derive(Clone)]
31pub struct PyGroupingSet {
32    grouping_set: GroupingSet,
33}
34
35#[pymethods]
36impl PyGroupingSet {
37    #[staticmethod]
38    #[pyo3(signature = (*exprs))]
39    fn rollup(exprs: Vec<PyExpr>) -> PyExpr {
40        Expr::GroupingSet(GroupingSet::Rollup(
41            exprs.into_iter().map(|e| e.expr).collect(),
42        ))
43        .into()
44    }
45
46    #[staticmethod]
47    #[pyo3(signature = (*exprs))]
48    fn cube(exprs: Vec<PyExpr>) -> PyExpr {
49        Expr::GroupingSet(GroupingSet::Cube(
50            exprs.into_iter().map(|e| e.expr).collect(),
51        ))
52        .into()
53    }
54
55    #[staticmethod]
56    #[pyo3(signature = (*expr_lists))]
57    fn grouping_sets(expr_lists: Vec<Vec<PyExpr>>) -> PyExpr {
58        Expr::GroupingSet(GroupingSet::GroupingSets(
59            expr_lists
60                .into_iter()
61                .map(|list| list.into_iter().map(|e| e.expr).collect())
62                .collect(),
63        ))
64        .into()
65    }
66}
67
68impl From<PyGroupingSet> for GroupingSet {
69    fn from(grouping_set: PyGroupingSet) -> Self {
70        grouping_set.grouping_set
71    }
72}
73
74impl From<GroupingSet> for PyGroupingSet {
75    fn from(grouping_set: GroupingSet) -> PyGroupingSet {
76        PyGroupingSet { grouping_set }
77    }
78}
datafusion_python/expr/grouping_set.rs

datafusion_python/expr/
grouping_set.rs
