YQL-19747: Support table aliases

commit_hash:6d67ec1fa5023083debd89aaa99950019ca37c90

6 files changed, 104 insertions, 8 deletions

diff --git a/yql/essentials/sql/v1/complete/analysis/global/column.cpp b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
index 7e185697d3f..8eed8a8dfd9 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/column.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
@@ -56,6 +56,26 @@ namespace NSQLComplete {
 
         class TInferenceVisitor: public SQLv1Antlr4BaseVisitor {
         public:
+            std::any visitNamed_single_source(SQLv1::Named_single_sourceContext* ctx) override {
+                SQLv1::Single_sourceContext* singleSource = ctx->single_source();
+                if (singleSource == nullptr) {
+                    return {};
+                }
+
+                std::any any = visit(singleSource);
+                if (!any.has_value()) {
+                    return {};
+                }
+                TColumnContext context = std::move(std::any_cast<TColumnContext>(any));
+
+                TMaybe<TString> alias = GetAlias(ctx);
+                if (alias.Empty()) {
+                    return context;
+                }
+
+                return Renamed(std::move(context), *alias);
+            }
+
             std::any visitTable_ref(SQLv1::Table_refContext* ctx) override {
                 TString cluster = GetId(ctx->cluster_expr()).GetOrElse("");
 
@@ -66,10 +86,27 @@ namespace NSQLComplete {
 
                 return TColumnContext{
                     .Tables = {
-                        {.Cluster = std::move(cluster), .Path = std::move(*path)},
+                        TTableId{std::move(cluster), std::move(*path)},
                     },
                 };
             }
+
+        private:
+            TMaybe<TString> GetAlias(SQLv1::Named_single_sourceContext* ctx) const {
+                TMaybe<TString> alias = GetId(ctx->an_id());
+                alias = alias.Defined() ? alias : GetId(ctx->an_id_as_compat());
+                return alias;
+            }
+
+            TColumnContext Renamed(TColumnContext context, TString alias) {
+                Y_ENSURE(!alias.empty());
+
+                for (TAliased<TTableId>& table : context.Tables) {
+                    table.Alias = alias;
+                }
+
+                return context;
+            }
         };
 
         class TVisitor: public TSQLv1NarrowingVisitor {
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.cpp b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
index 7574b6bdfc0..f73cad926e3 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/global.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
@@ -15,6 +15,15 @@
 
 namespace NSQLComplete {
 
+    TVector<TTableId> TColumnContext::TablesWithAlias(TStringBuf alias) const {
+        if (alias.empty()) {
+            return TVector<TTableId>(Tables.begin(), Tables.end());
+        }
+
+        auto filtered = NFuncTools::Filter([&](const auto& x) { return x.Alias == alias; }, Tables);
+        return TVector<TTableId>(filtered.begin(), filtered.end());
+    }
+
     class TErrorStrategy: public antlr4::DefaultErrorStrategy {
     public:
         antlr4::Token* singleTokenDeletion(antlr4::Parser* /* recognizer */) override {
@@ -117,6 +126,12 @@ namespace NSQLComplete {
 } // namespace NSQLComplete
 
 template <>
+void Out<NSQLComplete::TAliased<NSQLComplete::TTableId>>(IOutputStream& out, const NSQLComplete::TAliased<NSQLComplete::TTableId>& value) {
+    Out<NSQLComplete::TTableId>(out, value);
+    out << " AS " << value.Alias;
+}
+
+template <>
 void Out<NSQLComplete::TColumnContext>(IOutputStream& out, const NSQLComplete::TColumnContext& value) {
     out << "TColumnContext { ";
     out << "Tables: " << JoinSeq(", ", value.Tables);
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.h b/yql/essentials/sql/v1/complete/analysis/global/global.h
index 1ef1344e3c9..69eab86049d 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/global.h
+++ b/yql/essentials/sql/v1/complete/analysis/global/global.h
@@ -8,6 +8,7 @@
 #include <util/generic/maybe.h>
 #include <util/generic/string.h>
 #include <util/generic/vector.h>
+#include <util/generic/hash.h>
 
 namespace NSQLComplete {
 
@@ -16,8 +17,28 @@ namespace NSQLComplete {
         TString Cluster;
     };
 
+    template <std::regular T>
+    struct TAliased: T {
+        TString Alias;
+
+        TAliased(TString alias, T value)
+            : T(std::move(value))
+            , Alias(std::move(alias))
+        {
+        }
+
+        TAliased(T value)
+            : T(std::move(value))
+        {
+        }
+
+        friend bool operator==(const TAliased& lhs, const TAliased& rhs) = default;
+    };
+
     struct TColumnContext {
-        TVector<TTableId> Tables;
+        TVector<TAliased<TTableId>> Tables;
+
+        TVector<TTableId> TablesWithAlias(TStringBuf alias) const;
 
         friend bool operator==(const TColumnContext& lhs, const TColumnContext& rhs) = default;
     };
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp b/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp
index c8222580cce..d8b19786c11 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp
@@ -105,7 +105,7 @@ Y_UNIT_TEST_SUITE(GlobalAnalysisTests) {
 
             TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
 
-            TColumnContext expected = {.Tables = {{"plato", "Input"}}};
+            TColumnContext expected = {.Tables = {TTableId{"plato", "Input"}}};
             UNIT_ASSERT_VALUES_EQUAL(ctx.Column, expected);
         }
         {
@@ -113,7 +113,15 @@ Y_UNIT_TEST_SUITE(GlobalAnalysisTests) {
 
             TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
 
-            TColumnContext expected = {.Tables = {{"plato", "//home/input"}}};
+            TColumnContext expected = {.Tables = {TTableId{"plato", "//home/input"}}};
+            UNIT_ASSERT_VALUES_EQUAL(ctx.Column, expected);
+        }
+        {
+            TString query = "SELECT # FROM plato.Input AS x";
+
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+
+            TColumnContext expected = {.Tables = {TAliased<TTableId>("x", TTableId{"plato", "Input"})}};
             UNIT_ASSERT_VALUES_EQUAL(ctx.Column, expected);
         }
     }
diff --git a/yql/essentials/sql/v1/complete/analysis/local/local.cpp b/yql/essentials/sql/v1/complete/analysis/local/local.cpp
index 13e2fce69ad..014ac61e803 100644
--- a/yql/essentials/sql/v1/complete/analysis/local/local.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/local/local.cpp
@@ -107,7 +107,7 @@ namespace NSQLComplete {
             result.Hint = HintMatch(candidates);
             result.Object = ObjectMatch(context, candidates);
             result.Cluster = ClusterMatch(context, candidates);
-            result.Column = ColumnMatch(candidates);
+            result.Column = ColumnMatch(context, candidates);
             result.Binding = BindingMatch(candidates);
 
             return result;
@@ -322,8 +322,19 @@ namespace NSQLComplete {
             return cluster;
         }
 
-        bool ColumnMatch(const TC3Candidates& candidates) const {
-            return AnyOf(candidates.Rules, RuleAdapted(IsLikelyColumnStack));
+        TMaybe<TLocalSyntaxContext::TColumn> ColumnMatch(
+            const TCursorTokenContext& context, const TC3Candidates& candidates) const {
+            if (!AnyOf(candidates.Rules, RuleAdapted(IsLikelyColumnStack))) {
+                return Nothing();
+            }
+
+            TLocalSyntaxContext::TColumn column;
+            if (TMaybe<TRichParsedToken> begin;
+                (begin = context.MatchCursorPrefix({"ID_PLAIN", "DOT"})) ||
+                (begin = context.MatchCursorPrefix({"ID_PLAIN", "DOT", ""}))) {
+                column.Table = begin->Base->Content;
+            }
+            return column;
         }
 
         bool BindingMatch(const TC3Candidates& candidates) const {
diff --git a/yql/essentials/sql/v1/complete/analysis/local/local.h b/yql/essentials/sql/v1/complete/analysis/local/local.h
index cca182aacdb..c5d73d52018 100644
--- a/yql/essentials/sql/v1/complete/analysis/local/local.h
+++ b/yql/essentials/sql/v1/complete/analysis/local/local.h
@@ -48,6 +48,10 @@ namespace NSQLComplete {
             }
         };
 
+        struct TColumn {
+            TString Table;
+        };
+
         TKeywords Keywords;
         TMaybe<TPragma> Pragma;
         bool Type = false;
@@ -55,7 +59,7 @@ namespace NSQLComplete {
         TMaybe<THint> Hint;
         TMaybe<TObject> Object;
         TMaybe<TCluster> Cluster;
-        bool Column = false;
+        TMaybe<TColumn> Column;
         bool Binding = false;
         TEditRange EditRange;
     };