From a5ff9cacbe74da8fd6e6c560ebf6bf60cd4198a3 Mon Sep 17 00:00:00 2001
From: Vitaly Stoyan <vvvv@ydb.tech>
Date: Fri, 2 Feb 2024 18:07:15 +0000
Subject: [PATCH] init

---
 .../yql/core/issue/protos/issue_id.proto      |  1 +
 ydb/library/yql/core/issue/yql_issue.txt      |  4 +++
 ydb/library/yql/core/services/yql_lineage.cpp | 26 ++++++++++++++++---
 ydb/library/yql/core/services/yql_lineage.h   |  3 ++-
 .../core/services/yql_transform_pipeline.cpp  |  3 +--
 5 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/ydb/library/yql/core/issue/protos/issue_id.proto b/ydb/library/yql/core/issue/protos/issue_id.proto
index be254a95a234..196fafea2053 100644
--- a/ydb/library/yql/core/issue/protos/issue_id.proto
+++ b/ydb/library/yql/core/issue/protos/issue_id.proto
@@ -37,6 +37,7 @@ message TIssuesIds {
         CORE_LEGACY_RANK_FOR_NULLABLE_KEYS = 1109;
         CORE_LEGACY_REGEX_ENGINE = 1110;
         CORE_ALIAS_SHADOWS_COLUMN = 1111;
+        CORE_LINEAGE_INTERNAL_ERROR = 1112;
 
 // core errors
         CORE_GC_NODES_LIMIT_EXCEEDED = 1500;
diff --git a/ydb/library/yql/core/issue/yql_issue.txt b/ydb/library/yql/core/issue/yql_issue.txt
index 675f1ae66991..3db40634cd70 100644
--- a/ydb/library/yql/core/issue/yql_issue.txt
+++ b/ydb/library/yql/core/issue/yql_issue.txt
@@ -651,3 +651,7 @@ ids {
     code: YQL_UNTYPED_STRING_LITERALS
     severity: S_WARNING
 }
+ids {
+    code: CORE_LINEAGE_INTERNAL_ERROR
+    severity: S_WARNING
+}
diff --git a/ydb/library/yql/core/services/yql_lineage.cpp b/ydb/library/yql/core/services/yql_lineage.cpp
index 9b7d9b05be19..bba3df585b51 100644
--- a/ydb/library/yql/core/services/yql_lineage.cpp
+++ b/ydb/library/yql/core/services/yql_lineage.cpp
@@ -12,9 +12,10 @@ namespace {
 
 class TLineageScanner {
 public:
-    TLineageScanner(const TExprNode& root, const TTypeAnnotationContext& ctx)
+    TLineageScanner(const TExprNode& root, const TTypeAnnotationContext& ctx, TExprContext& exprCtx)
         : Root_(root)
         , Ctx_(ctx)
+        , ExprCtx_(exprCtx)
     {}
 
     TString Process() {
@@ -209,7 +210,15 @@ class TLineageScanner {
             return &lineage;
         }
 
-        if (node.IsCallable({"Unordered", "UnorderedSubquery", "Right!", "Skip", "Take", "Sort", "AssumeSorted"})) {
+        if (node.IsCallable({
+            "Unordered", 
+            "UnorderedSubquery", 
+            "Right!", 
+            "Skip", 
+            "Take", 
+            "Sort", 
+            "AssumeSorted", 
+            "SkipNullMembers"})) {
             lineage = *CollectLineage(node.Head());
             return &lineage;
         } else if (node.IsCallable("ExtractMembers")) {
@@ -224,11 +233,19 @@ class TLineageScanner {
             HandleWindow(lineage, node);
         } else if (node.IsCallable("EquiJoin")) {
             HandleEquiJoin(lineage, node);
+        } else {
+            Warning(node, TStringBuilder() << node.Content() << " is not supported");
         }
 
         return &lineage;
     }
 
+    void Warning(const TExprNode& node, const TString& message) {
+        auto issue = TIssue(ExprCtx_.GetPosition(node.Pos()), message);
+        SetIssueCode(EYqlIssueCode::TIssuesIds_EIssueCode_CORE_LINEAGE_INTERNAL_ERROR, issue);
+        ExprCtx_.AddWarning(issue);
+    }
+
     void HandleExtractMembers(TLineage& lineage, const TExprNode& node) {
         auto innerLineage = *CollectLineage(node.Head());
         if (innerLineage.Fields.Defined()) {
@@ -763,6 +780,7 @@ class TLineageScanner {
 private:
     const TExprNode& Root_;
     const TTypeAnnotationContext& Ctx_;
+    TExprContext& ExprCtx_;
     TNodeMap<IDataProvider*> Reads_, Writes_;
     ui32 NextReadId_ = 0;
     ui32 NextWriteId_ = 0;
@@ -773,8 +791,8 @@ class TLineageScanner {
 
 }
 
-TString CalculateLineage(const TExprNode& root, const TTypeAnnotationContext& ctx) {
-    TLineageScanner scanner(root, ctx);
+TString CalculateLineage(const TExprNode& root, const TTypeAnnotationContext& ctx, TExprContext& exprCtx) {
+    TLineageScanner scanner(root, ctx, exprCtx);
     return scanner.Process();
 }
 
diff --git a/ydb/library/yql/core/services/yql_lineage.h b/ydb/library/yql/core/services/yql_lineage.h
index 8bacffa3b9c7..3be1ce2d203a 100644
--- a/ydb/library/yql/core/services/yql_lineage.h
+++ b/ydb/library/yql/core/services/yql_lineage.h
@@ -4,7 +4,8 @@
 namespace NYql {
 
 struct TTypeAnnotationContext;
+struct TExprContext;
 
-TString CalculateLineage(const TExprNode& root, const TTypeAnnotationContext& ctx);
+TString CalculateLineage(const TExprNode& root, const TTypeAnnotationContext& ctx, TExprContext& exprCtx);
 
 }
diff --git a/ydb/library/yql/core/services/yql_transform_pipeline.cpp b/ydb/library/yql/core/services/yql_transform_pipeline.cpp
index bc3dc7ec0c25..36a14f7af7f7 100644
--- a/ydb/library/yql/core/services/yql_transform_pipeline.cpp
+++ b/ydb/library/yql/core/services/yql_transform_pipeline.cpp
@@ -190,9 +190,8 @@ TTransformationPipeline& TTransformationPipeline::AddLineageOptimization(TMaybe<
     Transformers_.push_back(TTransformStage(
         CreateFunctorTransformer(
             [typeCtx = TypeAnnotationContext_, &lineageOut](const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
-                Y_UNUSED(ctx);
                 output = input;
-                lineageOut = CalculateLineage(*input, *typeCtx);
+                lineageOut = CalculateLineage(*input, *typeCtx, ctx);
                 return IGraphTransformer::TStatus::Ok;
             }
         ),