From c6c214cf2c263afc59fadecfc18a06193d09b532 Mon Sep 17 00:00:00 2001
From: Andrey Neporada <aneporada@ydb.tech>
Date: Mon, 3 Jun 2024 23:14:52 +0300
Subject: [PATCH] [YQL-18511] Detect GROUPING() column count mismatch (#5143)

---
 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 653fe0674d2f..785ebdfa9270 100644
--- a/ydb/library/yql/sql/v1/select.cpp
+++ b/ydb/library/yql/sql/v1/select.cpp
@@ -1318,6 +1318,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;
     }
 
@@ -1496,6 +1510,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();
@@ -2600,6 +2618,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 54c8c812a713..751fb0bf14af 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);