YQL-19747: Support table aliases

commit_hash:6d67ec1fa5023083debd89aaa99950019ca37c90

11 files changed, 176 insertions, 20 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;
     };
diff --git a/yql/essentials/sql/v1/complete/core/name.h b/yql/essentials/sql/v1/complete/core/name.h
index 3eac47f3871..380fabdf5a4 100644
--- a/yql/essentials/sql/v1/complete/core/name.h
+++ b/yql/essentials/sql/v1/complete/core/name.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <util/generic/string.h>
+#include <util/generic/hash.h>
 
 namespace NSQLComplete {
 
@@ -17,3 +18,10 @@ namespace NSQLComplete {
     };
 
 } // namespace NSQLComplete
+
+template <>
+struct THash<NSQLComplete::TTableId> {
+    inline size_t operator()(const NSQLComplete::TTableId& x) const {
+        return THash<std::tuple<TString, TString>>()(std::tie(x.Cluster, x.Path));
+    }
+};
diff --git a/yql/essentials/sql/v1/complete/name/service/name_service.h b/yql/essentials/sql/v1/complete/name/service/name_service.h
index 5139b27dfa3..090b96c4022 100644
--- a/yql/essentials/sql/v1/complete/name/service/name_service.h
+++ b/yql/essentials/sql/v1/complete/name/service/name_service.h
@@ -70,6 +70,8 @@ namespace NSQLComplete {
         struct TConstraints {
             TVector<TTableId> Tables;
         };
+
+        TTableId Table;
     };
 
     struct TBindingName: TIndentifier {
diff --git a/yql/essentials/sql/v1/complete/name/service/schema/name_service.cpp b/yql/essentials/sql/v1/complete/name/service/schema/name_service.cpp
index 2970f06d7cd..47b8e4c4457 100644
--- a/yql/essentials/sql/v1/complete/name/service/schema/name_service.cpp
+++ b/yql/essentials/sql/v1/complete/name/service/schema/name_service.cpp
@@ -18,7 +18,7 @@ namespace NSQLComplete {
                         .Apply(ToListNameResponse);
                 }
 
-                if (request.Constraints.Column) {
+                if (request.Constraints.Column && !request.Constraints.Column->Tables.empty()) {
                     Y_ENSURE(request.Constraints.Column->Tables.size() == 1, "Not Implemented");
                     TTableId table = request.Constraints.Column->Tables[0];
                     return Schema_
@@ -28,7 +28,9 @@ namespace NSQLComplete {
                             .ColumnPrefix = request.Prefix,
                             .ColumnsLimit = request.Limit,
                         })
-                        .Apply(ToDescribeNameResponse);
+                        .Apply([table = std::move(table)](auto f) {
+                            return ToDescribeNameResponse(std::move(f), std::move(table));
+                        });
                 }
 
                 return NThreading::MakeFuture<TNameResponse>({});
@@ -96,13 +98,16 @@ namespace NSQLComplete {
                 return name;
             }
 
-            static TNameResponse ToDescribeNameResponse(NThreading::TFuture<TDescribeTableResponse> f) {
-                TDescribeTableResponse table = f.ExtractValue();
+            static TNameResponse ToDescribeNameResponse(
+                NThreading::TFuture<TDescribeTableResponse> f,
+                TTableId table) {
+                TDescribeTableResponse info = f.ExtractValue();
 
                 TNameResponse response;
-                for (TString& column : table.Columns) {
+                for (TString& column : info.Columns) {
                     TColumnName name;
                     name.Indentifier = std::move(column);
+                    name.Table = table;
                     response.RankedNames.emplace_back(std::move(name));
                 }
                 return response;
diff --git a/yql/essentials/sql/v1/complete/sql_complete.cpp b/yql/essentials/sql/v1/complete/sql_complete.cpp
index 07e837ed3f4..36aecf06b20 100644
--- a/yql/essentials/sql/v1/complete/sql_complete.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete.cpp
@@ -68,8 +68,8 @@ namespace NSQLComplete {
 
             return MakeUnionNameService(std::move(children), MakeDummyRanking())
                 ->Lookup(std::move(request))
-                .Apply([this, input, context = std::move(context)](auto f) {
-                    return ToCompletion(input, context, f.ExtractValue());
+                .Apply([this, input, context = std::move(context), global = std::move(global)](auto f) {
+                    return ToCompletion(input, std::move(context), global, f.ExtractValue());
                 });
         }
 
@@ -148,7 +148,7 @@ namespace NSQLComplete {
 
             if (context.Column && global.Column) {
                 request.Constraints.Column = TColumnName::TConstraints();
-                request.Constraints.Column->Tables = std::move(global.Column->Tables);
+                request.Constraints.Column->Tables = global.Column->TablesWithAlias(context.Column->Table);
             }
 
             return request;
@@ -157,10 +157,11 @@ namespace NSQLComplete {
         TCompletion ToCompletion(
             TCompletionInput input,
             TLocalSyntaxContext context,
+            const TGlobalContext& global,
             TNameResponse response) const {
             TCompletion completion = {
                 .CompletedToken = GetCompletedToken(input, context.EditRange),
-                .Candidates = Convert(std::move(response.RankedNames), std::move(context)),
+                .Candidates = Convert(std::move(response.RankedNames), std::move(context), global),
             };
 
             if (response.NameHintLength) {
@@ -175,16 +176,33 @@ namespace NSQLComplete {
             return completion;
         }
 
-        static TVector<TCandidate> Convert(TVector<TGenericName> names, TLocalSyntaxContext context) {
+        static TVector<TCandidate> Convert(
+            TVector<TGenericName> names,
+            TLocalSyntaxContext context,
+            const TGlobalContext& global) {
             TVector<TCandidate> candidates;
             candidates.reserve(names.size());
             for (auto& name : names) {
-                candidates.emplace_back(Convert(std::move(name), context));
+                candidates.emplace_back(Convert(std::move(name), context, global));
             }
             return candidates;
         }
 
-        static TCandidate Convert(TGenericName name, TLocalSyntaxContext& context) {
+        // TODO(YQL-19747): extract to a separate file
+        static TCandidate Convert(
+            TGenericName name,
+            TLocalSyntaxContext& context,
+            const TGlobalContext& global) {
+            // TODO(YQL-19747): support multiple aliases for a single table
+            THashMap<TTableId, TString> aliasByTable;
+            global.Column.Transform([&](auto&& column) {
+                aliasByTable.reserve(column.Tables.size());
+                for (const auto& table : column.Tables) {
+                    aliasByTable[table] = table.Alias;
+                }
+                return std::monostate();
+            });
+
             return std::visit([&](auto&& name) -> TCandidate {
                 using T = std::decay_t<decltype(name)>;
 
@@ -241,6 +259,12 @@ namespace NSQLComplete {
                 }
 
                 if constexpr (std::is_base_of_v<TColumnName, T>) {
+                    const TString& alias = aliasByTable.at(name.Table);
+                    if (context.Column->Table.empty() && !alias.empty()) {
+                        name.Indentifier.prepend('.');
+                        name.Indentifier.prepend(alias);
+                    }
+
                     return {ECandidateKind::ColumnName, std::move(name.Indentifier)};
                 }
 
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
index 1f29751d017..f78d2f54cb0 100644
--- a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -1072,6 +1072,27 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
             };
             UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, "USE example; SELECT # FROM `/people`"), expected);
         }
+        {
+            TVector<TCandidate> expected = {
+                {ColumnName, "x.age"},
+                {ColumnName, "x.name"},
+            };
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, "SELECT # FROM example.`/people` AS x"), expected);
+        }
+        { // It is parsed into ``` SELECT x.FROM example.`/people` AS x ```
+            TVector<TCandidate> expected = {};
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, "SELECT x.# FROM example.`/people` AS x"), expected);
+        }
+        {
+            TVector<TCandidate> expected = {
+                {ColumnName, "age"},
+            };
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, "SELECT x.a# FROM example.`/people` AS x"), expected);
+        }
+        {
+            TVector<TCandidate> expected = {};
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, "SELECT y.a# FROM example.`/people` AS x"), expected);
+        }
     }
 
     Y_UNIT_TEST(Typing) {