From 9d294e9a65d6110d2b0aafec1cd0f8d9aacb711f Mon Sep 17 00:00:00 2001
From: Andrey Neporada <aneporada@ydb.tech>
Date: Tue, 4 Jun 2024 09:22:20 +0300
Subject: [PATCH] [YQL-18511] Detect GROUPING() column count mismatch (merge to
 stable) (#5146)

---
 ydb/library/yql/sql/v1/select.cpp | 22 ++++++++++++++++++++++
 ydb/library/yql/sql/v1/source.cpp |  4 ++++
 ydb/library/yql/sql/v1/source.h   |  1 +
 3 files changed, 27 insertions(+)

diff --git a/ydb/library/yql/sql/v1/select.cpp b/ydb/library/yql/sql/v1/select.cpp
index 64c300b28a6f..d9cf20f20ecc 100644
--- a/ydb/library/yql/sql/v1/select.cpp
+++ b/ydb/library/yql/sql/v1/select.cpp
@@ -1316,6 +1316,20 @@ class TCompositeSelect: public IRealSource {
                 return false;
             }
         }
+
+        TMaybe<size_t> groupingColumnsCount;
+        size_t idx = 0;
+        for (const auto& select : Subselects) {
+            size_t count = select->GetGroupingColumnsCount();
+            if (!groupingColumnsCount.Defined()) {
+                groupingColumnsCount = count;
+            } else if (*groupingColumnsCount != count) {
+                ctx.Error(select->GetPos()) << TStringBuilder() << "Mismatch GROUPING() column count in composite select input #"
+                    << idx << ": expected " << *groupingColumnsCount << ", got: " << count << ". Please submit bug report";
+                return false;
+            }
+            ++idx;
+        }
         return true;
     }
 
@@ -1494,6 +1508,10 @@ class TSelectCore: public IRealSource, public IComposableSource {
         ISource::GetInputTables(tableList);
     }
 
+    size_t GetGroupingColumnsCount() const override {
+        return Source->GetGroupingColumnsCount();
+    }
+
     bool DoInit(TContext& ctx, ISource* initSrc) override {
         if (AsInner) {
             Source->UseAsInner();
@@ -2598,6 +2616,10 @@ class TNestedProxySource: public IProxySource {
         return true;
     }
 
+    size_t GetGroupingColumnsCount() const override {
+        return Hints.size();
+    }
+
     TNodePtr BuildGroupingColumns(const TString& label) override {
         if (Hints.empty()) {
             return nullptr;
diff --git a/ydb/library/yql/sql/v1/source.cpp b/ydb/library/yql/sql/v1/source.cpp
index f5c59806a88d..aa736f84b1c3 100644
--- a/ydb/library/yql/sql/v1/source.cpp
+++ b/ydb/library/yql/sql/v1/source.cpp
@@ -441,6 +441,10 @@ bool ISource::AddGrouping(TContext& ctx, const TVector<TString>& columns, TStrin
     return false;
 }
 
+size_t ISource::GetGroupingColumnsCount() const {
+    return 0;
+}
+
 TNodePtr ISource::BuildFilter(TContext& ctx, const TString& label) {
     return Filters.empty() ? nullptr : Y(ctx.UseUnordered(*this) ? "OrderedFilter" : "Filter", label, BuildFilterLambda());
 }
diff --git a/ydb/library/yql/sql/v1/source.h b/ydb/library/yql/sql/v1/source.h
index c7835e716366..ee805ac76dd8 100644
--- a/ydb/library/yql/sql/v1/source.h
+++ b/ydb/library/yql/sql/v1/source.h
@@ -75,6 +75,7 @@ namespace NSQLTranslationV1 {
         virtual bool SetSamplingOptions(TContext& ctx, TPosition pos, ESampleClause clause, ESampleMode mode, TNodePtr samplingRate, TNodePtr samplingSeed);
         virtual bool SetTableHints(TContext& ctx, TPosition pos, const TTableHints& hints, const TTableHints& contextHints);
         virtual bool AddGrouping(TContext& ctx, const TVector<TString>& columns, TString& groupingColumn);
+        virtual size_t GetGroupingColumnsCount() const;
         virtual TNodePtr BuildFilter(TContext& ctx, const TString& label);
         virtual TNodePtr BuildFilterLambda();
         virtual TNodePtr BuildFlattenByColumns(const TString& label);