YQL-19542: Save lineage in qstorage and check at replay

commit_hash:3693b0059b05bcf38d734e2d395dc3378902675c

9 files changed, 103 insertions, 31 deletions

diff --git a/yql/essentials/core/facade/yql_facade.cpp b/yql/essentials/core/facade/yql_facade.cpp
index dfb44c1af3d..0b6f9eca6b6 100644
--- a/yql/essentials/core/facade/yql_facade.cpp
+++ b/yql/essentials/core/facade/yql_facade.cpp
@@ -1192,7 +1192,7 @@ TProgram::TFutureStatus TProgram::OptimizeAsync(
                        .AddTypeAnnotation(TIssuesIds::CORE_TYPE_ANN, true)
                        .AddPostTypeAnnotation()
                        .Add(TExprOutputTransformer::Sync(ExprRoot_, traceOut), "ExprOutput")
-                       .AddOptimization()
+                       .AddOptimizationWithLineage(EnableLineage_)
                        .Add(TExprOutputTransformer::Sync(ExprRoot_, exprOut, withTypes), "AstOutput")
                        .Build();
 
@@ -1262,7 +1262,7 @@ TProgram::TFutureStatus TProgram::OptimizeAsyncWithConfig(
     pipeline.AddPostTypeAnnotation();
     pipelineConf.AfterTypeAnnotation(&pipeline);
 
-    pipeline.AddOptimization();
+    pipeline.AddOptimizationWithLineage(EnableLineage_);
     if (EnableRangeComputeFor_) {
         pipeline.Add(MakeExpandRangeComputeForTransformer(pipeline.GetTypeAnnotationContext()),
                      "ExpandRangeComputeFor", TIssuesIds::CORE_EXEC);
@@ -1402,7 +1402,7 @@ TProgram::TFutureStatus TProgram::RunAsync(
     pipeline.AddTypeAnnotation(TIssuesIds::CORE_TYPE_ANN, true);
     pipeline.AddPostTypeAnnotation();
     pipeline.Add(TExprOutputTransformer::Sync(ExprRoot_, traceOut), "ExprOutput");
-    pipeline.AddOptimizationWithLineage();
+    pipeline.AddOptimizationWithLineage(EnableLineage_);
     if (EnableRangeComputeFor_) {
         pipeline.Add(MakeExpandRangeComputeForTransformer(pipeline.GetTypeAnnotationContext()),
                      "ExpandRangeComputeFor", TIssuesIds::CORE_EXEC);
@@ -1480,7 +1480,7 @@ TProgram::TFutureStatus TProgram::RunAsyncWithConfig(
     pipeline.AddPostTypeAnnotation();
     pipelineConf.AfterTypeAnnotation(&pipeline);
 
-    pipeline.AddOptimizationWithLineage();
+    pipeline.AddOptimizationWithLineage(EnableLineage_);
     if (EnableRangeComputeFor_) {
         pipeline.Add(MakeExpandRangeComputeForTransformer(pipeline.GetTypeAnnotationContext()),
                      "ExpandRangeComputeFor", TIssuesIds::CORE_EXEC);
diff --git a/yql/essentials/core/facade/yql_facade.h b/yql/essentials/core/facade/yql_facade.h
index 018367a3b37..0f6767562ad 100644
--- a/yql/essentials/core/facade/yql_facade.h
+++ b/yql/essentials/core/facade/yql_facade.h
@@ -353,6 +353,10 @@ public:
 
     TString GetSourceCode() const;
 
+    void SetEnableLineage() {
+        EnableLineage_ = true;
+    }
+
 private:
     TProgram(
         const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry,
@@ -483,6 +487,7 @@ private:
     TMaybe<TString> GatewaysForMerge_;
     TIssues FinalIssues_;
     TMaybe<TIssue> ParametersIssue_;
+    bool EnableLineage_ = false;
 };
 
 void UpdateSqlFlagsFromQContext(const TQContext& qContext, THashSet<TString>& flags, TMaybe<TString> gatewaysPatch = {});
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);
diff --git a/yql/essentials/tools/yql_facade_run/yql_facade_run.cpp b/yql/essentials/tools/yql_facade_run/yql_facade_run.cpp
index c754381e910..5062d712f23 100644
--- a/yql/essentials/tools/yql_facade_run/yql_facade_run.cpp
+++ b/yql/essentials/tools/yql_facade_run/yql_facade_run.cpp
@@ -777,6 +777,9 @@ int TFacadeRunner::DoRun(TProgramFactory& factory) {
     }
     program->SetUseTableMetaFromGraph(RunOptions_.UseMetaFromGrpah);
     program->SetValidateOptions(RunOptions_.ValidateMode);
+    if (RunOptions_.EnableLineage) {
+        program->SetEnableLineage();
+    }
 
     bool fail = false;
     if (RunOptions_.ProgramType != EProgramType::SExpr) {
diff --git a/yql/essentials/tools/yql_facade_run/yql_facade_run.h b/yql/essentials/tools/yql_facade_run/yql_facade_run.h
index b7fdcbbb0ae..b0a08d370aa 100644
--- a/yql/essentials/tools/yql_facade_run/yql_facade_run.h
+++ b/yql/essentials/tools/yql_facade_run/yql_facade_run.h
@@ -141,6 +141,7 @@ public:
     bool EnableQPlayer = false;
     bool OptimizeLibs = true;
     bool CustomTests = false;
+    bool EnableLineage = false;
 
     void Parse(int argc, const char* argv[]);
 
diff --git a/yql/tools/yqlrun/lib/yqlrun_lib.cpp b/yql/tools/yqlrun/lib/yqlrun_lib.cpp
index aad7657ea4a..c0f23c972d3 100644
--- a/yql/tools/yqlrun/lib/yqlrun_lib.cpp
+++ b/yql/tools/yqlrun/lib/yqlrun_lib.cpp
@@ -49,6 +49,7 @@ TYqlRunTool::TYqlRunTool()
     GetRunOptions().UseRepeatableRandomAndTimeProviders = true;
     GetRunOptions().ResultsFormat = NYson::EYsonFormat::Pretty;
     GetRunOptions().CustomTests = true;
+    GetRunOptions().EnableLineage = true;
 
     GetRunOptions().AddOptExtension([this](NLastGetopt::TOpts& opts) {
         opts.AddLongOption('t', "table", "Table mapping").RequiredArgument("table@file")