Allow DISCARD in SUBQUERY

6 files changed, 49 insertions, 10 deletions

diff --git a/ydb/library/yql/sql/v1/node.cpp b/ydb/library/yql/sql/v1/node.cpp
index 04543e2cd7..b7cc07e3de 100644
--- a/ydb/library/yql/sql/v1/node.cpp
+++ b/ydb/library/yql/sql/v1/node.cpp
@@ -278,6 +278,10 @@ bool INode::IsSelect() const {
     return false;
 }
 
+bool INode::HasSelectResult() const {
+    return false;
+}
+
 const TString* INode::FuncName() const {
     return nullptr;
 }
diff --git a/ydb/library/yql/sql/v1/node.h b/ydb/library/yql/sql/v1/node.h
index 15d28dc9c7..6d540b8253 100644
--- a/ydb/library/yql/sql/v1/node.h
+++ b/ydb/library/yql/sql/v1/node.h
@@ -174,6 +174,7 @@ namespace NSQLTranslationV1 {
         void UseAsInner();
         virtual bool UsedSubquery() const;
         virtual bool IsSelect() const;
+        virtual bool HasSelectResult() const;
         virtual const TString* FuncName() const;
         virtual const TString* ModuleName() const;
         virtual bool HasSkip() const;
diff --git a/ydb/library/yql/sql/v1/select.cpp b/ydb/library/yql/sql/v1/select.cpp
index 28228731d3..f90ad44beb 100644
--- a/ydb/library/yql/sql/v1/select.cpp
+++ b/ydb/library/yql/sql/v1/select.cpp
@@ -1164,6 +1164,10 @@ public:
         return Settings;
     }
 
+    bool HasSelectResult() const final {
+        return !Settings.Discard;
+    }
+
     TPtr DoClone() const final {
         return new TReduceSource(Pos, Mode, Source->CloneSource(), CloneContainer(OrderBy),
             CloneContainer(Keys), CloneContainer(Args), SafeClone(Udf), SafeClone(Having), Settings,
@@ -1380,6 +1384,10 @@ public:
         return Settings;
     }
 
+    bool HasSelectResult() const override {
+        return !Settings.Discard;
+    }
+
     TNodePtr DoClone() const final {
         auto newSource = MakeIntrusive<TCompositeSelect>(Pos, Source->CloneSource(), OriginalSource->CloneSource(), Settings);
         newSource->SetSubselects(CloneContainer(Subselects), CloneContainer(Grouping), CloneContainer(GroupByExpr));
@@ -1752,6 +1760,10 @@ public:
         return true;
     }
 
+    bool HasSelectResult() const override {
+        return !Settings.Discard;
+    }
+
     bool IsStream() const override {
         return Source->IsStream();
     }
@@ -2358,6 +2370,10 @@ public:
         return false;
     }
 
+    bool HasSelectResult() const override {
+        return !Settings.Discard;
+    }
+
     bool IsStream() const override {
         return Source->IsStream();
     }
@@ -2662,6 +2678,10 @@ public:
         return true;
     }
 
+    bool HasSelectResult() const override {
+        return !Settings.Discard;
+    }
+
     TWriteSettings GetWriteSettings() const override {
         return Settings;
     }
@@ -2856,6 +2876,10 @@ public:
         return Source->IsSelect();
     }
 
+    bool HasSelectResult() const override {
+        return Source->HasSelectResult();
+    }
+
     TPtr DoClone() const final {
         return MakeIntrusive<TSelect>(Pos, Source->CloneSource(), SafeClone(SkipTake));
     }
@@ -2891,6 +2915,10 @@ public:
         return true;
     }
 
