YQ-3784 support computation graph invalidation in purecalc (#10799)

4 files changed, 155 insertions, 2 deletions

diff --git a/ydb/library/yql/public/purecalc/common/worker.cpp b/ydb/library/yql/public/purecalc/common/worker.cpp
index f7d1ad1e723..e04e2062ba8 100644
--- a/ydb/library/yql/public/purecalc/common/worker.cpp
+++ b/ydb/library/yql/public/purecalc/common/worker.cpp
@@ -526,6 +526,18 @@ void TPushStreamWorker::FeedToConsumer() {
     }
 }
 
+NYql::NUdf::IBoxedValue* TPushStreamWorker::GetPushStream() const {
+    auto& ctx = Graph_.ComputationGraph_->GetContext();
+    NUdf::TUnboxedValue pushStream = SelfNode_->GetValue(ctx);
+
+    if (Y_UNLIKELY(pushStream.IsInvalid())) {
+        SelfNode_->SetValue(ctx, Graph_.ComputationGraph_->GetHolderFactory().Create<TPushStream>());
+        pushStream = SelfNode_->GetValue(ctx);
+    }
+
+    return pushStream.AsBoxed().Get();
+}
+
 void TPushStreamWorker::SetConsumer(THolder<IConsumer<const NKikimr::NUdf::TUnboxedValue*>> consumer) {
     auto guard = Guard(GetScopedAlloc());
     const auto inputsCount = Graph_.SelfNodes_.size();
@@ -553,7 +565,7 @@ void TPushStreamWorker::Push(NKikimr::NUdf::TUnboxedValue&& value) {
     YQL_ENSURE(!Finished_, "OnFinish has already been sent to the consumer; no new values can be pushed");
 
     if (Y_LIKELY(SelfNode_)) {
-        static_cast<TPushStream*>(SelfNode_->GetValue(Graph_.ComputationGraph_->GetContext()).AsBoxed().Get())->SetValue(std::move(value));
+        static_cast<TPushStream*>(GetPushStream())->SetValue(std::move(value));
     }
 
     FeedToConsumer();
@@ -564,7 +576,7 @@ void TPushStreamWorker::OnFinish() {
     YQL_ENSURE(!Finished_, "already finished");
 
     if (Y_LIKELY(SelfNode_)) {
-        static_cast<TPushStream*>(SelfNode_->GetValue(Graph_.ComputationGraph_->GetContext()).AsBoxed().Get())->SetFinished();
+        static_cast<TPushStream*>(GetPushStream())->SetFinished();
     }
 
     FeedToConsumer();
diff --git a/ydb/library/yql/public/purecalc/common/worker.h b/ydb/library/yql/public/purecalc/common/worker.h
index 3f700ad353f..55eb440656b 100644
--- a/ydb/library/yql/public/purecalc/common/worker.h
+++ b/ydb/library/yql/public/purecalc/common/worker.h
@@ -164,6 +164,7 @@ namespace NYql {
 
         private:
             void FeedToConsumer();
+            NYql::NUdf::IBoxedValue* GetPushStream() const;
 
         public:
             void SetConsumer(THolder<IConsumer<const NKikimr::NUdf::TUnboxedValue*>>) override;
diff --git a/ydb/library/yql/public/purecalc/ut/test_mixed_allocators.cpp b/ydb/library/yql/public/purecalc/ut/test_mixed_allocators.cpp
new file mode 100644
index 00000000000..3038b5e4b2a
--- /dev/null
+++ b/ydb/library/yql/public/purecalc/ut/test_mixed_allocators.cpp
@@ -0,0 +1,139 @@
+#include <library/cpp/testing/unittest/registar.h>
+
+#include <ydb/library/yql/minikql/computation/mkql_computation_node_holders.h>
+#include <ydb/library/yql/minikql/mkql_string_util.h>
+
+#include <ydb/library/yql/public/purecalc/io_specs/protobuf/spec.h>
+#include <ydb/library/yql/public/purecalc/ut/protos/test_structs.pb.h>
+
+using namespace NYql::NPureCalc;
+
+namespace {
+    class TStatelessInputSpec : public TInputSpecBase {
+    public:
+        TStatelessInputSpec()
+            : Schemas_({NYT::TNode::CreateList()
+                .Add("StructType")
+                .Add(NYT::TNode::CreateList()
+                    .Add(NYT::TNode::CreateList()
+                        .Add("InputValue")
+                        .Add(NYT::TNode::CreateList()
+                            .Add("DataType")
+                            .Add("Utf8")
+                        )
+                    )
+                )
+            })
+        {};
+
+        const TVector<NYT::TNode>& GetSchemas() const override {
+            return Schemas_;
+        }
+
+    private:
+        const TVector<NYT::TNode> Schemas_;
+    };
+
+    class TStatelessInputConsumer : public IConsumer<const NYql::NUdf::TUnboxedValue&> {
+    public:
+        TStatelessInputConsumer(TWorkerHolder<IPushStreamWorker> worker)
+            : Worker_(std::move(worker))
+        {}
+
+        void OnObject(const NYql::NUdf::TUnboxedValue& value) override {
+            with_lock (Worker_->GetScopedAlloc()) {
+                NYql::NUdf::TUnboxedValue* items = nullptr;
+                NYql::NUdf::TUnboxedValue result = Worker_->GetGraph().GetHolderFactory().CreateDirectArrayHolder(1, items);
+
+                items[0] = value;
+
+                Worker_->Push(std::move(result));
+
+                // Clear graph after each object because
+                // values allocated on another allocator and should be released
+                Worker_->GetGraph().Invalidate();
+            }
+        }
+
+        void OnFinish() override {
+            with_lock(Worker_->GetScopedAlloc()) {
+                Worker_->OnFinish();
+            }
+        }
+
+    private:
+        TWorkerHolder<IPushStreamWorker> Worker_;
+    };
+
+    class TStatelessConsumer : public IConsumer<NPureCalcProto::TStringMessage*> {
+        const TString ExpectedData_;
+        const ui64 ExpectedRows_;
+        ui64 RowId_ = 0;
+
+    public:
+        TStatelessConsumer(const TString& expectedData, ui64 expectedRows)
+            : ExpectedData_(expectedData)
+            , ExpectedRows_(expectedRows)
+        {}
+
+        void OnObject(NPureCalcProto::TStringMessage* message) override {
+            UNIT_ASSERT_VALUES_EQUAL_C(ExpectedData_, message->GetX(), RowId_);
+            RowId_++;
+        }
+
+        void OnFinish() override {
+            UNIT_ASSERT_VALUES_EQUAL(ExpectedRows_, RowId_);
+        }
+    };
+}
+
+template <>
+struct TInputSpecTraits<TStatelessInputSpec> {
+    static constexpr bool IsPartial = false;
+    static constexpr bool SupportPushStreamMode = true;
+
+    using TConsumerType = THolder<IConsumer<const NYql::NUdf::TUnboxedValue&>>;
+
+    static TConsumerType MakeConsumer(const TStatelessInputSpec&, TWorkerHolder<IPushStreamWorker> worker) {
+        return MakeHolder<TStatelessInputConsumer>(std::move(worker));
+    }
+};
+
+Y_UNIT_TEST_SUITE(TestMixedAllocators) {
+    Y_UNIT_TEST(TestPushStream) {
+        const auto targetString = "large string >= 14 bytes";
+        const auto factory = MakeProgramFactory();
+        const auto sql = TStringBuilder() << "SELECT InputValue AS X FROM Input WHERE InputValue = \"" << targetString << "\";";
+
+        const auto program = factory->MakePushStreamProgram(
+            TStatelessInputSpec(),
+            TProtobufOutputSpec<NPureCalcProto::TStringMessage>(),
+            sql
+        );
+
+        const ui64 numberRows = 5;
+        const auto inputConsumer = program->Apply(MakeHolder<TStatelessConsumer>(targetString, numberRows));
+        NKikimr::NMiniKQL::TScopedAlloc alloc(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), true, false);
+
+        const auto pushString = [&](TString inputValue) {
+            NYql::NUdf::TUnboxedValue stringValue;
+            with_lock(alloc) {
+                stringValue = NKikimr::NMiniKQL::MakeString(inputValue);
+                alloc.Ref().LockObject(stringValue);
+            }
+
+            inputConsumer->OnObject(stringValue);
+
+            with_lock(alloc) {
+                alloc.Ref().UnlockObject(stringValue);
+                stringValue.Clear();
+            }
+        };
+
+        for (ui64 i = 0; i < numberRows; ++i) {
+            pushString(targetString);
+            pushString("another large string >= 14 bytes");
+        }
+        inputConsumer->OnFinish();
+    }
+}
diff --git a/ydb/library/yql/public/purecalc/ut/ya.make b/ydb/library/yql/public/purecalc/ut/ya.make
index 87f90a25ccf..6f23dcd7f4a 100644
--- a/ydb/library/yql/public/purecalc/ut/ya.make
+++ b/ydb/library/yql/public/purecalc/ut/ya.make
@@ -12,6 +12,7 @@ SRCS(
     test_user_data.cpp
     test_eval.cpp
     test_pool.cpp
+    test_mixed_allocators.cpp
 )
 
 PEERDIR(