[yql] Isolate Evaluation states

4 files changed, 29 insertions, 12 deletions

diff --git a/ydb/library/yql/core/services/yql_eval_expr.cpp b/ydb/library/yql/core/services/yql_eval_expr.cpp
index bd3e32bb54..06f9aabff5 100644
--- a/ydb/library/yql/core/services/yql_eval_expr.cpp
+++ b/ydb/library/yql/core/services/yql_eval_expr.cpp
@@ -201,6 +201,24 @@ private:
     TNodeSet CurrentEvalNodes;
 };
 
+struct TEvalScope {
+    TEvalScope(TTypeAnnotationContext& types)
+        : Types(types)
+    {
+        ++Types.EvaluationInProgress;
+        for (auto& dataProvider : Types.DataSources) {
+            dataProvider->EnterEvaluation(Types.EvaluationInProgress);
+        }
+    }
+
+    ~TEvalScope() {
+        for (auto& dataProvider : Types.DataSources) {
+            dataProvider->LeaveEvaluation(Types.EvaluationInProgress);
+        }
+        --Types.EvaluationInProgress;
+    }
+    TTypeAnnotationContext& Types;
+};
 
 bool ValidateCalcWorlds(const TExprNode& node, const TTypeAnnotationContext& types, TNodeSet& visited) {
     if (!visited.emplace(&node).second) {
@@ -948,10 +966,9 @@ IGraphTransformer::TStatus EvaluateExpression(const TExprNode::TPtr& input, TExp
             calcWorldRoot.Drop();
             fullTransformer->Rewind();
             auto prevSteps = ctx.Step;
-            ++types.EvaluationInProgress;
+            TEvalScope scope(types);
             status = SyncTransform(*fullTransformer, clonedArg, ctx);
             ctx.Step = prevSteps;
-            --types.EvaluationInProgress;
             if (status.Level == IGraphTransformer::TStatus::Error) {
                 return nullptr;
             }
@@ -999,11 +1016,6 @@ IGraphTransformer::TStatus EvaluateExpression(const TExprNode::TPtr& input, TExp
             delegatedNode->SetState(TExprNode::EState::ConstrComplete);
 
             status = SyncTransform(calcTransfomer ? *calcTransfomer : (*calcProvider.Get())->GetCallableExecutionTransformer(), delegatedNode, ctx);
-
-            for (auto& dataProvider : types.DataSources) {
-                dataProvider->UndoEvaluationChanges();
-            }
-
             if (status.Level == IGraphTransformer::TStatus::Error) {
                 return nullptr;
             }
diff --git a/ydb/library/yql/core/yql_data_provider.h b/ydb/library/yql/core/yql_data_provider.h
index b6f5b56b23..42d87d691f 100644
--- a/ydb/library/yql/core/yql_data_provider.h
+++ b/ydb/library/yql/core/yql_data_provider.h
@@ -142,7 +142,8 @@ public:
     virtual bool CanPullResult(const TExprNode& node, TSyncMap& syncList, bool& canRef) = 0;
     virtual bool GetExecWorld(const TExprNode::TPtr& node, TExprNode::TPtr& root) = 0;
     virtual bool CanEvaluate(const TExprNode& node) = 0;
-    virtual void UndoEvaluationChanges() = 0;
+    virtual void EnterEvaluation(ui64 id) = 0;
+    virtual void LeaveEvaluation(ui64 id) = 0;
     virtual TExprNode::TPtr CleanupWorld(const TExprNode::TPtr& node, TExprContext& ctx) = 0;
     virtual TExprNode::TPtr OptimizePull(const TExprNode::TPtr& source, const TFillSettings& fillSettings, TExprContext& ctx,
         IOptimizationContext& optCtx) = 0;
diff --git a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp
index c78ca70e55..c8187623fe 100644
--- a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp
+++ b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.cpp
@@ -216,9 +216,12 @@ bool TDataProviderBase::CanEvaluate(const TExprNode& node) {
     return false;
 }
 
-void TDataProviderBase::UndoEvaluationChanges() {
-    GetCallableExecutionTransformer().Rewind();
-    GetEpochsTransformer().Rewind();
+void TDataProviderBase::EnterEvaluation(ui64 id) {
+    Y_UNUSED(id);
+}
+
+void TDataProviderBase::LeaveEvaluation(ui64 id) {
+    Y_UNUSED(id);
 }
 
 TExprNode::TPtr TDataProviderBase::CleanupWorld(const TExprNode::TPtr& node, TExprContext& ctx) {
diff --git a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h
index 5f30ced450..a102d0a423 100644
--- a/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h
+++ b/ydb/library/yql/providers/common/provider/yql_data_provider_impl.h
@@ -66,7 +66,8 @@ public:
     bool CanPullResult(const TExprNode& node, TSyncMap& syncList, bool& canRef) override;
     bool GetExecWorld(const TExprNode::TPtr& node, TExprNode::TPtr& root) override;
     bool CanEvaluate(const TExprNode& node) override;
-    void UndoEvaluationChanges() override;
+    void EnterEvaluation(ui64 id) override;
+    void LeaveEvaluation(ui64 id) override;
     TExprNode::TPtr CleanupWorld(const TExprNode::TPtr& node, TExprContext& ctx) override;
     TExprNode::TPtr OptimizePull(const TExprNode::TPtr& source, const TFillSettings& fillSettings, TExprContext& ctx,
         IOptimizationContext& optCtx) override;