From a7ea30179e9ea22192e321d289bf30701557caad Mon Sep 17 00:00:00 2001
From: Dmitry Novik <n0vik@clickhouse.com>
Date: Tue, 15 Nov 2022 16:56:45 +0000
Subject: [PATCH] Optimize keys of GROUPING SETS

---
 .../OptimizeGroupByFunctionKeysPass.cpp       | 29 ++++++++++++++-----
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/src/Analyzer/Passes/OptimizeGroupByFunctionKeysPass.cpp b/src/Analyzer/Passes/OptimizeGroupByFunctionKeysPass.cpp
index da25467cd94..27f8a056a61 100644
--- a/src/Analyzer/Passes/OptimizeGroupByFunctionKeysPass.cpp
+++ b/src/Analyzer/Passes/OptimizeGroupByFunctionKeysPass.cpp
@@ -65,6 +65,19 @@ class OptimizeGroupByFunctionKeysVisitor : public InDepthQueryTreeVisitor<Optimi
         return true;
     }
 
+    static void optimizeGroupingSet(QueryTreeNodes & grouping_set)
+    {
+        QueryTreeNodes new_group_by_keys;
+        new_group_by_keys.reserve(grouping_set.size());
+        for (auto & group_by_elem : grouping_set)
+        {
+            if (!canBeEliminated(group_by_elem, grouping_set))
+                new_group_by_keys.push_back(group_by_elem);
+        }
+
+        grouping_set = std::move(new_group_by_keys);
+    }
+
 public:
     static bool needChildVisit(QueryTreeNodePtr & /*parent*/, QueryTreeNodePtr & child)
     {
@@ -81,16 +94,16 @@ public:
             return;
 
         auto & group_by = query->getGroupBy().getNodes();
-
-        QueryTreeNodes new_group_by_keys;
-        new_group_by_keys.reserve(group_by.size());
-        for (auto & group_by_elem : group_by)
+        if (query->isGroupByWithGroupingSets())
         {
-            if (!canBeEliminated(group_by_elem, group_by))
-                new_group_by_keys.push_back(group_by_elem);
+            for (auto & set : group_by)
+            {
+                auto & grouping_set = set->as<ListNode>()->getNodes();
+                optimizeGroupingSet(grouping_set);
+            }
         }
-
-        group_by = std::move(new_group_by_keys);
+        else
+            optimizeGroupingSet(group_by);
     }
 };