datafusion_sql/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 crate::planner::{ContextProvider, PlannerContext, SqlToRel};
19use datafusion_common::plan_err;
20use datafusion_common::{DFSchema, Result};
21use datafusion_expr::{Expr, GroupingSet};
22use sqlparser::ast::Expr as SQLExpr;
23
24impl<S: ContextProvider> SqlToRel<'_, S> {
25    pub(super) fn sql_grouping_sets_to_expr(
26        &self,
27        exprs: Vec<Vec<SQLExpr>>,
28        schema: &DFSchema,
29        planner_context: &mut PlannerContext,
30    ) -> Result<Expr> {
31        let args: Result<Vec<Vec<_>>> = exprs
32            .into_iter()
33            .map(|v| {
34                v.into_iter()
35                    .map(|e| self.sql_expr_to_logical_expr(e, schema, planner_context))
36                    .collect()
37            })
38            .collect();
39        Ok(Expr::GroupingSet(GroupingSet::GroupingSets(args?)))
40    }
41
42    pub(super) fn sql_rollup_to_expr(
43        &self,
44        exprs: Vec<Vec<SQLExpr>>,
45        schema: &DFSchema,
46        planner_context: &mut PlannerContext,
47    ) -> Result<Expr> {
48        let args: Result<Vec<_>> = exprs
49            .into_iter()
50            .map(|v| {
51                if v.len() != 1 {
52                    plan_err!(
53                        "Tuple expressions are not supported for Rollup expressions"
54                    )
55                } else {
56                    self.sql_expr_to_logical_expr(v[0].clone(), schema, planner_context)
57                }
58            })
59            .collect();
60        Ok(Expr::GroupingSet(GroupingSet::Rollup(args?)))
61    }
62
63    pub(super) fn sql_cube_to_expr(
64        &self,
65        exprs: Vec<Vec<SQLExpr>>,
66        schema: &DFSchema,
67        planner_context: &mut PlannerContext,
68    ) -> Result<Expr> {
69        let args: Result<Vec<_>> = exprs
70            .into_iter()
71            .map(|v| {
72                if v.len() != 1 {
73                    plan_err!("Tuple expressions not are supported for Cube expressions")
74                } else {
75                    self.sql_expr_to_logical_expr(v[0].clone(), schema, planner_context)
76                }
77            })
78            .collect();
79        Ok(Expr::GroupingSet(GroupingSet::Cube(args?)))
80    }
81}