Merge 1ae2d56 into 4a621d8

Author: ivanmorozov333

ydb/core/formats/arrow/accessor/abstract/constructor.cpp

@@ -0,0 +1,11 @@
+#include "constructor.h"
+#include <ydb/core/formats/arrow/accessor/plain/constructor.h>
+
+namespace NKikimr::NArrow::NAccessor {
+
+TConstructorContainer TConstructorContainer::GetDefaultConstructor() {
+    static std::shared_ptr<IConstructor> result = std::make_shared<NPlain::TConstructor>();
+    return result;
+}
+
+}

ydb/core/formats/arrow/accessor/abstract/constructor.h

@@ -0,0 +1,69 @@
+#pragma once
+#include <ydb/core/formats/arrow/accessor/common/chunk_data.h>
+#include <ydb/core/formats/arrow/common/accessor.h>
+#include <ydb/core/formats/arrow/protos/accessor.pb.h>
+
+#include <ydb/services/bg_tasks/abstract/interface.h>
+
+#include <library/cpp/object_factory/object_factory.h>
+
+namespace NKikimr::NArrow::NAccessor {
+
+class IConstructor {
+public:
+    using TFactory = NObjectFactory::TObjectFactory<IConstructor, TString>;
+    using TProto = NKikimrArrowAccessorProto::TConstructor;
+
+private:
+    virtual TConclusion<std::shared_ptr<NArrow::NAccessor::IChunkedArray>> DoConstruct(
+        const std::shared_ptr<arrow::RecordBatch>& originalData, const TChunkConstructionData& externalInfo) const = 0;
+    virtual NKikimrArrowAccessorProto::TConstructor DoSerializeToProto() const = 0;
+    virtual bool DoDeserializeFromProto(const NKikimrArrowAccessorProto::TConstructor& proto) = 0;
+    virtual std::shared_ptr<arrow::Schema> DoGetExpectedSchema(const std::shared_ptr<arrow::Field>& resultColumn) const = 0;
+    virtual TString DoDebugString() const {
+        return "";
+    }
+
+public:
+    virtual ~IConstructor() = default;
+
+    TString DebugString() const {
+        return TStringBuilder() << GetClassName() << ":" << DoDebugString();
+    }
+
+    TConclusion<std::shared_ptr<NArrow::NAccessor::IChunkedArray>> Construct(
+        const std::shared_ptr<arrow::RecordBatch>& originalData, const TChunkConstructionData& externalInfo) const {
+        return DoConstruct(originalData, externalInfo);
+    }
+
+    bool DeserializeFromProto(const NKikimrArrowAccessorProto::TConstructor& proto) {
+        return DoDeserializeFromProto(proto);
+    }
+
+    NKikimrArrowAccessorProto::TConstructor SerializeToProto() const {
+        return DoSerializeToProto();
+    }
+
+    void SerializeToProto(NKikimrArrowAccessorProto::TConstructor& proto) const {
+        proto = DoSerializeToProto();
+    }
+
+    std::shared_ptr<arrow::Schema> GetExpectedSchema(const std::shared_ptr<arrow::Field>& resultColumn) const {
+        AFL_VERIFY(resultColumn);
+        return DoGetExpectedSchema(resultColumn);
+    }
+
+    virtual TString GetClassName() const = 0;
+};
+
+class TConstructorContainer: public NBackgroundTasks::TInterfaceProtoContainer<IConstructor> {
+private:
+    using TBase = NBackgroundTasks::TInterfaceProtoContainer<IConstructor>;
+
+public:
+    using TBase::TBase;
+
+    static TConstructorContainer GetDefaultConstructor();
+};
+
+}   // namespace NKikimr::NArrow::NAccessor

ydb/core/formats/arrow/accessor/abstract/request.cpp

