YQL-19526: Use closed upvalues to run a callable

commit_hash:e58d49c7f50d1a4db0deb7b3f0f71f0b20e76b5b

1 files changed, 27 insertions, 0 deletions

diff --git a/yql/essentials/minikql/computation/mkql_computation_node.cpp b/yql/essentials/minikql/computation/mkql_computation_node.cpp
index e7e351cc979..01d72cf926a 100644
--- a/yql/essentials/minikql/computation/mkql_computation_node.cpp
+++ b/yql/essentials/minikql/computation/mkql_computation_node.cpp
@@ -26,6 +26,33 @@
 namespace NKikimr {
 namespace NMiniKQL {
 
+TComputationUpvalues::TComputationUpvalues(TComputationContext& ctx, IComputationNode* lambdaNode,
+                                           const TComputationExternalNodePtrVector& argNodes) {
+    std::set<const IComputationNode*> argSet(argNodes.cbegin(), argNodes.cend());
+    for (const auto uv : lambdaNode->GetUpvalues()) {
+        if (!argSet.contains(uv)) {
+            UpvalueNodes_.push_back(uv);
+        }
+    }
+    for (const auto uv : UpvalueNodes_) {
+        ClosedUpvalues_.push_back(uv->GetValue(ctx));
+    }
+    PreservedUpvalues_.resize(ClosedUpvalues_.size());
+}
+
+void TComputationUpvalues::SetUpvalues(TComputationContext& ctx) const {
+    for (size_t i = 0; i < UpvalueNodes_.size(); i++) {
+        PreservedUpvalues_[i] = UpvalueNodes_[i]->GetValue(ctx);
+        UpvalueNodes_[i]->SetValue(ctx, NUdf::TUnboxedValue(ClosedUpvalues_[i]));
+    }
+}
+
+void TComputationUpvalues::RestoreUpvalues(TComputationContext& ctx) const {
+    for (size_t i = 0; i < UpvalueNodes_.size(); i++) {
+        UpvalueNodes_[i]->SetValue(ctx, std::move(PreservedUpvalues_[i]));
+    }
+}
+
 std::unique_ptr<IArrowKernelComputationNode> IComputationNode::PrepareArrowKernelComputationNode(TComputationContext& ctx) const {
     Y_UNUSED(ctx);
     return {};