refactor: extract WriteBuffer from MergeTreeWriter

Extract batch buffering logic from MergeTreeWriter into a dedicated
WriteBuffer class. WriteBuffer manages batch_vec_, row_kinds_vec_,
and current_memory_in_bytes_, and provides Write() and Flush() methods.

This is a pure internal refactoring that:
- Does not change any external API or behavior
- Prepares for future Spill-to-Disk functionality
- Improves separation of concerns in MergeTreeWriter

Author: zjw1111

src/paimon/CMakeLists.txt

@@ -239,6 +239,7 @@ set(PAIMON_CORE_SRCS
     core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp
     core/mergetree/compact/changelog_merge_tree_rewriter.cpp
     core/mergetree/merge_tree_writer.cpp
+    core/mergetree/write_buffer.cpp
     core/mergetree/levels.cpp
     core/mergetree/lookup_levels.cpp
     core/migrate/file_meta_utils.cpp
@@ -592,6 +593,7 @@ if(PAIMON_BUILD_TESTS)
                     core/mergetree/lookup/persist_processor_test.cpp
                     core/mergetree/drop_delete_reader_test.cpp
                     core/mergetree/merge_tree_writer_test.cpp
+                    core/mergetree/write_buffer_test.cpp
                     core/mergetree/sorted_run_test.cpp
                     core/migrate/file_meta_utils_test.cpp
                     core/operation/metrics/compaction_metrics_test.cpp

src/paimon/core/mergetree/merge_tree_writer.cpp

@@ -23,42 +23,30 @@
 #include <utility>
 
 #include "arrow/api.h"
-#include "arrow/array/array_base.h"
-#include "arrow/array/array_binary.h"
-#include "arrow/array/array_nested.h"
 #include "arrow/c/abi.h"
-#include "arrow/c/bridge.h"
 #include "arrow/c/helpers.h"
-#include "arrow/util/checked_cast.h"
-#include "fmt/format.h"
 #include "paimon/common/metrics/metrics_impl.h"
 #include "paimon/common/table/special_fields.h"
-#include "paimon/common/types/data_field.h"
 #include "paimon/common/utils/arrow/status_utils.h"
 #include "paimon/common/utils/scope_guard.h"
 #include "paimon/core/io/async_key_value_producer_and_consumer.h"
 #include "paimon/core/io/compact_increment.h"
 #include "paimon/core/io/data_file_path_factory.h"
 #include "paimon/core/io/data_increment.h"
 #include "paimon/core/io/key_value_data_file_writer.h"
-#include "paimon/core/io/key_value_in_memory_record_reader.h"
 #include "paimon/core/io/key_value_meta_projection_consumer.h"
 #include "paimon/core/io/key_value_record_reader.h"
 #include "paimon/core/io/row_to_arrow_array_converter.h"
 #include "paimon/core/io/single_file_writer.h"
 #include "paimon/core/manifest/file_source.h"
 #include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h"
+#include "paimon/core/mergetree/write_buffer.h"
 #include "paimon/core/utils/commit_increment.h"
-#include "paimon/data/decimal.h"
 #include "paimon/format/file_format.h"
 #include "paimon/format/writer_builder.h"
 #include "paimon/metrics.h"
 
 namespace paimon {
-class FieldsComparator;
-class MemoryPool;
-template <typename T>
-class MergeFunctionWrapper;
 class FormatStatsExtractor;
 
 MergeTreeWriter::MergeTreeWriter(
@@ -71,7 +59,6 @@ MergeTreeWriter::MergeTreeWriter(
     const CoreOptions& options, const std::shared_ptr<CompactManager>& compact_manager,
     const std::shared_ptr<MemoryPool>& pool)
     : last_sequence_number_(last_sequence_number + 1),
-      current_memory_in_bytes_(0),
       pool_(pool),
       trimmed_primary_keys_(trimmed_primary_keys),
       options_(options),
@@ -80,10 +67,12 @@ MergeTreeWriter::MergeTreeWriter(
       user_defined_seq_comparator_(user_defined_seq_comparator),
       merge_function_wrapper_(merge_function_wrapper),
       schema_id_(schema_id),
-      value_type_(arrow::struct_(value_schema->fields())),
       compact_manager_(compact_manager),
       metrics_(std::make_shared<MetricsImpl>()) {
     write_schema_ = SpecialFields::CompleteSequenceAndValueKindField(value_schema);
+    write_buffer_ = std::make_unique<WriteBuffer>(
+        arrow::struct_(value_schema->fields()), trimmed_primary_keys_,
+        options_.GetSequenceField(), key_comparator_, merge_function_wrapper_, pool_);
 }
 
 Status MergeTreeWriter::DoClose() {
@@ -114,8 +103,7 @@ Status MergeTreeWriter::DoClose() {
         [[maybe_unused]] auto s = options_.GetFileSystem()->Delete(path_factory_->ToPath(file));
     }
 
-    batch_vec_.clear();
-    row_kinds_vec_.clear();
+    write_buffer_->Clear();
     new_files_.clear();
     deleted_files_.clear();
     compact_before_.clear();
@@ -129,23 +117,8 @@ Status MergeTreeWriter::DoClose() {
 }
 
 Status MergeTreeWriter::Write(std::unique_ptr<RecordBatch>&& moved_batch) {
-    if (ArrowArrayIsReleased(moved_batch->GetData())) {
-        return Status::Invalid("invalid batch: data is released");
-    }
-    std::unique_ptr<RecordBatch> batch = std::move(moved_batch);
-    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array,
-                                      arrow::ImportArray(batch->GetData(), value_type_));
-    auto value_struct_array =
-        arrow::internal::checked_pointer_cast<arrow::StructArray>(arrow_array);
-    if (value_struct_array == nullptr) {
-        return Status::Invalid("invalid RecordBatch: cannot cast to StructArray");
-    }
-    PAIMON_ASSIGN_OR_RAISE(int64_t memory_in_bytes, EstimateMemoryUse(value_struct_array));
-    current_memory_in_bytes_ += memory_in_bytes;
-
-    batch_vec_.push_back(std::move(value_struct_array));
-    row_kinds_vec_.push_back(batch->GetRowKind());
-    if (current_memory_in_bytes_ >= options_.GetWriteBufferSize()) {
+    PAIMON_RETURN_NOT_OK(write_buffer_->Write(std::move(moved_batch)));
+    if (write_buffer_->GetMemoryUsage() >= options_.GetWriteBufferSize()) {
         return Flush(/*wait_for_latest_compaction=*/false, /*forced_full_compaction=*/false);
     }
     return Status::OK();
@@ -247,25 +220,13 @@ Result<bool> MergeTreeWriter::CompactNotCompleted() {
 }
 
 Status MergeTreeWriter::Flush(bool wait_for_latest_compaction, bool forced_full_compaction) {
-    if (!batch_vec_.empty()) {
+    if (!write_buffer_->IsEmpty()) {
         if (compact_manager_->ShouldWaitForLatestCompaction()) {
             wait_for_latest_compaction = true;
         }
-        // 1. create key value iter for each record batch
-        std::vector<std::unique_ptr<KeyValueRecordReader>> readers;
-        readers.reserve(batch_vec_.size());
-        for (size_t i = 0; i < batch_vec_.size(); ++i) {
-            int64_t sequence_number = last_sequence_number_;
-            last_sequence_number_ += batch_vec_[i]->length();
-            auto in_memory_reader = std::make_unique<KeyValueInMemoryRecordReader>(
-                sequence_number, std::move(batch_vec_[i]), std::move(row_kinds_vec_[i]),
-                trimmed_primary_keys_, options_.GetSequenceField(), key_comparator_,
-                merge_function_wrapper_, pool_);
-            readers.push_back(std::move(in_memory_reader));
-        }
-        batch_vec_.clear();
-        row_kinds_vec_.clear();
-        current_memory_in_bytes_ = 0;
+        // 1. flush write buffer to get in-memory readers
+        PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> readers,
+                               write_buffer_->Flush(last_sequence_number_));
         // 2. prepare loser tree sort merge reader
         auto sort_merge_reader = std::make_unique<SortMergeReaderWithLoserTree>(
             std::move(readers), key_comparator_, user_defined_seq_comparator_,
@@ -354,68 +315,4 @@ MergeTreeWriter::CreateRollingRowWriter() const {
         options_.GetTargetFileSize(/*has_primary_key=*/true), create_file_writer);
 }
 
-Result<int64_t> MergeTreeWriter::EstimateMemoryUse(const std::shared_ptr<arrow::Array>& array) {
-    arrow::Type::type type = array->type()->id();
-    int64_t null_bits_size_in_bytes = (array->length() + 7) / 8;
-    switch (type) {
-        case arrow::Type::type::BOOL:
-            return null_bits_size_in_bytes + array->length() * sizeof(bool);
-        case arrow::Type::type::INT8:
-            return null_bits_size_in_bytes + array->length() * sizeof(int8_t);
-        case arrow::Type::type::INT16:
-            return null_bits_size_in_bytes + array->length() * sizeof(int16_t);
-        case arrow::Type::type::INT32:
-            return null_bits_size_in_bytes + array->length() * sizeof(int32_t);
-        case arrow::Type::type::DATE32:
-            return null_bits_size_in_bytes + array->length() * sizeof(int32_t);
-        case arrow::Type::type::INT64:
-            return null_bits_size_in_bytes + array->length() * sizeof(int64_t);
-        case arrow::Type::type::FLOAT:
-            return null_bits_size_in_bytes + array->length() * sizeof(float);
-        case arrow::Type::type::DOUBLE:
-            return null_bits_size_in_bytes + array->length() * sizeof(double);
-        case arrow::Type::type::TIMESTAMP:
-            return null_bits_size_in_bytes + array->length() * sizeof(int64_t);
-        case arrow::Type::type::DECIMAL:
-            return null_bits_size_in_bytes + array->length() * sizeof(Decimal::int128_t);
-        case arrow::Type::type::STRING:
-        case arrow::Type::type::BINARY: {
-            auto binary_array =
-                arrow::internal::checked_cast<const arrow::BinaryArray*>(array.get());
-            assert(binary_array);
-            int64_t value_length = binary_array->total_values_length();
-            int64_t offset_length = array->length() * sizeof(int32_t);
-            return null_bits_size_in_bytes + value_length + offset_length;
-        }
-        case arrow::Type::type::LIST: {
-            auto list_array = arrow::internal::checked_cast<const arrow::ListArray*>(array.get());
-            assert(list_array);
-            PAIMON_ASSIGN_OR_RAISE(int64_t value_mem, EstimateMemoryUse(list_array->values()));
-            return null_bits_size_in_bytes + value_mem;
-        }
-        case arrow::Type::type::MAP: {
-            auto map_array = arrow::internal::checked_cast<const arrow::MapArray*>(array.get());
-            assert(map_array);
-            PAIMON_ASSIGN_OR_RAISE(int64_t key_mem, EstimateMemoryUse(map_array->keys()));
-            PAIMON_ASSIGN_OR_RAISE(int64_t item_mem, EstimateMemoryUse(map_array->items()));
-            return null_bits_size_in_bytes + key_mem + item_mem;
-        }
-        case arrow::Type::type::STRUCT: {
-            auto struct_array =
-                arrow::internal::checked_cast<const arrow::StructArray*>(array.get());
-            assert(struct_array);
-            int64_t struct_mem = 0;
-            for (const auto& field : struct_array->fields()) {
-                PAIMON_ASSIGN_OR_RAISE(int64_t field_mem, EstimateMemoryUse(field));
-                struct_mem += field_mem;
-            }
-            return null_bits_size_in_bytes + struct_mem;
-        }
-        default:
-            assert(false);
-            return Status::Invalid(fmt::format("Do not support type {} in EstimateMemoryUse",
-                                               array->type()->ToString()));
-    }
-}
-
 }  // namespace paimon

src/paimon/core/mergetree/merge_tree_writer.h

@@ -28,6 +28,7 @@
 #include "paimon/core/io/rolling_file_writer.h"
 #include "paimon/core/key_value.h"
 #include "paimon/core/mergetree/compact/merge_function_wrapper.h"
+#include "paimon/core/mergetree/write_buffer.h"
 #include "paimon/core/utils/batch_writer.h"
 #include "paimon/core/utils/commit_increment.h"
 #include "paimon/core/utils/fields_comparator.h"
@@ -95,11 +96,8 @@ class MergeTreeWriter : public BatchWriter {
     Status UpdateCompactResult(const std::shared_ptr<CompactResult>& compact_result);
     Status UpdateCompactDeletionFile(const std::shared_ptr<CompactDeletionFile>& new_deletion_file);
 
-    static Result<int64_t> EstimateMemoryUse(const std::shared_ptr<arrow::Array>& array);
-
  private:
     int64_t last_sequence_number_;
-    int64_t current_memory_in_bytes_;
     std::shared_ptr<MemoryPool> pool_;
     std::vector<std::string> trimmed_primary_keys_;
     CoreOptions options_;
@@ -109,13 +107,11 @@ class MergeTreeWriter : public BatchWriter {
     std::shared_ptr<MergeFunctionWrapper<KeyValue>> merge_function_wrapper_;
     int64_t schema_id_;
     // write_schema = value_schema + special fields
-    std::shared_ptr<arrow::DataType> value_type_;
     std::shared_ptr<arrow::Schema> write_schema_;
 
     std::shared_ptr<CompactManager> compact_manager_;
 
-    std::vector<std::shared_ptr<arrow::StructArray>> batch_vec_;
-    std::vector<std::vector<RecordBatch::RowKind>> row_kinds_vec_;
+    std::unique_ptr<WriteBuffer> write_buffer_;
 
     std::shared_ptr<Metrics> metrics_;
 

src/paimon/core/mergetree/merge_tree_writer_test.cpp

@@ -21,17 +21,14 @@
 #include <map>
 #include <optional>
 #include <utility>
-#include <variant>
 
 #include "arrow/api.h"
 #include "arrow/array/array_base.h"
-#include "arrow/array/array_nested.h"
 #include "arrow/c/abi.h"
 #include "arrow/c/bridge.h"
 #include "arrow/ipc/json_simple.h"
 #include "gtest/gtest.h"
 #include "paimon/common/factories/io_hook.h"
-#include "paimon/common/fs/external_path_provider.h"
 #include "paimon/common/table/special_fields.h"
 #include "paimon/common/types/data_field.h"
 #include "paimon/common/utils/scope_guard.h"
@@ -50,7 +47,6 @@
 #include "paimon/fs/file_system.h"
 #include "paimon/fs/local/local_file_system.h"
 #include "paimon/memory/memory_pool.h"
-#include "paimon/metrics.h"
 #include "paimon/testing/utils/binary_row_generator.h"
 #include "paimon/testing/utils/io_exception_helper.h"
 #include "paimon/testing/utils/read_result_collector.h"
@@ -88,14 +84,21 @@ class MergeTreeWriterTest : public ::testing::Test {
     }
     void TearDown() override {}
 
-    void WriteBatch(const std::shared_ptr<arrow::Array>& array,
-                    const std::vector<RecordBatch::RowKind>& row_kinds,
-                    MergeTreeWriter* writer) const {
+    std::unique_ptr<RecordBatch> CreateBatch(
+        const std::shared_ptr<arrow::Array>& array,
+        const std::vector<RecordBatch::RowKind>& row_kinds) const {
         ::ArrowArray c_array;
-        ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok());
+        EXPECT_TRUE(arrow::ExportArray(*array, &c_array).ok());
         RecordBatchBuilder batch_builder(&c_array);
         batch_builder.SetRowKinds(row_kinds);
-        ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch, batch_builder.Finish());
+        EXPECT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch, batch_builder.Finish());
+        return batch;
+    }
+
+    void WriteBatch(const std::shared_ptr<arrow::Array>& array,
+                    const std::vector<RecordBatch::RowKind>& row_kinds,
+                    MergeTreeWriter* writer) const {
+        auto batch = CreateBatch(array, row_kinds);
         ASSERT_OK(writer->Write(std::move(batch)));
     }
 
@@ -761,76 +764,6 @@ TEST_F(MergeTreeWriterTest, TestIOException) {
     ASSERT_TRUE(run_complete);
 }
 
-TEST_F(MergeTreeWriterTest, TestEstimateMemoryUse) {
-    {
-        // test simple
-        std::shared_ptr<arrow::Array> array =
-            arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([
-          ["Lucy", 20, 1, 14.1],
-          ["Paul", 20, 1, null],
-          ["Alice", 10, 0, 13.1]
-        ])")
-                .ValueOrDie();
-        ASSERT_OK_AND_ASSIGN(int64_t memory_use, MergeTreeWriter::EstimateMemoryUse(array));
-        int64_t expected_memory_use =
-            1 + (13 + 3 * 4 + 1) + (3 * 4 + 1) + (3 * 4 + 1) + (3 * 8 + 1);
-        ASSERT_EQ(memory_use, expected_memory_use);
-    }
-    {
-        // test primitive type
-        arrow::FieldVector fields = {arrow::field("v0", arrow::boolean()),
-                                     arrow::field("v1", arrow::int8()),
-                                     arrow::field("v2", arrow::int16()),
-                                     arrow::field("v3", arrow::int32()),
-                                     arrow::field("v4", arrow::int64()),
-                                     arrow::field("v5", arrow::float32()),
-                                     arrow::field("v6", arrow::float64()),
-                                     arrow::field("v7", arrow::date32()),
-                                     arrow::field("v8", arrow::timestamp(arrow::TimeUnit::NANO)),
-                                     arrow::field("v9", arrow::decimal128(30, 20)),
-                                     arrow::field("v10", arrow::utf8()),
-                                     arrow::field("v11", arrow::binary())};
-
-        auto array = std::dynamic_pointer_cast<arrow::StructArray>(
-            arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
-        [true, 10, 200, 65536, 123456789, 0.0, 0.0, 2000, -86399999999500, "2134.48690000000000000009", "difference", "Alice"],
-        [false, -128, -32768, -2147483648, -9223372036854775808, -3.4028235E38, -1.7976931348623157E308, -719528, -9223372036854775808, "-999999999999999999.99999999999999999999", "Alice", "Two"],
-        [true, 127, 32767, 2147483647, 9223372036854775807, 3.4028235E38, 1.7976931348623157E308, 2932896, 9223372036854775807, "999999999999999999.99999999999999999999", "Alice", "made"],
-        [true, 0, 0, 0, 0, 1.4E-45, 4.9E-324, 0, 0, "0.00000000000000000000", "Alice", "wood"]
-])")
-                .ValueOrDie());
-        ASSERT_OK_AND_ASSIGN(int64_t memory_use, MergeTreeWriter::EstimateMemoryUse(array));
-        int64_t expected_memory_use = 1 + (4 + 1) + (4 + 1) + (2 * 4 + 1) + (4 * 4 + 1) +
-                                      (8 * 4 + 1) + (4 * 4 + 1) + (8 * 4 + 1) + (4 * 4 + 1) +
-                                      (8 * 4 + 1) + (4 * 16 + 1) + (25 + 4 * 4 + 1) +
-                                      (16 + 4 * 4 + 1);
-        ASSERT_EQ(memory_use, expected_memory_use);
-    }
-    {
-        // test nested type
-        arrow::FieldVector fields = {
-            arrow::field("f0", arrow::list(arrow::int32())),
-            arrow::field("f1", arrow::map(arrow::utf8(), arrow::int64())),
-            arrow::field("f2", arrow::struct_({arrow::field("sub1", arrow::int64()),
-                                               arrow::field("sub2", arrow::float64()),
-                                               arrow::field("sub3", arrow::boolean())})),
-        };
-        auto array = std::dynamic_pointer_cast<arrow::StructArray>(
-            arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields}), R"([
-        [[1, 2, 3],    [["apple", 3], ["banana", 4]],          [10, 10.1, false]],
-        [[4, 5],       [["cat", 5], ["dog", 6], ["mouse", 7]], [20, 20.1, true]],
-        [[6],          [["elephant", 7], ["fox", 8]],          [null, 30.1, true]]
-    ])")
-                .ValueOrDie());
-        ASSERT_OK_AND_ASSIGN(int64_t memory_use, MergeTreeWriter::EstimateMemoryUse(array));
-        int64_t list_mem = 1 + (4 * 6 + 1);
-        int64_t map_mem = 1 + (33 + 4 * 7 + 1) + (8 * 7 + 1);
-        int64_t struct_mem = 1 + (8 * 3 + 1) + (8 * 3 + 1) + (1 * 3 + 1);
-        int64_t expected_memory_use = 1 + list_mem + map_mem + struct_mem;
-        ASSERT_EQ(memory_use, expected_memory_use);
-    }
-}
-
 TEST_F(MergeTreeWriterTest, TestBulkData) {
     // each batch is a file due to WRITE_BUFFER_SIZE
     ASSERT_OK_AND_ASSIGN(