Merge pull request #27203 from ydb-platform/merge-rightlib-251021-0051HEADmain

4 files changed, 89 insertions, 27 deletions

diff --git a/yql/essentials/core/services/yql_lineage.cpp b/yql/essentials/core/services/yql_lineage.cpp
index 57e3e02192f..35721515add 100644
--- a/yql/essentials/core/services/yql_lineage.cpp
+++ b/yql/essentials/core/services/yql_lineage.cpp
@@ -5,6 +5,7 @@
 #include <yql/essentials/core/yql_opt_utils.h>
 #include <yql/essentials/core/yql_join.h>
 
+#include <library/cpp/yson/node/node_io.h>
 #include <util/system/env.h>
 
 namespace NYql {
@@ -966,4 +967,26 @@ TString CalculateLineage(const TExprNode& root, const TTypeAnnotationContext& ct
     return scanner.Process();
 }
 
+TString NormalizeLineage(const TString& lineageStr) {
+    THashMap<i64, TString> idToPath;
+    auto lineageNode = NYT::NodeFromYsonString(lineageStr);
+    auto& readsSection = lineageNode.AsMap()["Reads"];
+    for (auto& readNode : readsSection.AsList()) {
+        auto& readMap = readNode.AsMap();
+        idToPath[readMap["Id"].AsInt64()] = readMap["Name"].AsString();
+        readMap["Id"] = readMap["Name"];
+    }
+    auto& writesSection = lineageNode.AsMap()["Writes"];
+    for (auto& writeNode : writesSection.AsList()) {
+        auto& writeMap = writeNode.AsMap();
+        writeMap["Id"] = writeMap["Name"];
+        for (auto& [fieldName, fieldLineage] : writeMap["Lineage"].AsMap()) {
+            for (auto& inputField : fieldLineage.AsList()) {
+                inputField["Input"] = idToPath[inputField["Input"].AsInt64()];
+            }
+        }
+    }
+    return NYT::NodeToCanonicalYsonString(lineageNode);
+}
+
 } // namespace NYql
diff --git a/yql/essentials/core/services/yql_lineage.h b/yql/essentials/core/services/yql_lineage.h
index c1a71ae185b..2b66c064a4d 100644
--- a/yql/essentials/core/services/yql_lineage.h
+++ b/yql/essentials/core/services/yql_lineage.h
@@ -8,4 +8,7 @@ struct TExprContext;
 
 TString CalculateLineage(const TExprNode& root, const TTypeAnnotationContext& ctx, TExprContext& exprCtx, bool standalone);
 
+// Replace input and output table's IDs with pathes for checking lineage equality
+TString NormalizeLineage(const TString& lineageStr);
+
 } // namespace NYql
