YQL-19747: Complete columns from JOIN

vityaman · vityaman · commit b4ef2d519178 · 2025-06-23T17:48:51.000+03:00
Collect tables with aliases during global
analysis and send a multiple requests
to a schema provider.
commit_hash:51cbc40cee6cda09d5bfe256a0425e7c809c04f5
diff --git a/yql/essentials/sql/v1/complete/analysis/global/column.cpp b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
@@ -56,6 +56,20 @@ namespace NSQLComplete {
 
         class TInferenceVisitor: public SQLv1Antlr4BaseVisitor {
         public:
+            std::any visitJoin_source(SQLv1::Join_sourceContext* ctx) override {
+                TColumnContext united;
+                for (SQLv1::Flatten_sourceContext* ctx : ctx->flatten_source()) {
+                    std::any any = visit(ctx);
+                    if (!any.has_value()) {
+                        continue;
+                    }
+
+                    TColumnContext child = std::move(std::any_cast<TColumnContext>(any));
+                    united = United(std::move(united), std::move(child));
+                }
+                return united;
+            }
+
             std::any visitNamed_single_source(SQLv1::Named_single_sourceContext* ctx) override {
                 SQLv1::Single_sourceContext* singleSource = ctx->single_source();
                 if (singleSource == nullptr) {
@@ -107,6 +121,13 @@ namespace NSQLComplete {
 
                 return context;
             }
+
+            static TColumnContext United(TColumnContext&& lhs, TColumnContext&& rhs) {
+                lhs.Tables.reserve(lhs.Tables.size() + rhs.Tables.size());
+                std::move(rhs.Tables.begin(), rhs.Tables.end(), std::back_inserter(lhs.Tables));
+                SortUnique(lhs.Tables);
+                return lhs;
+            }
         };
 
         class TVisitor: public TSQLv1NarrowingVisitor {
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.h b/yql/essentials/sql/v1/complete/analysis/global/global.h
@@ -17,7 +17,9 @@ namespace NSQLComplete {
         TString Cluster;
     };
 
-    template <std::regular T>
+    template <class T>
+        requires std::regular<T> &&
+                 requires(T x) { {x < x} -> std::convertible_to<bool>; }
     struct TAliased: T {
         TString Alias;
 
@@ -32,6 +34,10 @@ namespace NSQLComplete {
         {
         }
 
+        friend bool operator<(const TAliased& lhs, const TAliased& rhs) {
+            return std::tie(lhs.Alias, static_cast<const T&>(lhs)) < std::tie(rhs.Alias, static_cast<const T&>(rhs));
+        }
+
         friend bool operator==(const TAliased& lhs, const TAliased& 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
@@ -126,4 +126,27 @@ Y_UNIT_TEST_SUITE(GlobalAnalysisTests) {
         }
     }
 
+    Y_UNIT_TEST(Join) {
+        IGlobalAnalysis::TPtr global = MakeGlobalAnalysis();
+        {
+            TString query = R"(
+                SELECT #
+                FROM q.a AS x, p.b, c
+                JOIN p.d AS y ON x.key = y.key
+            )";
+
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+
+            TColumnContext expected = {
+                .Tables = {
+                    TAliased<TTableId>("", {"", "c"}),
+                    TAliased<TTableId>("", {"p", "b"}),
+                    TAliased<TTableId>("x", {"q", "a"}),
+                    TAliased<TTableId>("y", {"p", "d"}),
+                },
+            };
+            UNIT_ASSERT_VALUES_EQUAL(ctx.Column, expected);
+        }
+    }
+
 } // Y_UNIT_TEST_SUITE(GlobalAnalysisTests)
diff --git a/yql/essentials/sql/v1/complete/core/name.cpp b/yql/essentials/sql/v1/complete/core/name.cpp
@@ -2,6 +2,14 @@
 
 #include <util/stream/output.h>
 
+namespace NSQLComplete {
+
+    bool operator<(const TTableId& lhs, const TTableId& rhs) {
+        return std::tie(lhs.Cluster, lhs.Path) < std::tie(rhs.Cluster, rhs.Path);
+    }
+
+} // namespace NSQLComplete
+
 template <>
 void Out<NSQLComplete::TTableId>(IOutputStream& out, const NSQLComplete::TTableId& value) {
     out << value.Cluster << ".`" << value.Path << "`";
diff --git a/yql/essentials/sql/v1/complete/core/name.h b/yql/essentials/sql/v1/complete/core/name.h
@@ -14,6 +14,7 @@ namespace NSQLComplete {
         TString Cluster;
         TString Path;
 
+        friend bool operator<(const TTableId& lhs, const TTableId& rhs);
         friend bool operator==(const TTableId& lhs, const TTableId& rhs) = default;
     };
 
diff --git a/yql/essentials/sql/v1/complete/name/service/name_service.h b/yql/essentials/sql/v1/complete/name/service/name_service.h
@@ -12,6 +12,7 @@
 
 namespace NSQLComplete {
 
+    // TODO(YQL-19747): Rename to Identifier
     struct TIndentifier {
         TString Indentifier;
     };
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
@@ -1,5 +1,8 @@
 #include "name_service.h"
 
+#include <library/cpp/threading/future/wait/wait.h>
+#include <library/cpp/iterator/functools.h>
+
 namespace NSQLComplete {
 
     namespace {
@@ -19,24 +22,47 @@ namespace NSQLComplete {
                 }
 
                 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_
-                        ->Describe({
-                            .TableCluster = table.Cluster,
-                            .TablePath = table.Path,
-                            .ColumnPrefix = request.Prefix,
-                            .ColumnsLimit = request.Limit,
-                        })
-                        .Apply([table = std::move(table)](auto f) {
-                            return ToDescribeNameResponse(std::move(f), std::move(table));
-                        });
+                    return BatchDescribe(
+                        std::move(request.Constraints.Column->Tables),
+                        request.Prefix,
+                        request.Limit);
                 }
 
                 return NThreading::MakeFuture<TNameResponse>({});
             }
 
         private:
+            NThreading::TFuture<TNameResponse> BatchDescribe(
+                TVector<TTableId> tables, TString prefix, ui64 limit) const {
+                TVector<NThreading::TFuture<TDescribeTableResponse>> futures;
+                for (const auto& table : tables) {
+                    TDescribeTableRequest request = {
+                        .TableCluster = table.Cluster,
+                        .TablePath = table.Path,
+                        .ColumnPrefix = prefix,
+                        .ColumnsLimit = limit,
+                    };
+                    futures.emplace_back(Schema_->Describe(request));
+                }
+
+                return NThreading::WaitAll(futures).Apply([tables, futures](auto) mutable {
+                    TNameResponse response;
+
+                    for (auto [table, f] : NFuncTools::Zip(tables, futures)) {
+                        TDescribeTableResponse description = f.ExtractValue();
+                        for (TString& column : description.Columns) {
+                            TColumnName name;
+                            name.Indentifier = std::move(column);
+                            name.Table = table;
+
+                            response.RankedNames.emplace_back(std::move(name));
+                        }
+                    }
+
+                    return response;
+                });
+            }
+
             static TListRequest ToListRequest(TNameRequest request) {
                 return {
                     .Cluster = ClusterName(*request.Constraints.Object),
@@ -98,21 +124,6 @@ namespace NSQLComplete {
                 return name;
             }
 
-            static TNameResponse ToDescribeNameResponse(
-                NThreading::TFuture<TDescribeTableResponse> f,
-                TTableId table) {
-                TDescribeTableResponse info = f.ExtractValue();
-
-                TNameResponse response;
-                for (TString& column : info.Columns) {
-                    TColumnName name;
-                    name.Indentifier = std::move(column);
-                    name.Table = table;
-                    response.RankedNames.emplace_back(std::move(name));
-                }
-                return response;
-            }
-
             ISchema::TPtr Schema_;
         };
 
diff --git a/yql/essentials/sql/v1/complete/name/service/schema/ya.make b/yql/essentials/sql/v1/complete/name/service/schema/ya.make
@@ -7,6 +7,7 @@ SRCS(
 PEERDIR(
     yql/essentials/sql/v1/complete/name/object
     yql/essentials/sql/v1/complete/name/service
+    library/cpp/iterator
 )
 
 END()
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -113,7 +113,11 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
                     "Age": {}
                 }},
                 "yql": { "type": "Folder", "entries": {
-                    "tutorial": { "type": "Table", "columns": {} }
+                    "tutorial": { "type": "Table", "columns": {
+                       "course": {},
+                       "room": {},
+                       "time": {}
+                    }}
                 }}
             }},
             "saurus": { "type": "Folder", "entries": {
@@ -1117,6 +1121,26 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         }
     }
 
+    Y_UNIT_TEST(ColumnsAtJoin) {
+        auto engine = MakeSqlCompletionEngineUT();
+        {
+            TString query = R"(
+                SELECT #
+                FROM example.`/people` AS ep
+                JOIN example.`/yql/tutorial` AS et ON 1 = 1
+            )";
+
+            TVector<TCandidate> expected = {
+                {ColumnName, "ep.Age"},
+                {ColumnName, "ep.Name"},
+                {ColumnName, "et.course"},
+                {ColumnName, "et.room"},
+                {ColumnName, "et.time"},
+            };
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(5, engine, query), expected);
+        }
+    }
+
     Y_UNIT_TEST(Typing) {
         const auto queryUtf16 = TUtf16String::FromUtf8(
             "SELECT \n"

Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@ SRCS(`
`7`	`7`	`PEERDIR(`
`8`	`8`	`yql/essentials/sql/v1/complete/name/object`
`9`	`9`	`yql/essentials/sql/v1/complete/name/service`
	`10`	`+ library/cpp/iterator`
`10`	`11`	`)`
`11`	`12`
`12`	`13`	`END()`