Add pragma CompactGroupBy

5 files changed, 50 insertions, 4 deletions

diff --git a/ydb/library/yql/sql/v1/context.cpp b/ydb/library/yql/sql/v1/context.cpp
index 4440080666..de912d1e83 100644
--- a/ydb/library/yql/sql/v1/context.cpp
+++ b/ydb/library/yql/sql/v1/context.cpp
@@ -66,6 +66,7 @@ THashMap<TStringBuf, TPragmaMaybeField> CTX_PRAGMA_MAYBE_FIELDS = {
     {"AnsiRankForNullableKeys", &TContext::AnsiRankForNullableKeys},
     {"AnsiOrderByLimitInUnionAll", &TContext::AnsiOrderByLimitInUnionAll},
     {"AnsiInForEmptyOrNullableItemsCollections", &TContext::AnsiInForEmptyOrNullableItemsCollections},
+    {"CompactGroupBy", &TContext::CompactGroupBy},
 };
 
 } // namespace
diff --git a/ydb/library/yql/sql/v1/context.h b/ydb/library/yql/sql/v1/context.h
index 19e155de1f..980d7e262d 100644
--- a/ydb/library/yql/sql/v1/context.h
+++ b/ydb/library/yql/sql/v1/context.h
@@ -299,6 +299,7 @@ namespace NSQLTranslationV1 {
         bool UseBlocks = false;
         bool AnsiLike = false;
         bool FeatureR010 = false; //Row pattern recognition: FROM clause
+        TMaybe<bool> CompactGroupBy;
     };
 
     class TColumnRefScope {
diff --git a/ydb/library/yql/sql/v1/sql_group_by.cpp b/ydb/library/yql/sql/v1/sql_group_by.cpp
index eed36c4768..dfdcbd0516 100644
--- a/ydb/library/yql/sql/v1/sql_group_by.cpp
+++ b/ydb/library/yql/sql/v1/sql_group_by.cpp
@@ -11,10 +11,14 @@ const TString TGroupByClause::AutogenerateNamePrefix = "group";
 
 bool TGroupByClause::Build(const TRule_group_by_clause& node) {
     // 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()));
-        CompactGroupBy = AnyOf(hints, [](const NSQLTranslation::TSQLHint& hint) { return to_lower(hint.Name) == "compact"; });
+    if (Ctx.CompactGroupBy.Defined()) {
+        CompactGroupBy = *Ctx.CompactGroupBy;
+    } else {
+        CompactGroupBy = node.HasBlock2();
+        if (!CompactGroupBy) {
+            auto hints = Ctx.PullHintForToken(Ctx.TokenPosition(node.GetToken1()));
+            CompactGroupBy = AnyOf(hints, [](const NSQLTranslation::TSQLHint& hint) { return to_lower(hint.Name) == "compact"; });
+        }
     }
     TPosition distinctPos;
     if (IsDistinctOptSet(node.GetRule_opt_set_quantifier4(), distinctPos)) {
diff --git a/ydb/library/yql/sql/v1/sql_query.cpp b/ydb/library/yql/sql/v1/sql_query.cpp
index 18b7c31291..b8ced318ce 100644
--- a/ydb/library/yql/sql/v1/sql_query.cpp
+++ b/ydb/library/yql/sql/v1/sql_query.cpp
@@ -1968,6 +1968,12 @@ TNodePtr TSqlQuery::PragmaStatement(const TRule_pragma_stmt& stmt, bool& success
                 return {};
             }
             Ctx.IncrementMonCounter("sql_pragma", "FeatureR010");
+        } else if (normalizedPragma == "compactgroupby") {
+            Ctx.CompactGroupBy = true;
+            Ctx.IncrementMonCounter("sql_pragma", "CompactGroupBy");
+        } else if (normalizedPragma == "disablecompactgroupby") {
+            Ctx.CompactGroupBy = false;
+            Ctx.IncrementMonCounter("sql_pragma", "DisableCompactGroupBy");
         } else {
             Error() << "Unknown pragma: " << pragma;
             Ctx.IncrementMonCounter("sql_errors", "UnknownPragma");
diff --git a/ydb/library/yql/sql/v1/sql_ut.cpp b/ydb/library/yql/sql/v1/sql_ut.cpp
index 12dab539ee..5e51cf21e6 100644
--- a/ydb/library/yql/sql/v1/sql_ut.cpp
+++ b/ydb/library/yql/sql/v1/sql_ut.cpp
@@ -2349,6 +2349,40 @@ Y_UNIT_TEST_SUITE(SqlParsingOnly) {
 
         UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Write"]);
     }
+
+    Y_UNIT_TEST(PragmaCompactGroupBy) {
+        auto req = "PRAGMA CompactGroupBy; SELECT key, COUNT(*) FROM plato.Input GROUP BY key;";
+        auto res = SqlToYql(req);
+        UNIT_ASSERT(res.Root);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "Aggregate") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos, line.find("'('compact)"));
+            }
+        };
+
+        TWordCountHive elementStat = { {TString("Aggregate"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Aggregate"]);
+    }
+
+    Y_UNIT_TEST(PragmaDisableCompactGroupBy) {
+        auto req = "PRAGMA DisableCompactGroupBy; SELECT key, COUNT(*) FROM plato.Input GROUP /*+ compact() */ BY key;";
+        auto res = SqlToYql(req);
+        UNIT_ASSERT(res.Root);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "Aggregate") {
+                UNIT_ASSERT_VALUES_EQUAL(TString::npos, line.find("'('compact)"));
+            }
+        };
+
+        TWordCountHive elementStat = { {TString("Aggregate"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Aggregate"]);
+    }
 }
 
 Y_UNIT_TEST_SUITE(ExternalFunction) {