YQL-19264 relaxed check for assume order by

commit_hash:29ee6c53b63771357492cc2e4f7ceac15c5a86b1

6 files changed, 48 insertions, 2 deletions

diff --git a/yql/essentials/sql/v1/builtin.cpp b/yql/essentials/sql/v1/builtin.cpp
index 94d8a3bc16d..5d33b2061fa 100644
--- a/yql/essentials/sql/v1/builtin.cpp
+++ b/yql/essentials/sql/v1/builtin.cpp
@@ -1907,6 +1907,10 @@ public:
         return new TTableRow<Join>(Pos, ArgsCount);
     }
 
+    bool IsTableRow() const final {
+        return true;
+    }
+
 private:
     const size_t ArgsCount;
     TNodePtr Node;
diff --git a/yql/essentials/sql/v1/node.cpp b/yql/essentials/sql/v1/node.cpp
index b2851429098..31d72385ca7 100644
--- a/yql/essentials/sql/v1/node.cpp
+++ b/yql/essentials/sql/v1/node.cpp
@@ -172,6 +172,14 @@ const TString* INode::GetColumnName() const {
     return nullptr;
 }
 
+bool INode::IsPlainColumn() const {
+    return GetColumnName() != nullptr;
+}
+
+bool INode::IsTableRow() const {
+    return false;
+}
+
 void INode::AssumeColumn() {
 }
 
@@ -466,6 +474,14 @@ const TString* IProxyNode::GetColumnName() const {
     return Inner->GetColumnName();
 }
 
+bool IProxyNode::IsPlainColumn() const {
+    return Inner->IsPlainColumn();
+}
+
+bool IProxyNode::IsTableRow() const {
+    return Inner->IsTableRow();
+}
+
 void IProxyNode::AssumeColumn() {
     Inner->AssumeColumn();
 }
@@ -2480,6 +2496,18 @@ public:
         return ColumnOnly ? Ids[0].Expr->GetColumnName() : nullptr;
     }
 
+    bool IsPlainColumn() const override {
+        if (GetColumnName()) {
+            return true;
+        }
+
+        if (Ids[0].Expr->IsTableRow()) {
+            return true;
+        }
+
+        return false;
+    }
+
     const TString* GetSourceName() const override {
         return Ids[0].Expr->GetSourceName();
     }
diff --git a/yql/essentials/sql/v1/node.h b/yql/essentials/sql/v1/node.h
index 99f42f353cd..c2c46c60775 100644
--- a/yql/essentials/sql/v1/node.h
+++ b/yql/essentials/sql/v1/node.h
@@ -168,6 +168,8 @@ namespace NSQLTranslationV1 {
         virtual TString GetOpName() const;
         virtual const TString* GetLiteral(const TString& type) const;
         virtual const TString* GetColumnName() const;
+        virtual bool IsPlainColumn() const;
+        virtual bool IsTableRow() const;
         virtual void AssumeColumn();
         virtual const TString* GetSourceName() const;
         virtual const TString* GetAtomContent() const;
@@ -299,6 +301,8 @@ namespace NSQLTranslationV1 {
         virtual TString GetOpName() const override;
         virtual const TString* GetLiteral(const TString &type) const override;
         virtual const TString* GetColumnName() const override;
+        virtual bool IsPlainColumn() const override;
+        virtual bool IsTableRow() const override;
         virtual void AssumeColumn() override;
         virtual const TString* GetSourceName() const override;
         virtual const TString* GetAtomContent() const override;
diff --git a/yql/essentials/sql/v1/select.cpp b/yql/essentials/sql/v1/select.cpp
index 4eee1efb266..552a74cb07a 100644
--- a/yql/essentials/sql/v1/select.cpp
+++ b/yql/essentials/sql/v1/select.cpp
@@ -912,7 +912,7 @@ TSourcePtr BuildInnerSource(TPosition pos, TNodePtr node, const TString& service
 }
 
 static bool IsComparableExpression(TContext& ctx, const TNodePtr& expr, bool assume, const char* sqlConstruction) {
-    if (assume && !expr->GetColumnName()) {
+    if (assume && !expr->IsPlainColumn()) {
         ctx.Error(expr->GetPos()) << "Only column names can be used in " << sqlConstruction;
         return false;
     }
@@ -925,7 +925,7 @@ static bool IsComparableExpression(TContext& ctx, const TNodePtr& expr, bool ass
         ctx.Error(expr->GetPos()) << "Unable to " << sqlConstruction << " aggregated values";
         return false;
     }
-    if (expr->GetColumnName()) {
+    if (expr->IsPlainColumn()) {
         return true;
     }
     if (expr->GetOpName().empty()) {
diff --git a/yql/essentials/sql/v1/sql_ut.cpp b/yql/essentials/sql/v1/sql_ut.cpp
index 62d5d2c83ad..2acc9e722a2 100644
--- a/yql/essentials/sql/v1/sql_ut.cpp
+++ b/yql/essentials/sql/v1/sql_ut.cpp
@@ -630,6 +630,11 @@ Y_UNIT_TEST_SUITE(SqlParsingOnly) {
         UNIT_ASSERT(res.Root);
     }
 
+    Y_UNIT_TEST(SelectAssumeOrderByTableRowAccess) {
+        NYql::TAstParseResult res = SqlToYql("$key = 'foo';select * from plato.Input assume order by TableRow().$key");
+        UNIT_ASSERT(res.Root);
+    }
+
     Y_UNIT_TEST(SelectOrderByDuplicateLabels) {
         NYql::TAstParseResult res = SqlToYql("select a from plato.Input order by a, a");
         UNIT_ASSERT(res.Root);
diff --git a/yql/essentials/sql/v1/sql_ut_antlr4.cpp b/yql/essentials/sql/v1/sql_ut_antlr4.cpp
index 202ac513916..1cd05a7ea5a 100644
--- a/yql/essentials/sql/v1/sql_ut_antlr4.cpp
+++ b/yql/essentials/sql/v1/sql_ut_antlr4.cpp
@@ -630,6 +630,11 @@ Y_UNIT_TEST_SUITE(SqlParsingOnly) {
         UNIT_ASSERT(res.Root);
     }
 
+    Y_UNIT_TEST(SelectAssumeOrderByTableRowAccess) {
+        NYql::TAstParseResult res = SqlToYql("$key = 'foo';select * from plato.Input assume order by TableRow().$key");
+        UNIT_ASSERT(res.Root);
+    }
+
     Y_UNIT_TEST(SelectOrderByDuplicateLabels) {
         NYql::TAstParseResult res = SqlToYql("select a from plato.Input order by a, a");
         UNIT_ASSERT(res.Root);