@@ -0,0 +1,16 @@
+#include "request.h"
+
+namespace NKikimr::NArrow::NAccessor {
+
+TConclusionStatus TRequestedConstructorContainer::DeserializeFromRequest(NYql::TFeaturesExtractor& features) {
+    const std::optional<TString> className = features.Extract("DATA_ACCESSOR_CONSTRUCTOR.CLASS_NAME");
+    if (!className) {
+        return TConclusionStatus::Success();
+    }
+    if (!TBase::Initialize(*className)) {
+        return TConclusionStatus::Fail("don't know anything about class_name=" + *className);
+    }
+    return TBase::GetObjectPtr()->DeserializeFromRequest(features);
+}
+
+}

ydb/core/formats/arrow/accessor/abstract/request.h

@@ -0,0 +1,58 @@
+#pragma once
+#include "constructor.h"
+
+#include <ydb/core/formats/arrow/protos/accessor.pb.h>
+
+#include <ydb/services/bg_tasks/abstract/interface.h>
+#include <ydb/services/metadata/abstract/request_features.h>
+
+#include <library/cpp/object_factory/object_factory.h>
+
+namespace NKikimr::NArrow::NAccessor {
+
+class IRequestedConstructor {
+public:
+    using TFactory = NObjectFactory::TObjectFactory<IRequestedConstructor, TString>;
+    using TProto = NKikimrArrowAccessorProto::TRequestedConstructor;
+private:
+    virtual TConclusion<NArrow::NAccessor::TConstructorContainer> DoBuildConstructor() const = 0;
+    virtual NKikimrArrowAccessorProto::TRequestedConstructor DoSerializeToProto() const = 0;
+    virtual bool DoDeserializeFromProto(const NKikimrArrowAccessorProto::TRequestedConstructor& proto) = 0;
+    virtual TConclusionStatus DoDeserializeFromRequest(NYql::TFeaturesExtractor& features) = 0;
+
+public:
+    virtual ~IRequestedConstructor() = default;
+
+    NKikimrArrowAccessorProto::TRequestedConstructor SerializeToProto() const {
+        return DoSerializeToProto();
+    }
+
+    void SerializeToProto(NKikimrArrowAccessorProto::TRequestedConstructor& proto) const {
+        proto = DoSerializeToProto();
+    }
+
+    bool DeserializeFromProto(const NKikimrArrowAccessorProto::TRequestedConstructor& proto) {
+        return DoDeserializeFromProto(proto);
+    }
+
+    TConclusionStatus DeserializeFromRequest(NYql::TFeaturesExtractor& features) {
+        return DoDeserializeFromRequest(features);
+    }
+
+    TConclusion<TConstructorContainer> BuildConstructor() const {
+        return DoBuildConstructor();
+    }
+
+    virtual TString GetClassName() const = 0;
+};
+
+class TRequestedConstructorContainer: public NBackgroundTasks::TInterfaceProtoContainer<IRequestedConstructor> {
+private:
+    using TBase = NBackgroundTasks::TInterfaceProtoContainer<IRequestedConstructor>;
+
+public:
+    using TBase::TBase;
+    TConclusionStatus DeserializeFromRequest(NYql::TFeaturesExtractor& features);
+};
+
+}   // namespace NKikimr::NArrow::NAccessor

ydb/core/formats/arrow/accessor/abstract/ya.make

@@ -0,0 +1,16 @@
+LIBRARY()
+
+PEERDIR(
+    ydb/core/formats/arrow/protos
+    ydb/core/formats/arrow/accessor/common
+    contrib/libs/apache/arrow
+    ydb/library/conclusion
+    ydb/services/metadata/abstract
+)
+
+SRCS(
+    constructor.cpp
+    request.cpp
+)
+
+END()

ydb/core/formats/arrow/accessor/common/chunk_data.cpp

@@ -0,0 +1,5 @@
+#include "chunk_data.h"
+
+namespace NKikimr::NArrow::NAccessor {
+
+}

ydb/core/formats/arrow/accessor/common/chunk_data.h

