group by modes in SQL translation

7 files changed, 79 insertions, 10 deletions

diff --git a/ydb/library/yql/sql/v1/SQLv1.g.in b/ydb/library/yql/sql/v1/SQLv1.g.in
index 272421c07ed..a5b4e870fb7 100644
--- a/ydb/library/yql/sql/v1/SQLv1.g.in
+++ b/ydb/library/yql/sql/v1/SQLv1.g.in
@@ -368,7 +368,7 @@ order_by_clause: ORDER BY sort_specification_list;
 
 ext_order_by_clause: ASSUME? order_by_clause;
 
-group_by_clause: GROUP COMPACT? BY opt_set_quantifier grouping_element_list;
+group_by_clause: GROUP COMPACT? BY opt_set_quantifier grouping_element_list (WITH an_id)?;
 
 grouping_element_list: grouping_element (COMMA grouping_element)*;
 
diff --git a/ydb/library/yql/sql/v1/format/sql_format.cpp b/ydb/library/yql/sql/v1/format/sql_format.cpp
index 426d2507c12..53a1381dbcc 100644
--- a/ydb/library/yql/sql/v1/format/sql_format.cpp
+++ b/ydb/library/yql/sql/v1/format/sql_format.cpp
@@ -1482,6 +1482,23 @@ private:
         PopCurrentIndent();
     }
 
+    void VisitGroupByClause(const TRule_group_by_clause& msg) {
+        Visit(msg.GetToken1());
+        if (msg.HasBlock2()) {
+            Visit(msg.GetBlock2());
+        }
+
+        Visit(msg.GetToken3());
+        Visit(msg.GetRule_opt_set_quantifier4());
+        Visit(msg.GetRule_grouping_element_list5());
+        if (msg.HasBlock6()) {
+            NewLine();
+            PushCurrentIndent();
+            Visit(msg.GetBlock6());
+            PopCurrentIndent();
+        }
+    }
+
     void VisitWindowDefinitionList(const TRule_window_definition_list& msg) {
         NewLine();
         PushCurrentIndent();
@@ -1707,6 +1724,7 @@ TStaticData::TStaticData()
         {TRule_without_column_list::GetDescriptor(), MakeFunctor(&TVisitor::VisitWithoutColumnList)},
         {TRule_table_ref::GetDescriptor(), MakeFunctor(&TVisitor::VisitTableRef)},
         {TRule_grouping_element_list::GetDescriptor(), MakeFunctor(&TVisitor::VisitGroupingElementList)},
+        {TRule_group_by_clause::GetDescriptor(), MakeFunctor(&TVisitor::VisitGroupByClause)},
         {TRule_window_definition_list::GetDescriptor(), MakeFunctor(&TVisitor::VisitWindowDefinitionList)},
         {TRule_window_specification::GetDescriptor(), MakeFunctor(&TVisitor::VisitWindowSpecification)},
         {TRule_window_partition_clause::GetDescriptor(), MakeFunctor(&TVisitor::VisitWindowParitionClause)},
diff --git a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
index d28dc491519..71fc05b93f0 100644
--- a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
+++ b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
@@ -587,6 +587,8 @@ Y_UNIT_TEST_SUITE(CheckSqlFormatter) {
              "SELECT\n\t1\nFROM user\nGROUP BY\n\tkey;\n"},
             {"select 1 from user group compact by key, value as v",
              "SELECT\n\t1\nFROM user\nGROUP COMPACT BY\n\tkey,\n\tvalue AS v;\n"},
+            {"select 1 from user group by key with combine",
+             "SELECT\n\t1\nFROM user\nGROUP BY\n\tkey\n\tWITH combine;\n"},
             {"select 1 from user order by key asc",
              "SELECT\n\t1\nFROM user\nORDER BY\n\tkey ASC;\n"},
             {"select 1 from user order by key, value desc",
diff --git a/ydb/library/yql/sql/v1/node.cpp b/ydb/library/yql/sql/v1/node.cpp
index bb5e410e777..8a407d0768e 100644
--- a/ydb/library/yql/sql/v1/node.cpp
+++ b/ydb/library/yql/sql/v1/node.cpp
@@ -1339,6 +1339,10 @@ void ISource::SetCompactGroupBy(bool compactGroupBy) {
     CompactGroupBy = compactGroupBy;
 }
 
+void ISource::SetGroupBySuffix(const TString& suffix) {
+    GroupBySuffix = suffix;
+}
+
 bool ISource::AddExpressions(TContext& ctx, const TVector<TNodePtr>& expressions, EExprSeat exprSeat) {
     YQL_ENSURE(exprSeat < EExprSeat::Max);
     THashSet<TString> names;
@@ -1803,7 +1807,7 @@ TNodePtr ISource::BuildAggregation(const TString& label) {
             Q(Y(BuildQuotedAtom(Pos, SessionWindow->GetLabel()), sessionWindow->BuildTraits(label))))));
     }
 
