Merge 2450266 into d3fa399

Author: ivanmorozov333

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

@@ -5,6 +5,8 @@
 #include <ydb/library/actors/core/log.h>
 #include <ydb/core/formats/arrow/permutations.h>
 #include <ydb/core/formats/arrow/arrow_helpers.h>
+#include <ydb/core/formats/arrow/splitter/simple.h>
+#include <ydb/core/formats/arrow/save_load/saver.h>
 
 namespace NKikimr::NArrow::NAccessor {
 
@@ -72,56 +74,83 @@ const std::partial_ordering IChunkedArray::TAddress::Compare(const TAddress& ite
     return TComparator::TypedCompare<true>(*Array, Position, *item.Array, item.Position);
 }
 
+ TChunkedArraySerialized::TChunkedArraySerialized(const std::shared_ptr<IChunkedArray>& array, const TString& serializedData)
+    : Array(array)
+    , SerializedData(serializedData) {
+    AFL_VERIFY(serializedData);
+    AFL_VERIFY(Array);
+    AFL_VERIFY(Array->GetRecordsCount());
+}
+
+std::optional<ui64> TTrivialArray::DoGetRawSize() const {
+    return NArrow::GetArrayDataSize(Array);
+}
+
+std::vector<NKikimr::NArrow::NAccessor::TChunkedArraySerialized> TTrivialArray::DoSplitBySizes(
+    const TColumnSaver& saver, const TString& fullSerializedData, const std::vector<ui64>& splitSizes) {
+    auto schema = std::make_shared<arrow::Schema>(arrow::FieldVector({ std::make_shared<arrow::Field>("f", GetDataType()) }));
+    auto chunks = NArrow::NSplitter::TSimpleSplitter(saver).SplitBySizes(
+        arrow::RecordBatch::Make(schema, GetRecordsCount(), { Array }), fullSerializedData, splitSizes);
+    std::vector<TChunkedArraySerialized> result;
+    for (auto&& i : chunks) {
+        AFL_VERIFY(i.GetSlicedBatch()->num_columns() == 1);
+        result.emplace_back(std::make_shared<TTrivialArray>(i.GetSlicedBatch()->column(0)), i.GetSerializedChunk());
+    }
+    return result;
+}
+
+std::partial_ordering IChunkedArray::TCurrentChunkAddress::Compare(const ui64 position, const TCurrentChunkAddress& item, const ui64 itemPosition) const {
+    AFL_VERIFY(GetStartPosition() <= position)("pos", position)("start", GetStartPosition());
+    AFL_VERIFY(position < GetFinishPosition())("pos", position)("finish", GetFinishPosition());
+    AFL_VERIFY(item.GetStartPosition() <= itemPosition)("start", item.GetStartPosition())("item", itemPosition);
+    AFL_VERIFY(itemPosition < item.GetFinishPosition())("item", itemPosition)("finish", item.GetFinishPosition());
+    return TComparator::TypedCompare<true>(*Array, position - GetStartPosition(), *item.Array, itemPosition - item.GetStartPosition());
+}
+
+std::shared_ptr<arrow::Array> IChunkedArray::TCurrentChunkAddress::CopyRecord(const ui64 recordIndex) const {
+    AFL_VERIFY(GetStartPosition() <= recordIndex);
+    AFL_VERIFY(recordIndex < GetFinishPosition());
+    return NArrow::CopyRecords(Array, { recordIndex - GetStartPosition() });
+}
+
+TString IChunkedArray::TCurrentChunkAddress::DebugString(const ui64 position) const {
+    AFL_VERIFY(position < GetFinishPosition());
+    AFL_VERIFY(GetStartPosition() <= position);
+    return NArrow::DebugString(Array, position - GetStartPosition());
+}
+
 namespace {
 class TChunkAccessor {
 private:
     std::shared_ptr<arrow::ChunkedArray> ChunkedArray;
+    std::optional<IChunkedArray::TCurrentChunkAddress>* Result;
+
 public:
-    TChunkAccessor(const std::shared_ptr<arrow::ChunkedArray>& chunkedArray)
+    TChunkAccessor(const std::shared_ptr<arrow::ChunkedArray>& chunkedArray, std::optional<IChunkedArray::TCurrentChunkAddress>& result)
         : ChunkedArray(chunkedArray)
+        , Result(&result)
     {
-
     }
     ui64 GetChunksCount() const {
         return (ui64)ChunkedArray->num_chunks();
     }
     ui64 GetChunkLength(const ui32 idx) const {
         return (ui64)ChunkedArray->chunk(idx)->length();
     }
-    std::shared_ptr<arrow::Array> GetArray(const ui32 idx) const {
-        return ChunkedArray->chunk(idx);
+    void OnArray(const ui32 idx, const ui32 startPosition, const ui32 /*internalPosition*/) const {
+        *Result = IChunkedArray::TCurrentChunkAddress(ChunkedArray->chunk(idx), startPosition, idx);
     }
 };
 
-}
-
-std::optional<ui64> TTrivialArray::DoGetRawSize() const {
-    return NArrow::GetArrayDataSize(Array);
-}
-
-std::partial_ordering IChunkedArray::TCurrentChunkAddress::Compare(const ui64 position, const TCurrentChunkAddress& item, const ui64 itemPosition) const {
-    AFL_VERIFY(StartPosition <= position);
-    AFL_VERIFY(position < FinishPosition);
-    AFL_VERIFY(item.StartPosition <= itemPosition);
-    AFL_VERIFY(itemPosition < item.FinishPosition);
-    return TComparator::TypedCompare<true>(*Array, position - StartPosition, *item.Array, itemPosition - item.StartPosition);
-}
-
-std::shared_ptr<arrow::Array> IChunkedArray::TCurrentChunkAddress::CopyRecord(const ui64 recordIndex) const {
-    AFL_VERIFY(StartPosition <= recordIndex);
-    AFL_VERIFY(recordIndex < FinishPosition);
-    return NArrow::CopyRecords(Array, { recordIndex - StartPosition });
-}
-
-TString IChunkedArray::TCurrentChunkAddress::DebugString(const ui64 position) const {
-    AFL_VERIFY(position < FinishPosition);
-    AFL_VERIFY(StartPosition <= position);
-    return NArrow::DebugString(Array, position - StartPosition);
-}
+}   // namespace
 
-IChunkedArray::TCurrentChunkAddress TTrivialChunkedArray::DoGetChunk(const std::optional<TCurrentChunkAddress>& chunkCurrent, const ui64 position) const {
-    TChunkAccessor accessor(Array);
-    return SelectChunk(chunkCurrent, position, accessor);
+IChunkedArray::TCurrentChunkAddress TTrivialChunkedArray::DoGetChunk(
+    const std::optional<TCurrentChunkAddress>& chunkCurrent, const ui64 position) const {
+    std::optional<IChunkedArray::TCurrentChunkAddress> result;
+    TChunkAccessor accessor(Array, result);
+    SelectChunk(chunkCurrent, position, accessor);
+    AFL_VERIFY(result);
+    return *result;
 }
 
 std::optional<ui64> TTrivialChunkedArray::DoGetRawSize() const {