refactor(mergetree): introduce BinaryInMemorySortBuffer and MergedKeyValueRecordReader

Refactor write_buffer to use the new SortBuffer abstraction (BinaryInMemorySortBuffer / BinaryExternalSortBuffer) and introduce MergedKeyValueRecordReader for merged key-value reading. Also adapt key_value_in_memory_record_reader, merge_tree_writer, spill reader/writer and core_options accordingly.

Author: zjw1111

src/paimon/CMakeLists.txt

@@ -213,6 +213,7 @@ set(PAIMON_CORE_SRCS
     core/io/key_value_data_file_record_reader.cpp
     core/io/key_value_data_file_writer.cpp
     core/io/key_value_in_memory_record_reader.cpp
+    core/io/merged_key_value_record_reader.cpp
     core/io/key_value_meta_projection_consumer.cpp
     core/io/key_value_projection_consumer.cpp
     core/io/key_value_projection_reader.cpp
@@ -246,6 +247,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/binary_in_memory_sort_buffer.cpp
     core/mergetree/write_buffer.cpp
     core/mergetree/levels.cpp
     core/mergetree/lookup_file.cpp
@@ -564,6 +566,7 @@ if(PAIMON_BUILD_TESTS)
                     core/io/key_value_data_file_record_reader_test.cpp
                     core/io/key_value_projection_reader_test.cpp
                     core/io/key_value_in_memory_record_reader_test.cpp
+                    core/io/merged_key_value_record_reader_test.cpp
                     core/io/complete_row_tracking_fields_reader_test.cpp
                     core/io/data_file_meta_test.cpp
                     core/io/file_index_evaluator_test.cpp

src/paimon/core/append/append_only_writer.h

@@ -68,6 +68,10 @@ class AppendOnlyWriter : public BatchWriter {
     Status Compact(bool full_compaction) override {
         return Flush(/*wait_for_latest_compaction=*/true, full_compaction);
     }
+
+    Status FlushMemory() override {
+        return Flush(/*wait_for_latest_compaction=*/false, /*forced_full_compaction=*/false);
+    }
     Result<CommitIncrement> PrepareCommit(bool wait_compaction) override;
     Result<bool> CompactNotCompleted() override {
         PAIMON_RETURN_NOT_OK(compact_manager_->TriggerCompaction(/*full_compaction=*/false));

src/paimon/core/io/key_value_in_memory_record_reader.cpp

@@ -17,7 +17,6 @@
 #include "paimon/core/io/key_value_in_memory_record_reader.h"
 
 #include <cassert>
-#include <optional>
 #include <utility>
 
 #include "arrow/array/array_base.h"
@@ -28,60 +27,44 @@
 #include "arrow/util/checked_cast.h"
 #include "fmt/format.h"
 #include "paimon/common/data/columnar/columnar_row_ref.h"
-#include "paimon/common/data/internal_row.h"
 #include "paimon/common/types/row_kind.h"
 #include "paimon/common/utils/arrow/arrow_utils.h"
 #include "paimon/common/utils/arrow/status_utils.h"
 #include "paimon/common/utils/fields_comparator.h"
-#include "paimon/core/mergetree/compact/merge_function_wrapper.h"
 #include "paimon/status.h"
 namespace paimon {
 class MemoryPool;
 
 Result<KeyValue> KeyValueInMemoryRecordReader::Iterator::Next() {
-    reader_->merge_function_wrapper_->Reset();
-    std::shared_ptr<InternalRow> current_key;
-    while (cursor_ < reader_->value_struct_array_->length()) {
-        uint64_t index = reader_->sort_indices_->Value(cursor_);
-        const RowKind* row_kind = RowKind::Insert();
-        if (!reader_->row_kinds_.empty()) {
-            PAIMON_ASSIGN_OR_RAISE(
-                row_kind, RowKind::FromByteValue(static_cast<int8_t>(reader_->row_kinds_[index])));
-        }
-        // key must hold value_struct_array as min/max key may be used after projection
-        auto key = std::make_unique<ColumnarRowRef>(reader_->key_ctx_, index);
-        auto value = std::make_unique<ColumnarRowRef>(reader_->value_ctx_, index);
-        KeyValue kv(row_kind, reader_->last_sequence_num_ + index,
-                    /*level=*/KeyValue::UNKNOWN_LEVEL, std::move(key), std::move(value));
-        if (current_key == nullptr) {
-            current_key = kv.key;
-        } else if (reader_->key_comparator_->CompareTo(*current_key, *kv.key) != 0) {
-            break;
-        }
-        PAIMON_RETURN_NOT_OK(reader_->merge_function_wrapper_->Add(std::move(kv)));
-        cursor_++;
+    uint64_t index = reader_->sort_indices_->Value(cursor_++);
+    const RowKind* row_kind = RowKind::Insert();
+    if (!reader_->row_kinds_.empty()) {
+        PAIMON_ASSIGN_OR_RAISE(
+            row_kind, RowKind::FromByteValue(static_cast<int8_t>(reader_->row_kinds_[index])));
     }
-    PAIMON_ASSIGN_OR_RAISE(std::optional<KeyValue> result,
-                           reader_->merge_function_wrapper_->GetResult());
-    assert(result != std::nullopt);
-    return std::move(result).value();
+
+    // key must hold value_struct_array as min/max key may be used after projection
+    auto key = std::make_unique<ColumnarRowRef>(reader_->key_ctx_, index);
+    auto value = std::make_unique<ColumnarRowRef>(reader_->value_ctx_, index);
+    return KeyValue(row_kind, reader_->last_sequence_num_ + index,
+                    /*level=*/KeyValue::UNKNOWN_LEVEL, std::move(key), std::move(value));
 }
 
 KeyValueInMemoryRecordReader::KeyValueInMemoryRecordReader(
-    int64_t last_sequence_num, std::shared_ptr<arrow::StructArray>&& struct_array,
-    std::vector<RecordBatch::RowKind>&& row_kinds, const std::vector<std::string>& primary_keys,
-    const std::vector<std::string>& user_defined_sequence_fields,
+    int64_t last_sequence_num, const std::shared_ptr<arrow::StructArray>& struct_array,
+    const std::vector<RecordBatch::RowKind>& row_kinds,
+    const std::vector<std::string>& primary_keys,
+    const std::vector<std::string>& user_defined_sequence_fields, bool sequence_fields_ascending,
     const std::shared_ptr<FieldsComparator>& key_comparator,
-    const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
     const std::shared_ptr<MemoryPool>& pool)
     : last_sequence_num_(last_sequence_num),
       primary_keys_(primary_keys),
       user_defined_sequence_fields_(user_defined_sequence_fields),
+      sequence_fields_ascending_(sequence_fields_ascending),
       pool_(pool),
-      value_struct_array_(std::move(struct_array)),
-      row_kinds_(std::move(row_kinds)),
-      key_comparator_(key_comparator),
-      merge_function_wrapper_(merge_function_wrapper) {
+      value_struct_array_(struct_array),
+      row_kinds_(row_kinds),
+      key_comparator_(key_comparator) {
     assert(value_struct_array_);
     ArrowUtils::TraverseArray(value_struct_array_);
 }
@@ -113,6 +96,7 @@ Result<std::unique_ptr<KeyValueRecordReader::Iterator>> KeyValueInMemoryRecordRe
 }
 
 void KeyValueInMemoryRecordReader::Close() {
+    visited_ = true;
     value_struct_array_.reset();
     row_kinds_.clear();
     sort_indices_.reset();
@@ -127,8 +111,11 @@ KeyValueInMemoryRecordReader::SortBatch() const {
     for (const auto& name : primary_keys_) {
         sort_keys.emplace_back(name, arrow::compute::SortOrder::Ascending);
     }
+    const auto sequence_sort_order = sequence_fields_ascending_
+                                         ? arrow::compute::SortOrder::Ascending
+                                         : arrow::compute::SortOrder::Descending;
     for (const auto& name : user_defined_sequence_fields_) {
-        sort_keys.emplace_back(name, arrow::compute::SortOrder::Ascending);
+        sort_keys.emplace_back(name, sequence_sort_order);
     }
     auto sort_options =
         arrow::compute::SortOptions(sort_keys, arrow::compute::NullPlacement::AtStart);

src/paimon/core/io/key_value_in_memory_record_reader.h

@@ -28,7 +28,6 @@
 #include "paimon/common/utils/fields_comparator.h"
 #include "paimon/core/io/key_value_record_reader.h"
 #include "paimon/core/key_value.h"
-#include "paimon/core/mergetree/compact/merge_function_wrapper.h"
 #include "paimon/record_batch.h"
 #include "paimon/result.h"
 
@@ -41,18 +40,17 @@ struct ColumnarBatchContext;
 class FieldsComparator;
 class MemoryPool;
 class Metrics;
-template <typename T>
-class MergeFunctionWrapper;
 
 class KeyValueInMemoryRecordReader : public KeyValueRecordReader {
  public:
-    KeyValueInMemoryRecordReader(
-        int64_t last_sequence_num, std::shared_ptr<arrow::StructArray>&& struct_array,
-        std::vector<RecordBatch::RowKind>&& row_kinds, const std::vector<std::string>& 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);
+    KeyValueInMemoryRecordReader(int64_t last_sequence_num,
+                                 const std::shared_ptr<arrow::StructArray>& struct_array,
+                                 const std::vector<RecordBatch::RowKind>& row_kinds,
+                                 const std::vector<std::string>& primary_keys,
+                                 const std::vector<std::string>& user_defined_sequence_fields,
+                                 bool sequence_fields_ascending,
+                                 const std::shared_ptr<FieldsComparator>& key_comparator,
+                                 const std::shared_ptr<MemoryPool>& pool);
 
     class Iterator : public KeyValueRecordReader::Iterator {
      public:
@@ -83,11 +81,11 @@ class KeyValueInMemoryRecordReader : public KeyValueRecordReader {
     int64_t last_sequence_num_ = -1;
     std::vector<std::string> primary_keys_;
     std::vector<std::string> user_defined_sequence_fields_;
+    bool sequence_fields_ascending_ = true;
     std::shared_ptr<MemoryPool> pool_;
     std::shared_ptr<arrow::StructArray> value_struct_array_;
     std::vector<RecordBatch::RowKind> row_kinds_;
     std::shared_ptr<FieldsComparator> key_comparator_;
-    std::shared_ptr<MergeFunctionWrapper<KeyValue>> merge_function_wrapper_;
 
     std::shared_ptr<arrow::NumericArray<arrow::UInt64Type>> sort_indices_;
     std::shared_ptr<ColumnarBatchContext> key_ctx_;