Merge 6868cc4 into e77b182

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,38 +74,36 @@ const std::partial_ordering IChunkedArray::TAddress::Compare(const TAddress& ite
     return TComparator::TypedCompare<true>(*Array, Position, *item.Array, item.Position);
 }
 
-namespace {
-class TChunkAccessor {
-private:
-    std::shared_ptr<arrow::ChunkedArray> ChunkedArray;
-public:
-    TChunkAccessor(const std::shared_ptr<arrow::ChunkedArray>& chunkedArray)
-        : ChunkedArray(chunkedArray)
-    {
-
-    }
-    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);
-    }
-};
-
+ 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(StartPosition <= position);
-    AFL_VERIFY(position < FinishPosition);
-    AFL_VERIFY(item.StartPosition <= itemPosition);
-    AFL_VERIFY(itemPosition < item.FinishPosition);
+    AFL_VERIFY(StartPosition <= position)("pos", position)("start", StartPosition);
+    AFL_VERIFY(position < FinishPosition)("pos", position)("finish", FinishPosition);
+    AFL_VERIFY(item.StartPosition <= itemPosition)("start", item.StartPosition)("item", itemPosition);
+    AFL_VERIFY(itemPosition < item.FinishPosition)("item", itemPosition)("finish", item.FinishPosition);
     return TComparator::TypedCompare<true>(*Array, position - StartPosition, *item.Array, itemPosition - item.StartPosition);
 }
 
@@ -119,9 +119,38 @@ TString IChunkedArray::TCurrentChunkAddress::DebugString(const ui64 position) co
     return NArrow::DebugString(Array, position - StartPosition);
 }
 
-IChunkedArray::TCurrentChunkAddress TTrivialChunkedArray::DoGetChunk(const std::optional<TCurrentChunkAddress>& chunkCurrent, const ui64 position) const {
-    TChunkAccessor accessor(Array);
-    return SelectChunk(chunkCurrent, position, accessor);
+namespace {
+class TChunkAccessor {
+private:
+    std::shared_ptr<arrow::ChunkedArray> ChunkedArray;
+    std::optional<IChunkedArray::TCurrentChunkAddress>* Result;
+
+public:
+    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();
+    }
+    void OnArray(const ui32 idx, const ui32 startPosition, const ui32 /*internalPosition*/) const {
+        *Result = IChunkedArray::TCurrentChunkAddress(ChunkedArray->chunk(idx), startPosition, idx);
+    }
+};
+
+}   // namespace
+
+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 {