feat(spill): optimize external sort with leveled merger to reduce write amplification

zjw1111 · claude · zjw1111 · commit ac25c247b7eb · 2026-05-18T14:12:21.000+08:00
Introduce LeveledMerger (LSM-tree-like structure) to manage spill files in levels,
reducing read/write amplification from O(N/K) to O(log_K(N)). Also adds dynamic
estimation of spill parameters based on actual row sizes.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt
@@ -253,6 +253,7 @@ set(PAIMON_CORE_SRCS
     core/mergetree/merge_tree_writer.cpp
     core/mergetree/in_memory_sort_buffer.cpp
     core/mergetree/external_sort_buffer.cpp
+    core/mergetree/leveled_merger.cpp
     core/mergetree/write_buffer.cpp
     core/mergetree/levels.cpp
     core/mergetree/lookup_file.cpp
@@ -648,6 +649,7 @@ if(PAIMON_BUILD_TESTS)
                     core/mergetree/merge_tree_writer_test.cpp
                     core/mergetree/write_buffer_test.cpp
                     core/mergetree/sort_buffer_test.cpp
+                    core/mergetree/leveled_merger_test.cpp
                     core/mergetree/sorted_run_test.cpp
                     core/mergetree/spill_channel_manager_test.cpp
                     core/mergetree/spill_reader_writer_test.cpp
diff --git a/src/paimon/core/core_options.cpp b/src/paimon/core/core_options.cpp
@@ -498,6 +498,11 @@ struct CoreOptions::Impl {
         // Parse local-sort.max-num-file-handles - spill file handle cap for local merge
         PAIMON_RETURN_NOT_OK(parser.Parse(Options::LOCAL_SORT_MAX_NUM_FILE_HANDLES,
                                           &local_sort_max_num_file_handles));
+        if (local_sort_max_num_file_handles < kLocalSortFileHandlesMinimalLimit) {
+            return Status::Invalid(fmt::format(
+                "invalid '{}': {}, must be at least {}", Options::LOCAL_SORT_MAX_NUM_FILE_HANDLES,
+                local_sort_max_num_file_handles, kLocalSortFileHandlesMinimalLimit));
+        }
         // Parse spill-compression - compression codec for spill files
         PAIMON_RETURN_NOT_OK(
             parser.Parse(Options::SPILL_COMPRESSION, &spill_compress_options.compress));
diff --git a/src/paimon/core/core_options.h b/src/paimon/core/core_options.h
@@ -181,6 +181,8 @@ class PAIMON_EXPORT CoreOptions {
 
     const std::map<std::string, std::string>& ToMap() const;
 
+    static constexpr int32_t kLocalSortFileHandlesMinimalLimit = 2;
+
  private:
     std::optional<std::string> GetDataFileExternalPaths() const;
     std::optional<std::string> GetGlobalIndexExternalPath() const;
diff --git a/src/paimon/core/disk/file_io_channel.h b/src/paimon/core/disk/file_io_channel.h
@@ -17,7 +17,6 @@
 #pragma once
 
 #include <cstdint>
-#include <memory>
 #include <random>
 #include <string>
 
@@ -68,4 +67,9 @@ class PAIMON_EXPORT FileIOChannel {
     };
 };
 
+struct FileChannelInfo {
+    FileIOChannel::ID channel_id;
+    int64_t file_size;
+};
+
 }  // namespace paimon
diff --git a/src/paimon/core/mergetree/external_sort_buffer.cpp b/src/paimon/core/mergetree/external_sort_buffer.cpp
@@ -16,6 +16,7 @@
 
 #include "paimon/core/mergetree/external_sort_buffer.h"
 
+#include <algorithm>
 #include <cassert>
 #include <utility>
 
@@ -80,9 +81,13 @@ ExternalSortBuffer::ExternalSortBuffer(
       user_defined_seq_comparator_(user_defined_seq_comparator),
       write_schema_(SpecialFields::CompleteSequenceAndValueKindField(value_schema)),
       options_(options),
-      spill_channel_manager_(std::make_shared<SpillChannelManager>(
-          options_.GetFileSystem(), options_.GetLocalSortMaxNumFileHandles())),
-      spill_channel_enumerator_(spill_channel_enumerator) {}
+      max_fan_in_(options.GetLocalSortMaxNumFileHandles()),
+      spill_channel_manager_(
+          std::make_shared<SpillChannelManager>(options_.GetFileSystem(), max_fan_in_)),
+      leveled_merger_(std::make_unique<LeveledMerger>(max_fan_in_)),
+      spill_channel_enumerator_(spill_channel_enumerator),
+      actual_max_fan_in_(max_fan_in_),
+      spill_batch_size_(options_.GetWriteBatchSize()) {}
 
 ExternalSortBuffer::~ExternalSortBuffer() {
     DoClear();
@@ -94,7 +99,10 @@ bool ExternalSortBuffer::HasSpilledData() const {
 
 void ExternalSortBuffer::DoClear() {
     in_memory_buffer_->Clear();
-    CleanupSpillFiles();
+
+    spill_channel_manager_->Reset();
+    total_spill_disk_bytes_ = 0;
+    leveled_merger_->Clear();
 }
 
 void ExternalSortBuffer::Clear() {
@@ -105,18 +113,42 @@ uint64_t ExternalSortBuffer::GetMemorySize() const {
     return in_memory_buffer_->GetMemorySize();
 }
 
+void ExternalSortBuffer::EstimateSpillParameters() {
+    int64_t estimated_row_size = in_memory_buffer_->GetEstimateMemoryUseForEachRow();
+    if (estimated_row_size <= 0) {
+        return;
+    }
+
+    const int32_t max_batch_size = options_.GetWriteBatchSize();
+    const int32_t min_batch_size = std::min(kSpillMinBatchSize, max_batch_size);
+    const int64_t merge_budget = options_.GetWriteBufferSize();
+    const int64_t max_memory_use_per_handle = merge_budget / max_fan_in_;
+
+    spill_batch_size_ = max_memory_use_per_handle / estimated_row_size;
+    spill_batch_size_ = std::clamp(spill_batch_size_, min_batch_size, max_batch_size);
+
+    actual_max_fan_in_ = merge_budget / (spill_batch_size_ * estimated_row_size);
+    actual_max_fan_in_ =
+        std::clamp(actual_max_fan_in_, CoreOptions::kLocalSortFileHandlesMinimalLimit, max_fan_in_);
+
+    // Re-derive spill_batch_size_ from the clamped actual_max_fan_in_ to stay within merge_budget.
+    spill_batch_size_ = merge_budget / (actual_max_fan_in_ * estimated_row_size);
+    spill_batch_size_ = std::clamp(spill_batch_size_, 1, max_batch_size);
+
+    leveled_merger_->SetMaxFanIn(actual_max_fan_in_);
+}
+
 Result<bool> ExternalSortBuffer::FlushMemory() {
     if (!in_memory_buffer_->HasData()) {
         return true;
     }
 
-    int64_t max_spill_disk_size = options_.GetWriteBufferSpillMaxDiskSize();
-
+    EstimateSpillParameters();
     PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> memory_buffer_readers,
                            in_memory_buffer_->CreateReaders());
     PAIMON_RETURN_NOT_OK(SpillMemoryBuffer(std::move(memory_buffer_readers)));
     in_memory_buffer_->Clear();
-    return total_spill_disk_bytes_ < max_spill_disk_size;
+    return total_spill_disk_bytes_ < options_.GetWriteBufferSpillMaxDiskSize();
 }
 
 Result<bool> ExternalSortBuffer::Write(std::unique_ptr<RecordBatch>&& batch) {
@@ -128,11 +160,17 @@ Result<bool> ExternalSortBuffer::Write(std::unique_ptr<RecordBatch>&& batch) {
 }
 
 Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> ExternalSortBuffer::CreateReaders() {
-    PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> readers,
-                           CollectSpillReaders());
     PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> memory_readers,
                            in_memory_buffer_->CreateReaders());
+    if (!HasSpilledData()) {
+        return memory_readers;
+    }
 
+    int32_t max_spill_files = actual_max_fan_in_ - 1;
+    PAIMON_RETURN_NOT_OK(
+        leveled_merger_->RunFinalCleanupIfNeeded(max_spill_files, CreateMergeFn()));
+    PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> readers,
+                           CreateSpillReaders(leveled_merger_->GetAllFiles()));
     readers.insert(readers.end(), std::make_move_iterator(memory_readers.begin()),
                    std::make_move_iterator(memory_readers.end()));
     return readers;
@@ -142,26 +180,20 @@ bool ExternalSortBuffer::HasData() const {
     return in_memory_buffer_->HasData() || HasSpilledData();
 }
 
-void ExternalSortBuffer::CleanupSpillFiles() {
-    spill_channel_manager_->Reset();
-    total_spill_disk_bytes_ = 0;
-}
-
-Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> ExternalSortBuffer::CollectSpillReaders()
-    const {
+Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> ExternalSortBuffer::CreateSpillReaders(
+    const std::vector<FileChannelInfo>& files) const {
     std::vector<std::unique_ptr<KeyValueRecordReader>> readers;
-    const auto& channel_ids = spill_channel_manager_->GetChannels();
-    readers.reserve(channel_ids.size());
-    for (const auto& channel_id : channel_ids) {
-        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<SpillReader> spill_reader,
+    readers.reserve(files.size());
+    for (const auto& file : files) {
+        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<SpillReader> reader,
                                SpillReader::Create(options_.GetFileSystem(), key_schema_,
-                                                   value_schema_, pool_, channel_id));
-        readers.push_back(std::move(spill_reader));
+                                                   value_schema_, pool_, file.channel_id));
+        readers.push_back(std::move(reader));
     }
     return readers;
 }
 
-Result<int64_t> ExternalSortBuffer::SpillToDisk(
+Result<FileChannelInfo> ExternalSortBuffer::SpillToDisk(
     std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers, int32_t write_batch_size) {
     const auto& spill_compress_options = options_.GetSpillCompressOptions();
     PAIMON_ASSIGN_OR_RAISE(
@@ -202,41 +234,37 @@ Result<int64_t> ExternalSortBuffer::SpillToDisk(
     PAIMON_RETURN_NOT_OK(spill_writer->Close());
     PAIMON_ASSIGN_OR_RAISE(int64_t spilled_file_size, spill_writer->GetFileSize());
     cleanup_guard.Release();
-    return spilled_file_size;
+    return FileChannelInfo{spill_writer->GetChannelId(), spilled_file_size};
 }
 
 Status ExternalSortBuffer::SpillMemoryBuffer(
     std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers) {
-    PAIMON_ASSIGN_OR_RAISE(int64_t spill_file_size,
-                           SpillToDisk(std::move(readers), options_.GetWriteBatchSize()));
-    total_spill_disk_bytes_ += spill_file_size;
-
-    if (options_.GetLocalSortMaxNumFileHandles() > 0 &&
-        static_cast<int32_t>(spill_channel_manager_->GetChannels().size()) >=
-            options_.GetLocalSortMaxNumFileHandles()) {
-        PAIMON_RETURN_NOT_OK(MergeSpilledFiles());
-    }
-    return Status::OK();
+    PAIMON_ASSIGN_OR_RAISE(FileChannelInfo file_info,
+                           SpillToDisk(std::move(readers), spill_batch_size_));
+    total_spill_disk_bytes_ += file_info.file_size;
+    leveled_merger_->AddFile(file_info);
+    return leveled_merger_->RunCompactionIfNeeded(CreateMergeFn());
 }
 
-Status ExternalSortBuffer::MergeSpilledFiles() {
-    if (spill_channel_manager_->GetChannels().size() < 2) {
-        return Status::OK();
-    }
-    auto spill_channel_ids_before_merge = spill_channel_manager_->GetChannels();
-    auto cleanup_guard = ScopeGuard([&]() {
-        for (const auto& spill_channel_id : spill_channel_ids_before_merge) {
-            [[maybe_unused]] auto status = spill_channel_manager_->DeleteChannel(spill_channel_id);
-        }
-    });
+LeveledMerger::MergeFn ExternalSortBuffer::CreateMergeFn() {
+    return [this](const std::vector<FileChannelInfo>& files) -> Result<FileChannelInfo> {
+        return MergeAndReplaceFiles(files);
+    };
+}
 
+Result<FileChannelInfo> ExternalSortBuffer::MergeAndReplaceFiles(
+    const std::vector<FileChannelInfo>& files) {
     PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> readers,
-                           CollectSpillReaders());
-    PAIMON_ASSIGN_OR_RAISE(int64_t merged_file_size,
-                           SpillToDisk(std::move(readers), options_.GetWriteBatchSize()));
-    total_spill_disk_bytes_ = merged_file_size;
+                           CreateSpillReaders(files));
+    PAIMON_ASSIGN_OR_RAISE(FileChannelInfo output,
+                           SpillToDisk(std::move(readers), spill_batch_size_));
+    total_spill_disk_bytes_ += output.file_size;
 
-    return Status::OK();
+    for (const auto& file : files) {
+        [[maybe_unused]] auto status = spill_channel_manager_->DeleteChannel(file.channel_id);
+        total_spill_disk_bytes_ -= file.file_size;
+    }
+    return output;
 }
 
 }  // namespace paimon
diff --git a/src/paimon/core/mergetree/external_sort_buffer.h b/src/paimon/core/mergetree/external_sort_buffer.h
@@ -25,6 +25,7 @@
 #include "paimon/core/core_options.h"
 #include "paimon/core/disk/file_io_channel.h"
 #include "paimon/core/mergetree/in_memory_sort_buffer.h"
+#include "paimon/core/mergetree/leveled_merger.h"
 #include "paimon/core/mergetree/sort_buffer.h"
 #include "paimon/record_batch.h"
 #include "paimon/result.h"
@@ -63,14 +64,18 @@ class ExternalSortBuffer : public SortBuffer {
     bool HasData() const override;
 
  private:
+    static constexpr int32_t kSpillMinBatchSize = 256;
+
     void DoClear();
+    void EstimateSpillParameters();
     bool HasSpilledData() const;
-    Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> CollectSpillReaders() const;
-    Result<int64_t> SpillToDisk(std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers,
-                                int32_t write_batch_size);
-    Status MergeSpilledFiles();
+    Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> CreateSpillReaders(
+        const std::vector<FileChannelInfo>& files) const;
+    Result<FileChannelInfo> SpillToDisk(
+        std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers, int32_t write_batch_size);
+    LeveledMerger::MergeFn CreateMergeFn();
+    Result<FileChannelInfo> MergeAndReplaceFiles(const std::vector<FileChannelInfo>& files);
     Status SpillMemoryBuffer(std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers);
-    void CleanupSpillFiles();
 
     ExternalSortBuffer(std::unique_ptr<InMemorySortBuffer>&& in_memory_buffer,
                        const std::shared_ptr<arrow::Schema>& key_schema,
@@ -90,10 +95,14 @@ class ExternalSortBuffer : public SortBuffer {
     const std::shared_ptr<FieldsComparator> user_defined_seq_comparator_;
     const std::shared_ptr<arrow::Schema> write_schema_;
     const CoreOptions options_;
+    const int32_t max_fan_in_;
     const std::shared_ptr<SpillChannelManager> spill_channel_manager_;
 
+    std::unique_ptr<LeveledMerger> leveled_merger_;
     std::shared_ptr<FileIOChannel::Enumerator> spill_channel_enumerator_;
     int64_t total_spill_disk_bytes_ = 0;
+    int32_t actual_max_fan_in_;
+    int32_t spill_batch_size_;
 };
 
 }  // namespace paimon
diff --git a/src/paimon/core/mergetree/in_memory_sort_buffer.cpp b/src/paimon/core/mergetree/in_memory_sort_buffer.cpp
@@ -53,6 +53,7 @@ InMemorySortBuffer::InMemorySortBuffer(int64_t last_sequence_number,
 void InMemorySortBuffer::Clear() {
     buffered_batches_.clear();
     current_memory_in_bytes_ = 0;
+    total_row_count_ = 0;
 }
 
 uint64_t InMemorySortBuffer::GetMemorySize() const {
@@ -82,7 +83,11 @@ Result<bool> InMemorySortBuffer::Write(std::unique_ptr<RecordBatch>&& moved_batc
     buffered_batch.struct_array = std::move(value_struct_array);
     buffered_batch.row_kinds = batch->GetRowKind();
     next_sequence_number_ += buffered_batch.struct_array->length();
+    total_row_count_ += buffered_batch.struct_array->length();
     buffered_batches_.push_back(std::move(buffered_batch));
+    if (total_row_count_ > 0) {
+        estimated_memory_use_for_each_row_ = current_memory_in_bytes_ / total_row_count_;
+    }
     return current_memory_in_bytes_ < write_buffer_size_;
 }
 
@@ -107,6 +112,10 @@ bool InMemorySortBuffer::HasData() const {
     return !buffered_batches_.empty();
 }
 
+uint64_t InMemorySortBuffer::GetEstimateMemoryUseForEachRow() const {
+    return estimated_memory_use_for_each_row_;
+}
+
 // TODO(jinli.zjw): Consider making the memory estimation more accurate.
 // https://github.com/alibaba/paimon-cpp/pull/206#discussion_r3021325389
 Result<int64_t> InMemorySortBuffer::EstimateMemoryUse(const std::shared_ptr<arrow::Array>& array) {
diff --git a/src/paimon/core/mergetree/in_memory_sort_buffer.h b/src/paimon/core/mergetree/in_memory_sort_buffer.h
@@ -63,11 +63,12 @@ class InMemorySortBuffer : public SortBuffer {
     Result<bool> Write(std::unique_ptr<RecordBatch>&& batch) override;
     Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> CreateReaders() override;
     bool HasData() const override;
+    uint64_t GetEstimateMemoryUseForEachRow() const;
 
+ private:
     /// Estimate memory usage of an Arrow array.
     static Result<int64_t> EstimateMemoryUse(const std::shared_ptr<arrow::Array>& array);
 
- private:
     const std::shared_ptr<MemoryPool> pool_;
     const std::shared_ptr<arrow::DataType> value_type_;
     const std::vector<std::string> trimmed_primary_keys_;
@@ -78,6 +79,8 @@ class InMemorySortBuffer : public SortBuffer {
 
     std::vector<BufferedWriteBatch> buffered_batches_;
     uint64_t current_memory_in_bytes_ = 0;
+    uint64_t estimated_memory_use_for_each_row_ = 0;
+    int64_t total_row_count_ = 0;
     int64_t next_sequence_number_ = 0;
 };
 
diff --git a/src/paimon/core/mergetree/leveled_merger.cpp b/src/paimon/core/mergetree/leveled_merger.cpp
diff --git a/src/paimon/core/mergetree/leveled_merger.h b/src/paimon/core/mergetree/leveled_merger.h
diff --git a/src/paimon/core/mergetree/leveled_merger_test.cpp b/src/paimon/core/mergetree/leveled_merger_test.cpp
diff --git a/src/paimon/core/mergetree/sort_buffer_test.cpp b/src/paimon/core/mergetree/sort_buffer_test.cpp