@@ -0,0 +1,23 @@
+#pragma once
+#include <ydb/library/accessor/accessor.h>
+#include <contrib/libs/apache/arrow/cpp/src/arrow/type.h>
+#include <contrib/libs/apache/arrow/cpp/src/arrow/scalar.h>
+
+namespace NKikimr::NArrow::NAccessor {
+
+class TChunkConstructionData {
+private:
+    YDB_READONLY(ui32, RecordsCount, 0);
+    YDB_READONLY_DEF(std::shared_ptr<arrow::Scalar>, DefaultValue);
+    YDB_READONLY_DEF(std::shared_ptr<arrow::DataType>, ColumnType);
+
+public:
+    TChunkConstructionData(
+        const ui32 recordsCount, const std::shared_ptr<arrow::Scalar>& defaultValue, const std::shared_ptr<arrow::DataType>& columnType)
+        : RecordsCount(recordsCount)
+        , DefaultValue(defaultValue)
+        , ColumnType(columnType) {
+    }
+};
+
+}   // namespace NKikimr::NArrow::NAccessor

ydb/core/formats/arrow/accessor/common/ya.make

@@ -0,0 +1,11 @@
+LIBRARY()
+
+PEERDIR(
+    contrib/libs/apache/arrow
+)
+
+SRCS(
+    chunk_data.cpp
+)
+
+END()

ydb/core/formats/arrow/accessor/composite/accessor.cpp

@@ -0,0 +1,58 @@
+#include "accessor.h"
+namespace NKikimr::NArrow::NAccessor {
+
+namespace {
+class TCompositeChunkAccessor {
+private:
+    const std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>>& Chunks;
+    std::optional<IChunkedArray::TCurrentChunkAddress>* ResultChunkAddress = nullptr;
+    std::optional<IChunkedArray::TCurrentArrayAddress>* ResultArrayAddress = nullptr;
+
+public:
+    TCompositeChunkAccessor(
+        const std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>>& chunks, std::optional<IChunkedArray::TCurrentChunkAddress>& result)
+        : Chunks(chunks)
+        , ResultChunkAddress(&result) {
+    }
+    TCompositeChunkAccessor(
+        const std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>>& chunks, std::optional<IChunkedArray::TCurrentArrayAddress>& result)
+        : Chunks(chunks)
+        , ResultArrayAddress(&result) {
+    }
+    ui64 GetChunksCount() const {
+        return Chunks.size();
+    }
+    ui64 GetChunkLength(const ui32 idx) const {
+        return Chunks[idx]->GetRecordsCount();
+    }
+    void OnArray(const ui32 chunkIdx, const ui32 startPosition, const ui32 internalPosition) const {
+        if (ResultChunkAddress) {
+            *ResultChunkAddress = NArrow::NAccessor::IChunkedArray::TCurrentChunkAddress(
+                Chunks[chunkIdx]->GetChunk({}, internalPosition).GetArray(), startPosition, chunkIdx);
+        }
+        if (ResultArrayAddress) {
+            *ResultArrayAddress = NArrow::NAccessor::IChunkedArray::TCurrentArrayAddress(Chunks[chunkIdx], startPosition, chunkIdx);
+        }
+    }
+};
+}   // namespace
+
+NArrow::NAccessor::IChunkedArray::TCurrentChunkAddress TCompositeChunkedArray::DoGetChunk(
+    const std::optional<TCurrentChunkAddress>& chunkCurrent, const ui64 position) const {
+    std::optional<IChunkedArray::TCurrentChunkAddress> result;
+    TCompositeChunkAccessor accessor(Chunks, result);
+    SelectChunk(chunkCurrent, position, accessor);
+    AFL_VERIFY(result);
+    return *result;
+}
+
+NArrow::NAccessor::IChunkedArray::TCurrentArrayAddress TCompositeChunkedArray::DoGetArray(
+    const std::optional<TCurrentArrayAddress>& chunkCurrent, const ui64 position, const std::shared_ptr<IChunkedArray>& /*selfPtr*/) const {
+    std::optional<IChunkedArray::TCurrentArrayAddress> result;
+    TCompositeChunkAccessor accessor(Chunks, result);
+    SelectChunk(chunkCurrent, position, accessor);
+    AFL_VERIFY(result);
+    return *result;
+}
+
+}   // namespace NKikimr::NArrow::NAccessor

