fix review

Author: zjw1111

src/paimon/core/mergetree/merge_tree_writer.cpp

@@ -226,7 +226,7 @@ Status MergeTreeWriter::Flush(bool wait_for_latest_compaction, bool forced_full_
         }
         // 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_));
+                               write_buffer_->DrainToReaders(&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_,

src/paimon/core/mergetree/write_buffer.cpp

@@ -39,12 +39,12 @@ WriteBuffer::WriteBuffer(
     const std::shared_ptr<FieldsComparator>& key_comparator,
     const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
     const std::shared_ptr<MemoryPool>& pool)
-    : value_type_(value_type),
+    : pool_(pool),
+      value_type_(value_type),
       trimmed_primary_keys_(trimmed_primary_keys),
       user_defined_sequence_fields_(user_defined_sequence_fields),
       key_comparator_(key_comparator),
-      merge_function_wrapper_(merge_function_wrapper),
-      pool_(pool) {}
+      merge_function_wrapper_(merge_function_wrapper) {}
 
 Status WriteBuffer::Write(std::unique_ptr<RecordBatch>&& moved_batch) {
     if (ArrowArrayIsReleased(moved_batch->GetData())) {
@@ -66,7 +66,7 @@ Status WriteBuffer::Write(std::unique_ptr<RecordBatch>&& moved_batch) {
     return Status::OK();
 }
 
-Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> WriteBuffer::Flush(
+Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> WriteBuffer::DrainToReaders(
     int64_t* last_sequence_number) {
     std::vector<std::unique_ptr<KeyValueRecordReader>> readers;
     if (batch_vec_.empty()) {
@@ -84,10 +84,7 @@ Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> WriteBuffer::Flush(
         readers.push_back(std::move(in_memory_reader));
     }
 
-    batch_vec_.clear();
-    row_kinds_vec_.clear();
-    current_memory_in_bytes_ = 0;
-
+    Clear();
     return readers;
 }
 
@@ -97,6 +94,8 @@ void WriteBuffer::Clear() {
     current_memory_in_bytes_ = 0;
 }
 
+// TODO(jinli.zjw): Consider making the memory estimation more accurate.
+// https://github.com/alibaba/paimon-cpp/pull/206#discussion_r3021325389
 Result<int64_t> WriteBuffer::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;

src/paimon/core/mergetree/write_buffer.h

@@ -57,10 +57,11 @@ class WriteBuffer {
     /// Does NOT check memory thresholds or trigger flush.
     Status Write(std::unique_ptr<RecordBatch>&& batch);
 
-    /// Flush all buffered batches into KeyValueInMemoryRecordReaders and clear the buffer.
-    /// @param[in,out] last_sequence_number current sequence number, updated after flush
+    /// Drain all buffered batches into KeyValueInMemoryRecordReaders and clear the buffer.
+    /// @param[in,out] last_sequence_number current sequence number, updated after draining
     /// @return list of KeyValueRecordReaders built from buffered data
-    Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> Flush(int64_t* last_sequence_number);
+    Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> DrainToReaders(
+        int64_t* last_sequence_number);
 
     /// Return current memory usage in bytes.
     int64_t GetMemoryUsage() const {
@@ -80,12 +81,12 @@ class WriteBuffer {
     static Result<int64_t> EstimateMemoryUse(const std::shared_ptr<arrow::Array>& array);
 
     // Immutable configuration
+    const std::shared_ptr<MemoryPool> pool_;
     const std::shared_ptr<arrow::DataType> value_type_;
     const std::vector<std::string> trimmed_primary_keys_;
     const std::vector<std::string> user_defined_sequence_fields_;
     const std::shared_ptr<FieldsComparator> key_comparator_;
     const std::shared_ptr<MergeFunctionWrapper<KeyValue>> merge_function_wrapper_;
-    const std::shared_ptr<MemoryPool> pool_;
 
     // Mutable buffer state
     std::vector<std::shared_ptr<arrow::StructArray>> batch_vec_;

src/paimon/core/mergetree/write_buffer_test.cpp

@@ -39,6 +39,11 @@ class MergeFunctionWrapper;
 }  // namespace paimon
 
 namespace paimon::test {
+struct ReaderResult {
+    std::vector<int64_t> sequence_numbers;
+    std::vector<int8_t> row_kind_values;
+};
+
 class WriteBufferTest : public ::testing::Test {
  public:
     void SetUp() override {
@@ -69,6 +74,18 @@ class WriteBufferTest : public ::testing::Test {
         return batch;
     }
 
+    Result<ReaderResult> ReadReaderResult(KeyValueRecordReader* reader) const {
+        PAIMON_ASSIGN_OR_RAISE(auto iterator, reader->NextBatch());
+
+        ReaderResult result;
+        while (iterator->HasNext()) {
+            PAIMON_ASSIGN_OR_RAISE(KeyValue key_value, iterator->Next());
+            result.sequence_numbers.push_back(key_value.sequence_number);
+            result.row_kind_values.push_back(key_value.value_kind->ToByteValue());
+        }
+        return result;
+    }
+
  protected:
     std::shared_ptr<MemoryPool> pool_;
     std::vector<DataField> value_fields_;
@@ -100,32 +117,23 @@ TEST_F(WriteBufferTest, TestFlushResetsStateAndAdvancesSequenceNumber) {
     ASSERT_GT(write_buffer.GetMemoryUsage(), 0);
 
     int64_t last_sequence_number = 10;
-    ASSERT_OK_AND_ASSIGN(auto readers, write_buffer.Flush(&last_sequence_number));
+    ASSERT_OK_AND_ASSIGN(auto readers, write_buffer.DrainToReaders(&last_sequence_number));
 
     ASSERT_EQ(readers.size(), 2);
     ASSERT_TRUE(write_buffer.IsEmpty());
     ASSERT_EQ(write_buffer.GetMemoryUsage(), 0);
     ASSERT_EQ(last_sequence_number, 13);
 
-    ASSERT_OK_AND_ASSIGN(auto first_iterator, readers[0]->NextBatch());
-    ASSERT_TRUE(first_iterator);
-    ASSERT_TRUE(first_iterator->HasNext());
-    ASSERT_OK_AND_ASSIGN(KeyValue first, first_iterator->Next());
-    ASSERT_EQ(first.sequence_number, 10);
-    ASSERT_EQ(first.value_kind->ToByteValue(), RowKind::Insert()->ToByteValue());
-    ASSERT_TRUE(first_iterator->HasNext());
-    ASSERT_OK_AND_ASSIGN(KeyValue second, first_iterator->Next());
-    ASSERT_EQ(second.sequence_number, 11);
-    ASSERT_EQ(second.value_kind->ToByteValue(), RowKind::Insert()->ToByteValue());
-    ASSERT_FALSE(first_iterator->HasNext());
-
-    ASSERT_OK_AND_ASSIGN(auto second_iterator, readers[1]->NextBatch());
-    ASSERT_TRUE(second_iterator);
-    ASSERT_TRUE(second_iterator->HasNext());
-    ASSERT_OK_AND_ASSIGN(KeyValue third, second_iterator->Next());
-    ASSERT_EQ(third.sequence_number, 12);
-    ASSERT_EQ(third.value_kind->ToByteValue(), RowKind::Insert()->ToByteValue());
-    ASSERT_FALSE(second_iterator->HasNext());
+    ASSERT_OK_AND_ASSIGN(auto first_result, ReadReaderResult(readers[0].get()));
+    ASSERT_EQ(first_result.sequence_numbers, (std::vector<int64_t>{10, 11}));
+    ASSERT_EQ(
+        first_result.row_kind_values,
+        (std::vector<int8_t>{RowKind::Insert()->ToByteValue(), RowKind::Insert()->ToByteValue()}));
+
+    ASSERT_OK_AND_ASSIGN(auto second_result, ReadReaderResult(readers[1].get()));
+    ASSERT_EQ(second_result.sequence_numbers, (std::vector<int64_t>{12}));
+    ASSERT_EQ(second_result.row_kind_values,
+              (std::vector<int8_t>{RowKind::Insert()->ToByteValue()}));
 }
 
 TEST_F(WriteBufferTest, TestFlushPreservesRowKinds) {
@@ -150,26 +158,17 @@ TEST_F(WriteBufferTest, TestFlushPreservesRowKinds) {
     ASSERT_OK(write_buffer.Write(CreateBatch(array, row_kinds)));
 
     int64_t last_sequence_number = 0;
-    ASSERT_OK_AND_ASSIGN(auto readers, write_buffer.Flush(&last_sequence_number));
+    ASSERT_OK_AND_ASSIGN(auto readers, write_buffer.DrainToReaders(&last_sequence_number));
     ASSERT_EQ(readers.size(), 1);
     ASSERT_EQ(last_sequence_number, 4);
 
-    ASSERT_OK_AND_ASSIGN(auto iterator, readers[0]->NextBatch());
-    ASSERT_TRUE(iterator);
-
-    std::vector<int8_t> actual_row_kind_values;
-    std::vector<int64_t> actual_sequence_numbers;
-    while (iterator->HasNext()) {
-        ASSERT_OK_AND_ASSIGN(KeyValue key_value, iterator->Next());
-        actual_row_kind_values.push_back(key_value.value_kind->ToByteValue());
-        actual_sequence_numbers.push_back(key_value.sequence_number);
-    }
+    ASSERT_OK_AND_ASSIGN(auto reader_result, ReadReaderResult(readers[0].get()));
 
-    ASSERT_EQ(actual_row_kind_values,
+    ASSERT_EQ(reader_result.row_kind_values,
               (std::vector<int8_t>{
                   RowKind::Insert()->ToByteValue(), RowKind::UpdateBefore()->ToByteValue(),
                   RowKind::UpdateAfter()->ToByteValue(), RowKind::Delete()->ToByteValue()}));
-    ASSERT_EQ(actual_sequence_numbers, (std::vector<int64_t>{0, 1, 2, 3}));
+    ASSERT_EQ(reader_result.sequence_numbers, (std::vector<int64_t>{0, 1, 2, 3}));
 }
 
 TEST_F(WriteBufferTest, TestEstimateMemoryUse) {