Transfer data from a topic to a table. (#14371)

Co-authored-by: Ilnaz Nizametdinov <i.nizametdinov@gmail.com>

Author: nshestakov

ydb/core/kqp/runtime/kqp_write_table.cpp

@@ -44,6 +44,10 @@ class TColumnBatch : public IDataBatch {
         return result;
     }
 
+    std::shared_ptr<void> ExtractBatch() override {
+        return std::dynamic_pointer_cast<void>(Extract());
+    }
+
     explicit TColumnBatch(const TRecordBatchPtr& data)
         : Data(data)
         , Memory(NArrow::GetBatchDataSize(Data)) {
@@ -75,6 +79,11 @@ class TRowBatch : public IDataBatch {
         return {std::move(Cells), std::move(Data)};
     }
 
+    std::shared_ptr<void> ExtractBatch() override {
+        auto r = std::make_shared<std::pair<std::vector<TCell>, std::vector<TCharVectorPtr>>>(std::move(Extract()));
+        return std::reinterpret_pointer_cast<void>(r);
+    }
+
     TRowBatch(std::vector<TCell>&& cells, std::vector<TCharVectorPtr>&& data, i64 size, ui32 rows, ui16 columns)
         : Cells(std::move(cells))
         , Data(std::move(data))

ydb/core/kqp/runtime/kqp_write_table.h

@@ -15,6 +15,8 @@ class IDataBatch : public TThrRefBase {
     virtual TString SerializeToString() const = 0;
     virtual i64 GetMemory() const = 0;
     virtual bool IsEmpty() const = 0;
+
+    virtual std::shared_ptr<void> ExtractBatch() = 0;
 };
 
 using IDataBatchPtr = TIntrusivePtr<IDataBatch>;

ydb/core/persqueue/purecalc/purecalc.cpp

@@ -0,0 +1,269 @@
+#include "purecalc.h"
+
+#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
+#include <yql/essentials/minikql/computation/mkql_custom_list.h>
+#include <yql/essentials/minikql/mkql_string_util.h>
+#include <yql/essentials/utils/yql_panic.h>
+
+namespace NYdb::NTopic::NPurecalc {
+
+namespace {
+
+using namespace NYql::NUdf;
+using namespace NKikimr::NMiniKQL;
+
+constexpr const char* DataFieldName = "_data";
+constexpr const char* OffsetFieldName = "_offset";
+
+constexpr const size_t FieldCount = 2; // Change it when change fields
+
+struct FieldPositions {
+    ui64 Data = 0;
+    ui64 Offset = 0;
+};
+
+
+NYT::TNode CreateTypeNode(const TString& fieldType) {
+    return NYT::TNode::CreateList()
+        .Add("DataType")
+        .Add(fieldType);
+}
+
+void AddField(NYT::TNode& node, const TString& fieldName, const TString& fieldType) {
+    node.Add(
+        NYT::TNode::CreateList()
+            .Add(fieldName)
+            .Add(CreateTypeNode(fieldType))
+    );
+}
+
+
+NYT::TNode CreateMessageScheme() {
+    auto structMembers = NYT::TNode::CreateList();
+    AddField(structMembers, DataFieldName, "String");
+    AddField(structMembers, OffsetFieldName, "Uint64");
+
+    return NYT::TNode::CreateList()
+        .Add("StructType")
+        .Add(std::move(structMembers));
+}
+
+static const TVector<NYT::TNode> InputSchema{ CreateMessageScheme() };
+
+struct TMessageWrapper {
+    const TMessage& Message;
+
+    NYql::NUdf::TUnboxedValuePod GetData() const {
+        return NKikimr::NMiniKQL::MakeString(Message.Data);
+    }
+
+    NYql::NUdf::TUnboxedValuePod GetOffset() const {
+        return NYql::NUdf::TUnboxedValuePod(Message.Offset);
+    }
+};
+
+class TInputConverter {
+protected:
+    IWorker* Worker_;
+    TPlainContainerCache Cache_;
+    FieldPositions Position;
+
+public:
+    explicit TInputConverter(IWorker* worker)
+        : Worker_(worker)
+    {
+        const TStructType* structType = worker->GetInputType();
+        const ui64 count = structType->GetMembersCount();
+ 
+        for (ui64 i = 0; i < count; ++i) { 
+            const auto name = structType->GetMemberName(i);
+            if (name == DataFieldName) {
+                Position.Data = i;
+            } else if (name == OffsetFieldName) {
+                Position.Offset = i;
+            }
+        }
+    }
+
+public:
+    void DoConvert(const TMessage* message, TUnboxedValue& result) {
+        auto& holderFactory = Worker_->GetGraph().GetHolderFactory();
+        TUnboxedValue* items = nullptr;
+        result = Cache_.NewArray(holderFactory, static_cast<ui32>(FieldCount), items);
+
+        TMessageWrapper wrap {*message};
+        items[Position.Data] = wrap.GetData();
+        items[Position.Offset] = wrap.GetOffset();
+    }
+
+    void ClearCache() {
+        Cache_.Clear();
+    }
+};
+
+/**
+ * List (or, better, stream) of unboxed values. Used as an input value in pull workers.
+ */
+class TMessageListValue final: public TCustomListValue {
+private:
+    mutable bool HasIterator_ = false;
+    THolder<IStream<TMessage*>> Underlying_;
+    TInputConverter Converter;
+    IWorker* Worker_;
+    TScopedAlloc& ScopedAlloc_;
+
+public:
+    TMessageListValue(
+        TMemoryUsageInfo* memInfo,
+        const TMessageInputSpec& /*inputSpec*/,
+        THolder<IStream<TMessage*>> underlying,
+        IWorker* worker
+    )
+        : TCustomListValue(memInfo)
+        , Underlying_(std::move(underlying))
+        , Converter(worker)
+        , Worker_(worker)
+        , ScopedAlloc_(Worker_->GetScopedAlloc())
+    {
+    }
+
+    ~TMessageListValue() override {
+        {
+            // This list value stored in the worker's computation graph and destroyed upon the computation
+            // graph's destruction. This brings us to an interesting situation: scoped alloc is acquired,
+            // worker and computation graph are half-way destroyed, and now it's our turn to die. The problem is,
+            // the underlying stream may own another worker. This happens when chaining programs. Now, to destroy
+            // that worker correctly, we need to release our scoped alloc (because that worker has its own
+            // computation graph and scoped alloc).
+            // By the way, note that we shouldn't interact with the worker here because worker is in the middle of
+            // its own destruction. So we're using our own reference to the scoped alloc. That reference is alive
+            // because scoped alloc destroyed after computation graph.
+            auto unguard = Unguard(ScopedAlloc_);
+            Underlying_.Destroy();
+        }
+    }
+
+public:
+    TUnboxedValue GetListIterator() const override {
+        YQL_ENSURE(!HasIterator_, "Only one pass over input is supported");
+        HasIterator_ = true;
+        return TUnboxedValuePod(const_cast<TMessageListValue*>(this));
+    }
+
+    bool Next(TUnboxedValue& result) override {
+        const TMessage* message;
+        {
+            auto unguard = Unguard(ScopedAlloc_);
+            message = Underlying_->Fetch();
+        }
+
+        if (!message) {
+            return false;
+        }
+
+        Converter.DoConvert(message, result);
+
+        return true;
+    }
+
+    EFetchStatus Fetch(TUnboxedValue& result) override {
+        if (Next(result)) {
+            return EFetchStatus::Ok;
+        } else {
+            return EFetchStatus::Finish;
+        }
+    }
+};
+
+class TMessageConsumerImpl final: public IConsumer<TMessage*> {
+private:
+    TWorkerHolder<IPushStreamWorker> WorkerHolder;
+    TInputConverter Converter;
+
+public:
+    TMessageConsumerImpl(
+        const TMessageInputSpec& /*inputSpec*/,
+        TWorkerHolder<IPushStreamWorker> worker
+    )
+        : WorkerHolder(std::move(worker))
+        , Converter(WorkerHolder.Get())
+    {
+    }
+
+    ~TMessageConsumerImpl() override {
+        with_lock(WorkerHolder->GetScopedAlloc()) {
+            Converter.ClearCache();
+        }
+    }
+
+public:
+    void OnObject(TMessage* message) override {
+        TBindTerminator bind(WorkerHolder->GetGraph().GetTerminator());
+
+        with_lock(WorkerHolder->GetScopedAlloc()) {
+            Y_DEFER {
+                // Clear cache after each object because
+                // values allocated on another allocator and should be released
+                Converter.ClearCache();
+                WorkerHolder->Invalidate();
+            };
+
+            TUnboxedValue result;
+            Converter.DoConvert(message, result);
+            WorkerHolder->Push(std::move(result));
+        }
+    }
+
+    void OnFinish() override {
+        TBindTerminator bind(WorkerHolder->GetGraph().GetTerminator());
+
+        with_lock(WorkerHolder->GetScopedAlloc()) {
+            WorkerHolder->OnFinish();
+        }
+    }
+};
+
+} // namespace
+
+const TVector<NYT::TNode>& TMessageInputSpec::GetSchemas() const {
+    return InputSchema;
+}
+
+} // namespace NYdb::NTopic::NPurecalc
+
+namespace NYql::NPureCalc {
+
+using namespace NYdb::NTopic::NPurecalc;
+
+using ConsumerType = TInputSpecTraits<TMessageInputSpec>::TConsumerType;
+
+void TInputSpecTraits<TMessageInputSpec>::PreparePullStreamWorker(
+    const TMessageInputSpec& inputSpec,
+    IPullStreamWorker* worker,
+    THolder<IStream<TMessage*>> stream
+) {
+    with_lock(worker->GetScopedAlloc()) {
+        worker->SetInput(
+            worker->GetGraph().GetHolderFactory().Create<TMessageListValue>(inputSpec, std::move(stream), worker), 0);
+    }
+}
+
+void TInputSpecTraits<TMessageInputSpec>::PreparePullListWorker(
+    const TMessageInputSpec& inputSpec,
+    IPullListWorker* worker,
+    THolder<IStream<TMessage*>> stream
+) {
+    with_lock(worker->GetScopedAlloc()) {
+        worker->SetInput(
+            worker->GetGraph().GetHolderFactory().Create<TMessageListValue>(inputSpec, std::move(stream), worker), 0);
+    }
+}
+
+ConsumerType TInputSpecTraits<TMessageInputSpec>::MakeConsumer(
+    const TMessageInputSpec& inputSpec,
+    TWorkerHolder<IPushStreamWorker> worker
+) {
+    return MakeHolder<TMessageConsumerImpl>(inputSpec, std::move(worker));
+}
+
+} // namespace NYql::NPureCalc

ydb/core/persqueue/purecalc/purecalc.h

@@ -0,0 +1,57 @@
+#pragma once
+
+#include <yql/essentials/public/purecalc/common/interface.h>
+
+namespace NYdb::NTopic::NPurecalc {
+
+using namespace NYql::NPureCalc;
+
+struct TMessage {
+    TMessage(const TString& data)
+        : Data(data) {
+    }
+
+    TMessage& WithOffset(ui64 offset) {
+        Offset = offset;
+        return *this;
+    }
+
+    const TString& Data;
+    ui64 Offset = 0;
+};
+
+class TMessageInputSpec: public TInputSpecBase {
+public:
+    /**
+     * Build input spec and associate the given message descriptor.
+     */
+    explicit TMessageInputSpec() =  default;
+
+public:
+    const TVector<NYT::TNode>& GetSchemas() const override;
+    bool ProvidesBlocks() const override { return false; }
+};
+
+}
+
+namespace NYql::NPureCalc {
+
+template<>
+struct TInputSpecTraits<NYdb::NTopic::NPurecalc::TMessageInputSpec> {
+
+    static const constexpr bool IsPartial = false;
+
+    static const constexpr bool SupportPullStreamMode = true;
+    static const constexpr bool SupportPullListMode = true;
+    static const constexpr bool SupportPushStreamMode = true;
+
+    using TInput = NYdb::NTopic::NPurecalc::TMessage;
+    using TInputSpecType = NYdb::NTopic::NPurecalc::TMessageInputSpec;
+    using TConsumerType = THolder<IConsumer<TInput*>>;
+
+    static void PreparePullStreamWorker(const TInputSpecType&, IPullStreamWorker*, THolder<IStream<TInput*>>);
+    static void PreparePullListWorker(const TInputSpecType&, IPullListWorker*, THolder<IStream<TInput*>>);
+    static TConsumerType MakeConsumer(const TInputSpecType&, TWorkerHolder<IPushStreamWorker>);
+};
+
+}

ydb/core/persqueue/purecalc/ya.make

@@ -0,0 +1,13 @@
+LIBRARY()
+
+SRCS(
+    purecalc.cpp
+)
+
+PEERDIR(
+    ydb/core/fq/libs/row_dispatcher/purecalc_no_pg_wrapper
+)
+
+YQL_LAST_ABI_VERSION()
+
+END()

ydb/core/persqueue/ya.make

@@ -85,6 +85,7 @@ RECURSE(
     config
     events
     partition_key_range
+    purecalc
     writer
 )