diff --git a/yql/essentials/core/services/yql_transform_pipeline.cpp b/yql/essentials/core/services/yql_transform_pipeline.cpp
index 1b2df336500..45c0a5ae6a8 100644
--- a/yql/essentials/core/services/yql_transform_pipeline.cpp
+++ b/yql/essentials/core/services/yql_transform_pipeline.cpp
@@ -21,6 +21,9 @@
 
 namespace NYql {
 
+const TString LineageComponent = "Lineage";
+const TString LineageResultLabel = "LineageResult";
+
 TTransformationPipeline::TTransformationPipeline(
     TIntrusivePtr<TTypeAnnotationContext> ctx,
     TTypeAnnCallableFactory typeAnnCallableFactory)
@@ -167,33 +170,66 @@ TTransformationPipeline& TTransformationPipeline::AddFinalCommonOptimization(EYq
     return *this;
 }
 
-TTransformationPipeline& TTransformationPipeline::AddOptimizationWithLineage(bool checkWorld, bool withFinalOptimization, EYqlIssueCode issueCode) {
+TTransformationPipeline& TTransformationPipeline::AddOptimizationWithLineage(bool enableLineage, bool checkWorld, bool withFinalOptimization, EYqlIssueCode issueCode) {
     AddCommonOptimization(false, issueCode);
-    Transformers_.push_back(TTransformStage(
-        CreateChoiceGraphTransformer(
-            [&typesCtx = std::as_const(*TypeAnnotationContext_)](const TExprNode::TPtr&, TExprContext&) {
-                return typesCtx.EnableLineage;
-            },
-            TTransformStage(
-                CreateFunctorTransformer(
-                    [typeCtx = TypeAnnotationContext_](const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
-                        output = input;
-                        try {
-                            CalculateLineage(*input, *typeCtx, ctx, false);
-                        } catch (const std::exception& e) {
-                            YQL_LOG(ERROR) << "CalculateLineage error: " << e.what();
-                            typeCtx->CorrectLineage = false;
-                        }
-                        return IGraphTransformer::TStatus::Ok;
-                    }),
-                "Lineage",
-                issueCode),
-            TTransformStage(
-                new TNullTransformer(),
-                "SkipLineage",
-                issueCode)),
-        "LineageCalculation",
-        issueCode));
+    if (enableLineage) {
+        Transformers_.push_back(TTransformStage(
+            CreateChoiceGraphTransformer(
+                [&typesCtx = std::as_const(*TypeAnnotationContext_)](const TExprNode::TPtr&, TExprContext&) {
+                    return typesCtx.EnableLineage;
+                },
+                TTransformStage(
+                    CreateSinglePassFunctorTransformer(
+                        [typeCtx = TypeAnnotationContext_](const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
+                            output = input;
+                            TString calculatedLineage, loadedLineage;
+                            if (typeCtx->QContext && typeCtx->QContext.CanRead()) {
+                                auto loaded = typeCtx->QContext.GetReader()->Get({LineageComponent, LineageResultLabel}).GetValueSync();
+                                if (loaded.Defined()) {
+                                    loadedLineage = loaded->Value;
+                                } else {
+                                    YQL_LOG(INFO) << "There is no lineage in QStorage, lineage calculation is skipped in replay mode";
+                                    return IGraphTransformer::TStatus::Ok;
+                                }
+                            }
+                            std::exception_ptr lineageError;
+                            try {
+                                calculatedLineage = CalculateLineage(*input, *typeCtx, ctx, false);
+                            } catch (const std::exception& e) {
+                                YQL_LOG(ERROR) << "Lineage calculation error: " << e.what();
+                                typeCtx->CorrectLineage = false;
+                                lineageError = std::current_exception();
+                            }
+                            if (!loadedLineage.empty()) {
+                                // if lineage calculation is failed, but loaded lineage exists, rethrow exception for replay mode
+                                if (lineageError) {
+                                    std::rethrow_exception(lineageError);
+                                }
+                                if (NormalizeLineage(calculatedLineage) != NormalizeLineage(loadedLineage)) {
+                                    YQL_LOG(INFO) << "Lineage in replay is different:"
+                                                  << "\nCalculated lineage:\n"
+                                                  << calculatedLineage
+                                                  << "\nLoaded lineage:\n"
+                                                  << loadedLineage;
+                                    throw yexception() << "Lineage in replay is different";
+                                }
+                                YQL_LOG(INFO) << "Lineage replay is the same";
+                            }
+                            if (typeCtx->QContext && typeCtx->QContext.CanWrite() && typeCtx->CorrectLineage) {
+                                typeCtx->QContext.GetWriter()->Put({LineageComponent, LineageResultLabel}, calculatedLineage).GetValueSync();
+                                YQL_LOG(INFO) << "Lineage is saved to QStorage";
+                            }
+                            return IGraphTransformer::TStatus::Ok;
+                        }),
+                    "Lineage",
+                    issueCode),
+                TTransformStage(
+                    new TNullTransformer(),
+                    "SkipLineage",
+                    issueCode)),
+            "LineageCalculation",
+            issueCode));
+    }
     AddProviderOptimization(issueCode);
     if (withFinalOptimization) {
         AddFinalCommonOptimization(issueCode);
diff --git a/yql/essentials/core/services/yql_transform_pipeline.h b/yql/essentials/core/services/yql_transform_pipeline.h
index 4cd60c76079..ca1dbfae83d 100644
--- a/yql/essentials/core/services/yql_transform_pipeline.h
+++ b/yql/essentials/core/services/yql_transform_pipeline.h
@@ -36,7 +36,7 @@ public:
     TTransformationPipeline& AddFinalCommonOptimization(EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddOptimization(bool checkWorld = true, bool withFinalOptimization = true, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddProviderOptimization(EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
-    TTransformationPipeline& AddOptimizationWithLineage(bool checkWorld = true, bool withFinalOptimization = true, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
+    TTransformationPipeline& AddOptimizationWithLineage(bool enableLineage, bool checkWorld = true, bool withFinalOptimization = true, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddLineageOptimization(TMaybe<TString>& lineageOut, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddCheckExecution(bool checkWorld = true, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddRun(TOperationProgressWriter writer, EYqlIssueCode issueCode = TIssuesIds::CORE_EXEC);