-    return Y("AssumeColumnOrderPartial", Y("Aggregate", label, Q(keysTuple), Q(aggrArgs), Q(options)), Q(keysTuple));
+    return Y("AssumeColumnOrderPartial", Y("Aggregate" + GroupBySuffix, label, Q(keysTuple), Q(aggrArgs), Q(options)), Q(keysTuple));
 }
 
 TMaybe<TString> ISource::FindColumnMistype(const TString& name) const {
diff --git a/ydb/library/yql/sql/v1/node.h b/ydb/library/yql/sql/v1/node.h
index 633143aa34c..1a23499c79a 100644
--- a/ydb/library/yql/sql/v1/node.h
+++ b/ydb/library/yql/sql/v1/node.h
@@ -826,6 +826,7 @@ namespace NSQLTranslationV1 {
         virtual bool AddFilter(TContext& ctx, TNodePtr filter);
         virtual bool AddGroupKey(TContext& ctx, const TString& column);
         virtual void SetCompactGroupBy(bool compactGroupBy);
+        virtual void SetGroupBySuffix(const TString& suffix);
         virtual TString MakeLocalName(const TString& name);
         virtual bool AddAggregation(TContext& ctx, TAggregationPtr aggr);
         virtual bool AddFuncOverWindow(TContext& ctx, TNodePtr expr);
@@ -905,6 +906,7 @@ namespace NSQLTranslationV1 {
         THashMap<TString, TString> GroupByColumnAliases;
         TVector<TNodePtr> Filters;
         bool CompactGroupBy = false;
+        TString GroupBySuffix;
         TSet<TString> GroupKeys;
         TVector<TString> OrderedGroupKeys;
         std::array<TVector<TNodePtr>, static_cast<unsigned>(EExprSeat::Max)> NamedExprs;
@@ -1287,6 +1289,7 @@ namespace NSQLTranslationV1 {
         const TVector<TNodePtr>& groupByExpr,
         const TVector<TNodePtr>& groupBy,
         bool compactGroupBy,
+        const TString& groupBySuffix,
         bool assumeSorted,
         const TVector<TSortSpecificationPtr>& orderBy,
         TNodePtr having,
diff --git a/ydb/library/yql/sql/v1/select.cpp b/ydb/library/yql/sql/v1/select.cpp
index e2dd68db4ee..b9c3d448a3a 100644
--- a/ydb/library/yql/sql/v1/select.cpp
+++ b/ydb/library/yql/sql/v1/select.cpp
@@ -1403,6 +1403,7 @@ public:
         const TVector<TNodePtr>& groupByExpr,
         const TVector<TNodePtr>& groupBy,
         bool compactGroupBy,
+        const TString& groupBySuffix,
         bool assumeSorted,
         const TVector<TSortSpecificationPtr>& orderBy,
         TNodePtr having,
@@ -1420,6 +1421,7 @@ public:
         , GroupBy(groupBy)
         , AssumeSorted(assumeSorted)
         , CompactGroupBy(compactGroupBy)
+        , GroupBySuffix(groupBySuffix)
         , OrderBy(orderBy)
         , Having(having)
         , WinSpecs(winSpecs)
@@ -1473,6 +1475,7 @@ public:
         }
 
         src->SetCompactGroupBy(CompactGroupBy);
+        src->SetGroupBySuffix(GroupBySuffix);
 
         for (auto& term: Terms) {
             term->CollectPreaggregateExprs(ctx, *src, DistinctAggrExpr);
@@ -1763,7 +1766,7 @@ public:
             newSpecs.emplace(cur.first, cur.second->Clone());
         }
         return new TSelectCore(Pos, Source->CloneSource(), CloneContainer(GroupByExpr),
-                CloneContainer(GroupBy), CompactGroupBy, AssumeSorted, CloneContainer(OrderBy),
+                CloneContainer(GroupBy), CompactGroupBy, GroupBySuffix, AssumeSorted, CloneContainer(OrderBy),
                 SafeClone(Having), newSpecs, SafeClone(HoppingWindowSpec),
                 CloneContainer(Terms), Distinct, Without, SelectStream, Settings);
     }
@@ -2085,6 +2088,7 @@ private:
     TVector<TNodePtr> GroupBy;
     bool AssumeSorted = false;
     bool CompactGroupBy = false;
+    TString GroupBySuffix;
     TVector<TSortSpecificationPtr> OrderBy;
     TNodePtr Having;
     TWinSpecs WinSpecs;
@@ -2465,6 +2469,7 @@ TSourcePtr DoBuildSelectCore(
     const TVector<TNodePtr>& groupByExpr,
     const TVector<TNodePtr>& groupBy,
     bool compactGroupBy,
+    const TString& groupBySuffix,
     bool assumeSorted,
     const TVector<TSortSpecificationPtr>& orderBy,
     TNodePtr having,
@@ -2477,14 +2482,14 @@ TSourcePtr DoBuildSelectCore(
     const TWriteSettings& settings
 ) {
     if (groupBy.empty() || !groupBy.front()->ContentListPtr()) {
-        return new TSelectCore(pos, std::move(source), groupByExpr, groupBy, compactGroupBy, assumeSorted,
+        return new TSelectCore(pos, std::move(source), groupByExpr, groupBy, compactGroupBy, groupBySuffix, assumeSorted,
             orderBy, having, winSpecs, hoppingWindowSpec, terms, distinct, without, selectStream, settings);
     }
     if (groupBy.size() == 1) {
         /// actualy no big idea to use grouping function in this case (result allways 0)
         auto contentPtr = groupBy.front()->ContentListPtr();
         source = new TNestedProxySource(pos, *contentPtr, source);
-        return DoBuildSelectCore(ctx, pos, originalSource, source, groupByExpr, *contentPtr, compactGroupBy,
+        return DoBuildSelectCore(ctx, pos, originalSource, source, groupByExpr, *contentPtr, compactGroupBy, groupBySuffix,
             assumeSorted, orderBy, having, std::move(winSpecs),
             hoppingWindowSpec, std::move(terms), distinct, std::move(without), selectStream, settings);
     }
@@ -2512,7 +2517,7 @@ TSourcePtr DoBuildSelectCore(
         }
         totalGroups += contentPtr->size();
         TSelectCore* selectCore = new TSelectCore(pos, std::move(proxySource), CloneContainer(groupByExpr),
-            CloneContainer(*contentPtr), compactGroupBy, assumeSorted, orderBy, SafeClone(having), winSpecs,
+            CloneContainer(*contentPtr), compactGroupBy, groupBySuffix, assumeSorted, orderBy, SafeClone(having), winSpecs,
             hoppingWindowSpec, terms, distinct, without, selectStream, settings);
         subselects.emplace_back(selectCore);
     }
@@ -2533,6 +2538,7 @@ TSourcePtr BuildSelectCore(
     const TVector<TNodePtr>& groupByExpr,
     const TVector<TNodePtr>& groupBy,
     bool compactGroupBy,
+    const TString& groupBySuffix,
     bool assumeSorted,
     const TVector<TSortSpecificationPtr>& orderBy,
     TNodePtr having,
@@ -2544,7 +2550,7 @@ TSourcePtr BuildSelectCore(
     bool selectStream,
     const TWriteSettings& settings)
 {
-    return DoBuildSelectCore(ctx, pos, source, source, groupByExpr, groupBy, compactGroupBy, assumeSorted, orderBy,
+    return DoBuildSelectCore(ctx, pos, source, source, groupByExpr, groupBy, compactGroupBy, groupBySuffix, assumeSorted, orderBy,
         having, std::move(winSpecs), hoppingWindowSpec, std::move(terms), distinct, std::move(without), selectStream, settings);
 }
 
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index 184c606cff3..53f345695ca 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -3601,6 +3601,7 @@ public:
     TMap<TString, TNodePtr>& Aliases();
     THoppingWindowSpecPtr GetHoppingWindow() const;
     bool IsCompactGroupBy() const;
+    TString GetSuffix() const;
 
 private:
     TMaybe<TVector<TNodePtr>> MultiplyGroupingSets(const TVector<TNodePtr>& lhs, const TVector<TNodePtr>& rhs) const;
@@ -3624,6 +3625,7 @@ private:
     THoppingWindowSpecPtr HoppingWindowSpec; // stream queries
     static const TString AutogenerateNamePrefix;
     bool CompactGroupBy;
+    TString Suffix;
 };
 
 const TString TGroupByClause::AutogenerateNamePrefix = "group";
@@ -6710,6 +6712,7 @@ TSourcePtr TSqlSelect::SelectCore(const TRule_select_core& node, const TWriteSet
     TVector<TNodePtr> groupByExpr, groupBy;
     THoppingWindowSpecPtr hoppingWindowSpec;
     bool compactGroupBy = false;
+    TString groupBySuffix;
     if (node.HasBlock11()) {
         TGroupByClause clause(Ctx, Mode);
         if (!clause.Build(node.GetBlock11().GetRule_group_by_clause1(), source->IsStream())) {
@@ -6723,6 +6726,7 @@ TSourcePtr TSqlSelect::SelectCore(const TRule_select_core& node, const TWriteSet
         clause.SetFeatures("sql_features");
         hoppingWindowSpec = clause.GetHoppingWindow();
         compactGroupBy = clause.IsCompactGroupBy();
+        groupBySuffix = clause.GetSuffix();
     }
 
     TNodePtr having;
@@ -6810,7 +6814,7 @@ TSourcePtr TSqlSelect::SelectCore(const TRule_select_core& node, const TWriteSet
     if (!ValidateSelectColumns(terms)) {
         return nullptr;
     }
-    return BuildSelectCore(Ctx, startPos, std::move(source), groupByExpr, groupBy, compactGroupBy, assumeSorted, orderBy, having,
+    return BuildSelectCore(Ctx, startPos, std::move(source), groupByExpr, groupBy, compactGroupBy, groupBySuffix, assumeSorted, orderBy, having,
         std::move(windowSpec), hoppingWindowSpec, std::move(terms), distinct, std::move(without), selectStream, settings);
 }
 
@@ -7185,7 +7189,7 @@ bool TSqlTranslation::OrderByClause(const TRule_order_by_clause& node, TVector<T
 }
 
 bool TGroupByClause::Build(const TRule_group_by_clause& node, bool stream) {
-    // group_by_clause: GROUP COMPACT? BY opt_set_quantifier grouping_element_list;
+    // group_by_clause: GROUP COMPACT? BY opt_set_quantifier grouping_element_list (WITH an_id)?;
     CompactGroupBy = node.HasBlock2();
     if (!CompactGroupBy) {
         auto hints = Ctx.PullHintForToken(Ctx.TokenPosition(node.GetToken1()));
@@ -7200,6 +7204,33 @@ bool TGroupByClause::Build(const TRule_group_by_clause& node, bool stream) {
     if (!ParseList(node.GetRule_grouping_element_list5(), EGroupByFeatures::Ordinary)) {
         return false;
     }
+
+    if (node.HasBlock6()) {
+        TString mode = Id(node.GetBlock6().GetRule_an_id2(), *this);
+        TMaybe<TIssue> normalizeError = NormalizeName(Ctx.Pos(), mode);
+        if (!normalizeError.Empty()) {
+            Error() << normalizeError->Message;
+            Ctx.IncrementMonCounter("sql_errors", "NormalizeGroupByModeError");
+            return false;
+        }
+
+        if (mode == "combine") {
+            Suffix = "Combine";
+        } else if (mode == "combinestate") {
+            Suffix = "CombineState";
+        } else if (mode == "mergestate") {
+            Suffix = "MergeState";
+        } else if (mode == "finalize") {
+            Suffix = "Finalize";
+        } else if (mode == "mergefinalize") {
+            Suffix = "MergeFinalize";
+        } else {
+            Ctx.Error() << "Unsupported group by mode: " << mode;
+            Ctx.IncrementMonCounter("sql_errors", "GroupByModeUnknown");
+            return false;
+        }
+    }
+
     if (!ResolveGroupByAndGrouping()) {
         return false;
     }
@@ -7263,6 +7294,10 @@ bool TGroupByClause::IsCompactGroupBy() const {
     return CompactGroupBy;
 }
 
+TString TGroupByClause::GetSuffix() const {
+    return Suffix;
+}
+
 TMaybe<TVector<TNodePtr>> TGroupByClause::MultiplyGroupingSets(const TVector<TNodePtr>& lhs, const TVector<TNodePtr>& rhs) const {
     TVector<TNodePtr> content;
     for (const auto& leftNode: lhs) {
@@ -7846,6 +7881,7 @@ TSourcePtr TSqlSelect::Build(const TRule& node, TPosition pos, TSelectKindResult
         TVector<TNodePtr> groupByExpr;
         TVector<TNodePtr> groupBy;
         bool compactGroupBy = false;
+        TString groupBySuffix = "";
         TNodePtr having;
         TWinSpecs winSpecs;
         THoppingWindowSpecPtr hoppingWindowSpec;
@@ -7856,7 +7892,7 @@ TSourcePtr TSqlSelect::Build(const TRule& node, TPosition pos, TSelectKindResult
         TVector<TNodePtr> terms;
         terms.push_back(BuildColumn(unionPos, "*", ""));
 
-        result = BuildSelectCore(Ctx, unionPos, std::move(result), groupByExpr, groupBy, compactGroupBy,
+        result = BuildSelectCore(Ctx, unionPos, std::move(result), groupByExpr, groupBy, compactGroupBy, groupBySuffix,
             assumeOrderBy, orderBy, having, std::move(winSpecs), hoppingWindowSpec, std::move(terms),
             distinct, std::move(without), stream, settings);
     } else {