ydb/core/formats/arrow/accessor/composite/accessor.h

@@ -0,0 +1,72 @@
+#pragma once
+#include <ydb/core/formats/arrow/common/accessor.h>
+#include <ydb/library/accessor/accessor.h>
+
+namespace NKikimr::NArrow::NAccessor {
+
+class TCompositeChunkedArray: public NArrow::NAccessor::IChunkedArray {
+private:
+    using TBase = NArrow::NAccessor::IChunkedArray;
+
+private:
+    std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>> Chunks;
+
+protected:
+    virtual TCurrentArrayAddress DoGetArray(const std::optional<TCurrentArrayAddress>& chunkCurrent, const ui64 /*position*/,
+        const std::shared_ptr<IChunkedArray>& selfPtr) const override;
+
+    virtual std::vector<NArrow::NAccessor::TChunkedArraySerialized> DoSplitBySizes(
+        const TColumnSaver& /*saver*/, const TString& /*fullSerializedData*/, const std::vector<ui64>& /*splitSizes*/) override {
+        AFL_VERIFY(false);
+        return {};
+    }
+
+    virtual std::shared_ptr<arrow::Scalar> DoGetScalar(const ui32 index) const override {
+        AFL_VERIFY(false)("problem", "cannot use method");
+        return nullptr;
+    }
+    virtual std::optional<ui64> DoGetRawSize() const override {
+        return {};
+    }
+    virtual std::shared_ptr<arrow::Scalar> DoGetMaxScalar() const override {
+        AFL_VERIFY(false);
+        return nullptr;
+    }
+    virtual TCurrentChunkAddress DoGetChunk(const std::optional<TCurrentChunkAddress>& chunkCurrent, const ui64 position) const override;
+    virtual std::shared_ptr<arrow::ChunkedArray> DoGetChunkedArray() const override {
+        AFL_VERIFY(false);
+        return nullptr;
+    }
+
+    TCompositeChunkedArray(std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>>&& chunks, const ui32 recordsCount,
+        const std::shared_ptr<arrow::DataType>& type)
+        : TBase(recordsCount, NArrow::NAccessor::IChunkedArray::EType::SerializedChunkedArray, type)
+        , Chunks(std::move(chunks)) {
+    }
+
+public:
+    class TBuilder {
+    private:
+        ui32 RecordsCount = 0;
+        std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>> Chunks;
+        const std::shared_ptr<arrow::DataType> Type;
+
+    public:
+        TBuilder(const std::shared_ptr<arrow::DataType>& type)
+            : Type(type) {
+            AFL_VERIFY(Type);
+        }
+
+        void AddChunk(const std::shared_ptr<NArrow::NAccessor::IChunkedArray>& arr) {
+            AFL_VERIFY(arr->GetDataType()->id() == Type->id())("incoming", arr->GetDataType()->ToString())("main", Type->ToString());
+            Chunks.emplace_back(arr);
+            RecordsCount += arr->GetRecordsCount();
+        }
+
+        std::shared_ptr<TCompositeChunkedArray> Finish() {
+            return std::shared_ptr<TCompositeChunkedArray>(new TCompositeChunkedArray(std::move(Chunks), RecordsCount, Type));
+        }
+    };
+};
+
+}   // namespace NKikimr::NArrow::NAccessor

ydb/core/formats/arrow/accessor/composite/ya.make

@@ -0,0 +1,12 @@
+LIBRARY()
+
+PEERDIR(
+    contrib/libs/apache/arrow
+    ydb/core/formats/arrow/common
+)
+
+SRCS(
+    accessor.cpp
+)
+
+END()