implement more

apache · May 6, 2022 · 1979238 · 1979238
1 parent 9d1eb18
commit 1979238
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 10 deletions.
diff --git a/datafusion/core/src/logical_plan/expr_rewriter.rs b/datafusion/core/src/logical_plan/expr_rewriter.rs
@@ -223,7 +223,14 @@ impl ExprRewritable for Expr {
                 GroupingSet::Cube(exprs) => {
                     Expr::GroupingSet(GroupingSet::Cube(rewrite_vec(exprs, rewriter)?))
                 }
-                GroupingSet::GroupingSets(_expr) => todo!(),
+                GroupingSet::GroupingSets(lists_of_exprs) => {
+                    Expr::GroupingSet(GroupingSet::GroupingSets(
+                        lists_of_exprs
+                            .iter()
+                            .map(|exprs| rewrite_vec(exprs.clone(), rewriter))
+                            .collect::<Result<Vec<_>>>()?,
+                    ))
+                }
             },
             Expr::AggregateUDF { args, fun } => Expr::AggregateUDF {
                 args: rewrite_vec(args, rewriter)?,

diff --git a/datafusion/core/src/optimizer/utils.rs b/datafusion/core/src/optimizer/utils.rs
@@ -328,7 +328,9 @@ pub fn expr_sub_expressions(expr: &Expr) -> Result<Vec<Expr>> {
         Expr::GroupingSet(grouping_set) => match grouping_set {
             GroupingSet::Rollup(exprs) => Ok(exprs.clone()),
             GroupingSet::Cube(exprs) => Ok(exprs.clone()),
-            GroupingSet::GroupingSets(_) => Err(DataFusionError::Plan("GroupingSets are not supported yet".to_string())),
+            GroupingSet::GroupingSets(_) => Err(DataFusionError::Plan(
+                "GroupingSets are not supported yet".to_string(),
+            )),
         },
         Expr::WindowFunction {
             args,
@@ -472,7 +474,9 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> {
             GroupingSet::Cube(_exprs) => {
                 Ok(Expr::GroupingSet(GroupingSet::Rollup(expressions.to_vec())))
             }
-            GroupingSet::GroupingSets(_) => Err(DataFusionError::Plan("GroupingSets are not supported yet".to_string())),
+            GroupingSet::GroupingSets(_) => Err(DataFusionError::Plan(
+                "GroupingSets are not supported yet".to_string(),
+            )),
         },
         Expr::Case { .. } => {
             let mut base_expr: Option<Box<Expr>> = None;

diff --git a/datafusion/core/src/physical_plan/planner.rs b/datafusion/core/src/physical_plan/planner.rs
@@ -184,7 +184,26 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
                     .collect::<Result<Vec<_>>>()?
                     .join(", ")
             )),
-            _ => todo!(),
+            GroupingSet::Cube(exprs) => Ok(format!(
+                "CUBE ({})",
+                exprs
+                    .iter()
+                    .map(|e| create_physical_name(e, false))
+                    .collect::<Result<Vec<_>>>()?
+                    .join(", ")
+            )),
+            GroupingSet::GroupingSets(lists_of_exprs) => {
+                let mut strings = vec![];
+                for exprs in lists_of_exprs {
+                    let exprs_str = exprs
+                        .iter()
+                        .map(|e| create_physical_name(e, false))
+                        .collect::<Result<Vec<_>>>()?
+                        .join(", ");
+                    strings.push(format!("({})", exprs_str));
+                }
+                Ok(format!("GROUPING SETS ({})", strings.join(", ")))
+            }
         },
 
         Expr::InList {

diff --git a/datafusion/core/src/sql/planner.rs b/datafusion/core/src/sql/planner.rs
@@ -4687,7 +4687,7 @@ mod tests {
         assert_expected_plan(sql, &expected);
     }
 
-    #[ignore] // not supported yet
+    #[ignore] // see https://github.com/apache/arrow-datafusion/issues/2469
     #[tokio::test]
     async fn aggregate_with_grouping_sets() {
         let sql = "SELECT id, state, age, COUNT(*) FROM person GROUP BY id, GROUPING SETS ((state), (state, age), (id, state))";

diff --git a/datafusion/core/src/sql/utils.rs b/datafusion/core/src/sql/utils.rs
@@ -27,8 +27,8 @@ use crate::{
     error::{DataFusionError, Result},
     logical_plan::{Column, ExpressionVisitor, Recursion},
 };
-use datafusion_expr::expr::GroupingSet;
 use datafusion_expr::expr::find_columns_referenced_by_expr;
+use datafusion_expr::expr::GroupingSet;
 use std::collections::HashMap;
 
 /// Collect all deeply nested `Expr::AggregateFunction` and

diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs
@@ -338,8 +338,7 @@ pub fn find_columns_referenced_by_expr(e: &Expr) -> Vec<Column> {
             .flat_map(find_columns_referenced_by_expr)
             .collect(),
         Expr::GroupingSet(set) => match set {
-            GroupingSet::Rollup(exprs)
-            | GroupingSet::Cube(exprs) => exprs
+            GroupingSet::Rollup(exprs) | GroupingSet::Cube(exprs) => exprs
                 .iter()
                 .flat_map(find_columns_referenced_by_expr)
                 .collect(),
@@ -351,8 +350,8 @@ pub fn find_columns_referenced_by_expr(e: &Expr) -> Vec<Column> {
                     }
                 }
                 cols
-            },
-        }
+            }
+        },
     }
 }