+    bool HasSelectResult() const override {
+        return Source->HasSelectResult();
+    }
+
     bool DoInit(TContext& ctx, ISource* src) override {
         if (!Source->Init(ctx, src)) {
             return false;
@@ -2959,7 +2987,7 @@ public:
             return false;
         }
         node = L(node, Y("let", "output", output));
-        if (WriteResult) {
+        if (WriteResult || writeSettings.Discard) {
             if (EOrderKind::None == Source->GetOrderKind() && ctx.UseUnordered(*Source)) {
                 node = L(node, Y("let", "output", Y("Unordered", "output")));
             }
diff --git a/ydb/library/yql/sql/v1/sql_query.cpp b/ydb/library/yql/sql/v1/sql_query.cpp
index c842be2757..9ed305223e 100644
--- a/ydb/library/yql/sql/v1/sql_query.cpp
+++ b/ydb/library/yql/sql/v1/sql_query.cpp
@@ -2181,19 +2181,22 @@ TNodePtr TSqlQuery::Build(const TSQLv1ParserAST& ast) {
     }
 
     ui32 topLevelSelects = 0;
+    bool hasTailOps = false;
     for (auto& block : blocks) {
-        if (block->IsSelect()) {
+        if (block->HasSelectResult()) {
             ++topLevelSelects;
+        } else if (topLevelSelects) {
+            hasTailOps = true;
         }
     }
 
-    if ((Mode == NSQLTranslation::ESqlMode::SUBQUERY || Mode == NSQLTranslation::ESqlMode::LIMITED_VIEW) && topLevelSelects != 1) {
-        Error() << "Strictly one select/process/reduce statement must be used in the "
+    if ((Mode == NSQLTranslation::ESqlMode::SUBQUERY || Mode == NSQLTranslation::ESqlMode::LIMITED_VIEW) && (topLevelSelects != 1 || hasTailOps)) {
+        Error() << "Strictly one select/process/reduce statement is expected at the end of "
             << (Mode == NSQLTranslation::ESqlMode::LIMITED_VIEW ? "view" : "subquery");
         return nullptr;
     }
 
-    if (!Ctx.PragmaAutoCommit && Ctx.Settings.EndOfQueryCommit && IsQueryMode(Mode)) {
+     if (!Ctx.PragmaAutoCommit && Ctx.Settings.EndOfQueryCommit && IsQueryMode(Mode)) {
         AddStatementToBlocks(blocks, BuildCommitClusters(Ctx.Pos()));
     }
 
diff --git a/ydb/library/yql/sql/v1/sql_select.cpp b/ydb/library/yql/sql/v1/sql_select.cpp
index b79d72ac10..551ed619a1 100644
--- a/ydb/library/yql/sql/v1/sql_select.cpp
+++ b/ydb/library/yql/sql/v1/sql_select.cpp
@@ -1211,8 +1211,8 @@ TSqlSelect::TSelectKindResult TSqlSelect::SelectKind(const TRule_select_kind& no
 {
     const bool discard = node.HasBlock1();
     const bool hasLabel = node.HasBlock3();
-    if ((discard || hasLabel) && (Mode == NSQLTranslation::ESqlMode::LIMITED_VIEW || Mode == NSQLTranslation::ESqlMode::SUBQUERY)) {
-        Ctx.Error() << "DISCARD and INTO RESULT are not allowed in current mode";
+    if (hasLabel && (Mode == NSQLTranslation::ESqlMode::LIMITED_VIEW || Mode == NSQLTranslation::ESqlMode::SUBQUERY)) {
+        Ctx.Error() << "INTO RESULT is not allowed in current mode";
         return {};
     }
 
diff --git a/ydb/library/yql/sql/v1/sql_translation.cpp b/ydb/library/yql/sql/v1/sql_translation.cpp
index edd313801d..a366e587ca 100644
--- a/ydb/library/yql/sql/v1/sql_translation.cpp
+++ b/ydb/library/yql/sql/v1/sql_translation.cpp
@@ -3970,14 +3970,17 @@ bool TSqlTranslation::DefineActionOrSubqueryStatement(const TRule_define_action_
     const bool hasValidBody = DefineActionOrSubqueryBody(query, innerBlocks, stmt.GetRule_define_action_or_subquery_body8());
 
     ui32 topLevelSelects = 0;
+    bool hasTailOps = false;
     for (auto& block : innerBlocks) {
-        if (block->IsSelect()) {
+        if (block->HasSelectResult()) {
             ++topLevelSelects;
+        } else if (topLevelSelects) {
+            hasTailOps = true;
         }
     }
 
-    if (isSubquery && topLevelSelects != 1) {
-        Error() << "Strictly one select/process/reduce statement must be used in the subquery, got: " << topLevelSelects;
+    if (isSubquery && (topLevelSelects != 1 || hasTailOps)) {
+        Error() << "Strictly one select/process/reduce statement is expected at the end of subquery";
         return false;
     }