fix review

Co-authored-by: Copilot <copilot@github.com>

Author: zjw1111

include/paimon/defs.h

@@ -182,7 +182,7 @@ struct PAIMON_EXPORT Options {
     static const char WRITE_BUFFER_SPILLABLE[];
 
     /// "write-buffer-spill.max-disk-size" - The max disk to use for write buffer spill. This only
-    /// work when the write buffer spill is enabled. Default value is unlimited.
+    /// works when the write buffer spill is enabled. Default value is unlimited.
     static const char WRITE_BUFFER_SPILL_MAX_DISK_SIZE[];
 
     /// "local-sort.max-num-file-handles" - The maximal fan-in for external merge sort. It limits

src/paimon/CMakeLists.txt

@@ -234,6 +234,7 @@ set(PAIMON_CORE_SRCS
     core/manifest/manifest_list.cpp
     core/manifest/partition_entry.cpp
     core/manifest/index_manifest_file_handler.cpp
+    core/memory/writer_memory_manager.cpp
     core/mergetree/compact/universal_compaction.cpp
     core/mergetree/compact/early_full_compaction.cpp
     core/mergetree/compact/aggregate/aggregate_merge_function.cpp
@@ -253,7 +254,6 @@ set(PAIMON_CORE_SRCS
     core/mergetree/in_memory_sort_buffer.cpp
     core/mergetree/external_sort_buffer.cpp
     core/mergetree/write_buffer.cpp
-    core/mergetree/writer_memory_manager.cpp
     core/mergetree/levels.cpp
     core/mergetree/lookup_file.cpp
     core/mergetree/lookup_levels.cpp
@@ -596,6 +596,7 @@ if(PAIMON_BUILD_TESTS)
                     core/manifest/file_entry_test.cpp
                     core/manifest/index_manifest_entry_serializer_test.cpp
                     core/manifest/index_manifest_file_handler_test.cpp
+                    core/memory/writer_memory_manager_test.cpp
                     core/mergetree/levels_test.cpp
                     core/mergetree/lookup_file_test.cpp
                     core/mergetree/lookup_levels_test.cpp
@@ -639,7 +640,6 @@ if(PAIMON_BUILD_TESTS)
                     core/mergetree/sorted_run_test.cpp
                     core/mergetree/spill_channel_manager_test.cpp
                     core/mergetree/spill_reader_writer_test.cpp
-                    core/mergetree/writer_memory_manager_test.cpp
                     core/migrate/file_meta_utils_test.cpp
                     core/operation/metrics/compaction_metrics_test.cpp
                     core/operation/data_evolution_file_store_scan_test.cpp

…core/mergetree/writer_memory_manager.cpp …on/core/memory/writer_memory_manager.cppsrc/paimon/core/mergetree/writer_memory_manager.cpp renamed to src/paimon/core/memory/writer_memory_manager.cpp src/paimon/core/mergetree/writer_memory_manager.cpp renamed to src/paimon/core/memory/writer_memory_manager.cpp

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "paimon/core/mergetree/writer_memory_manager.h"
+#include "paimon/core/memory/writer_memory_manager.h"
 
 #include <cassert>
 
@@ -44,7 +44,7 @@ void WriterMemoryManager::RefreshWriterMemory(BatchWriter* writer) {
 
 Status WriterMemoryManager::OnWriteCompleted(BatchWriter* writer) {
     UpdateWriterMemory(writer);
-    if (total_memory_ < max_memory_) {
+    if (total_memory_ < memory_limit_) {
         return Status::OK();
     }
 
@@ -68,13 +68,9 @@ void WriterMemoryManager::UpdateWriterMemory(BatchWriter* writer) {
     }
 }
 
-WriterMemoryManager::Candidate WriterMemoryManager::PickLargest(
-    const std::unordered_set<BatchWriter*>& skipped) const {
+WriterMemoryManager::Candidate WriterMemoryManager::PickLargest() const {
     Candidate candidate;
     for (const auto& [writer, memory] : writer_memory_) {
-        if (skipped.find(writer) != skipped.end()) {
-            continue;
-        }
         if (memory > candidate.memory) {
             candidate = {writer, memory};
         }
@@ -83,33 +79,29 @@ WriterMemoryManager::Candidate WriterMemoryManager::PickLargest(
 }
 
 Status WriterMemoryManager::ShrinkToLimit() {
-    std::unordered_set<BatchWriter*> no_progress_writers;
     while (true) {
-        if (total_memory_ < max_memory_) {
+        if (total_memory_ < memory_limit_) {
             return Status::OK();
         }
 
-        Candidate picked = PickLargest(no_progress_writers);
+        Candidate picked = PickLargest();
         if (picked.memory == 0) {
             return Status::Invalid(
                 fmt::format("Unable to release memory to below the write-buffer-size limit ({} "
                             "bytes), this might be a bug.",
-                            max_memory_));
+                            memory_limit_));
         }
-
         BatchWriter* candidate = picked.writer;
         uint64_t before_memory = picked.memory;
-
         PAIMON_RETURN_NOT_OK(candidate->FlushMemory());
 
         UpdateWriterMemory(candidate);
-
         uint64_t after_memory = candidate->GetMemoryUsage();
-
         if (after_memory >= before_memory) {
-            no_progress_writers.insert(candidate);
-        } else {
-            no_progress_writers.erase(candidate);
+            return Status::Invalid(fmt::format(
+                "Before flushing memory, writer had {} bytes of memory allocated, After flushing "
+                "memory, writer still has {} bytes of memory allocated, this might be a bug.",
+                before_memory, after_memory));
         }
     }
 }

…n/core/mergetree/writer_memory_manager.h …imon/core/memory/writer_memory_manager.hsrc/paimon/core/mergetree/writer_memory_manager.h renamed to src/paimon/core/memory/writer_memory_manager.h src/paimon/core/mergetree/writer_memory_manager.h renamed to src/paimon/core/memory/writer_memory_manager.h

@@ -18,24 +18,27 @@
 
 #include <cstdint>
 #include <unordered_map>
-#include <unordered_set>
 
 #include "paimon/status.h"
 
 namespace paimon {
 
 class BatchWriter;
 
-/// Coordinates global write-buffer memory across managed writers.
-/// NOTE: This class is not thread-safe.
+/// Coordinates global write-buffer memory across managed writers. Used in `AbstractFileStoreWrite`.
+/// @note This class is not thread-safe.
 class WriterMemoryManager {
  public:
-    explicit WriterMemoryManager(uint64_t max_memory) : max_memory_(max_memory) {}
+    explicit WriterMemoryManager(uint64_t memory_limit) : memory_limit_(memory_limit) {}
 
+    /// Register a writer when create a new `BatchWriter` in `AbstractFileStoreWrite::GetWriter()`
     void RegisterWriter(BatchWriter* writer);
+    /// Unregister a writer when the `BatchWriter` has been erased in `AbstractFileStoreWrite`
     void UnregisterWriter(BatchWriter* writer);
+    /// Refresh the memory usage of a writer when after `BatchWriter::PrepareCommit()`
     void RefreshWriterMemory(BatchWriter* writer);
-
+    /// Check if the total memory usage exceeds the limit after `BatchWriter::Write()`, and trigger
+    /// flush if needed.
     Status OnWriteCompleted(BatchWriter* writer);
 
  private:
@@ -45,10 +48,10 @@ class WriterMemoryManager {
     };
 
     void UpdateWriterMemory(BatchWriter* writer);
-    Candidate PickLargest(const std::unordered_set<BatchWriter*>& skipped) const;
+    Candidate PickLargest() const;
     Status ShrinkToLimit();
 
-    uint64_t max_memory_;
+    const uint64_t memory_limit_;
     uint64_t total_memory_ = 0;
     std::unordered_map<BatchWriter*, uint64_t> writer_memory_;
 };

…mergetree/writer_memory_manager_test.cpp …re/memory/writer_memory_manager_test.cppsrc/paimon/core/mergetree/writer_memory_manager_test.cpp renamed to src/paimon/core/memory/writer_memory_manager_test.cpp src/paimon/core/mergetree/writer_memory_manager_test.cpp renamed to src/paimon/core/memory/writer_memory_manager_test.cpp

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "paimon/core/mergetree/writer_memory_manager.h"
+#include "paimon/core/memory/writer_memory_manager.h"
 
 #include <algorithm>
 #include <memory>
@@ -35,8 +35,8 @@ namespace {
 
 class FakeBatchWriter : public BatchWriter {
  public:
-    explicit FakeBatchWriter(std::string name, std::vector<std::string>* flush_history = nullptr)
-        : name_(std::move(name)), flush_history_(flush_history) {}
+    FakeBatchWriter(const std::string& name, std::vector<std::string>* flush_history)
+        : name_(name), flush_history_(flush_history) {}
 
     void SetMemoryUsage(uint64_t memory_usage) {
         memory_usage_ = memory_usage;
@@ -52,7 +52,7 @@ class FakeBatchWriter : public BatchWriter {
 
     Status FlushMemory() override {
         uint64_t reduction = memory_usage_;
-        if (flush_calls_ < static_cast<int>(flush_reductions_.size())) {
+        if (flush_calls_ < static_cast<int32_t>(flush_reductions_.size())) {
             reduction = flush_reductions_[flush_calls_];
         }
         reduction = std::min(reduction, memory_usage_);
@@ -103,13 +103,13 @@ class FakeBatchWriter : public BatchWriter {
     std::vector<std::string>* flush_history_;
     std::vector<uint64_t> flush_reductions_;
     uint64_t memory_usage_ = 0;
-    int flush_calls_ = 0;
+    int32_t flush_calls_ = 0;
 };
 
 }  // namespace
 
 TEST(WriterMemoryManagerTest, DoesNotFlushWhenMemoryIsBelowLimit) {
-    WriterMemoryManager manager(/*max_memory=*/100);
+    WriterMemoryManager manager(/*memory_limit=*/100);
     std::vector<std::string> flush_history;
     FakeBatchWriter writer("writer", &flush_history);
     manager.RegisterWriter(&writer);
@@ -120,7 +120,7 @@ TEST(WriterMemoryManagerTest, DoesNotFlushWhenMemoryIsBelowLimit) {
 }
 
 TEST(WriterMemoryManagerTest, UnregisterWriterRemovesMemoryFromLedger) {
-    WriterMemoryManager manager(/*max_memory=*/100);
+    WriterMemoryManager manager(/*memory_limit=*/100);
     std::vector<std::string> flush_history;
     FakeBatchWriter writer_a("writer_a", &flush_history);
     manager.RegisterWriter(&writer_a);
@@ -137,7 +137,7 @@ TEST(WriterMemoryManagerTest, UnregisterWriterRemovesMemoryFromLedger) {
 }
 
 TEST(WriterMemoryManagerTest, RefreshWriterMemoryUpdatesLedgerWithoutFlushing) {
-    WriterMemoryManager manager(/*max_memory=*/80);
+    WriterMemoryManager manager(/*memory_limit=*/80);
     std::vector<std::string> flush_history;
     FakeBatchWriter writer_a("writer_a", &flush_history);
     manager.RegisterWriter(&writer_a);
@@ -155,7 +155,7 @@ TEST(WriterMemoryManagerTest, RefreshWriterMemoryUpdatesLedgerWithoutFlushing) {
 }
 
 TEST(WriterMemoryManagerTest, FlushWriterMemoryWithMultipleWriters) {
-    WriterMemoryManager manager(/*max_memory=*/100);
+    WriterMemoryManager manager(/*memory_limit=*/100);
     std::vector<std::string> flush_history;
     FakeBatchWriter writer_a("writer_a", &flush_history);
     manager.RegisterWriter(&writer_a);
@@ -182,7 +182,7 @@ TEST(WriterMemoryManagerTest, FlushWriterMemoryWithMultipleWriters) {
 }
 
 TEST(WriterMemoryManagerTest, ReclaimsCallerWhenCallerIsLargestWriter) {
-    WriterMemoryManager manager(/*max_memory=*/100);
+    WriterMemoryManager manager(/*memory_limit=*/100);
     std::vector<std::string> flush_history;
     FakeBatchWriter writer_a("writer_a", &flush_history);
     manager.RegisterWriter(&writer_a);
@@ -193,7 +193,6 @@ TEST(WriterMemoryManagerTest, ReclaimsCallerWhenCallerIsLargestWriter) {
     manager.RefreshWriterMemory(&writer_a);
 
     writer_b.SetMemoryUsage(90);
-    writer_b.SetFlushReductions({90});
     ASSERT_OK(manager.OnWriteCompleted(&writer_b));
 
     ASSERT_EQ(flush_history, std::vector<std::string>({"writer_b"}));
@@ -202,7 +201,7 @@ TEST(WriterMemoryManagerTest, ReclaimsCallerWhenCallerIsLargestWriter) {
 }
 
 TEST(WriterMemoryManagerTest, ContinuesReclaimingUntilBelowGlobalLimit) {
-    WriterMemoryManager manager(/*max_memory=*/61);
+    WriterMemoryManager manager(/*memory_limit=*/61);
     std::vector<std::string> flush_history;
     FakeBatchWriter writer_a("writer_a", &flush_history);
     manager.RegisterWriter(&writer_a);
@@ -224,22 +223,23 @@ TEST(WriterMemoryManagerTest, ContinuesReclaimingUntilBelowGlobalLimit) {
 }
 
 TEST(WriterMemoryManagerTest, ReturnsConfigurationErrorWhenNoWriterCanReleaseEnoughMemory) {
-    WriterMemoryManager manager(/*max_memory=*/100);
+    WriterMemoryManager manager(/*memory_limit=*/100);
     std::vector<std::string> flush_history;
     FakeBatchWriter writer_a("writer_a", &flush_history);
     manager.RegisterWriter(&writer_a);
     FakeBatchWriter writer_b("writer_b", &flush_history);
     manager.RegisterWriter(&writer_b);
 
     writer_b.SetMemoryUsage(20);
-    writer_b.SetFlushReductions({20});
     manager.RefreshWriterMemory(&writer_b);
 
     writer_a.SetMemoryUsage(120);
     writer_a.SetFlushReductions({10, 0});
-    ASSERT_NOK_WITH_MSG(manager.OnWriteCompleted(&writer_a),
-                        "Unable to release memory to below the write-buffer-size limit");
-    ASSERT_EQ(flush_history, std::vector<std::string>({"writer_a", "writer_b"}));
+    ASSERT_NOK_WITH_MSG(
+        manager.OnWriteCompleted(&writer_a),
+        "Before flushing memory, writer had 110 bytes of memory allocated, After flushing memory, "
+        "writer still has 110 bytes of memory allocated, this might be a bug.");
+    ASSERT_EQ(flush_history, std::vector<std::string>({"writer_a"}));
 }
 
 }  // namespace paimon::test

src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp

@@ -773,11 +773,13 @@ TEST_F(SortMergeReaderTest, TestRawSortNoMergeWithMinHeap) {
     // key: k0, user defined sequence field: ts, value: v0
     // Format: [_SEQUENCE_NUMBER, _VALUE_KIND, k0, ts, v0]
     // Reader1 (SEQUENCE_NUMBER 0..5):
-    //   batch1: [key=1,ts=1,v=1], [key=1,ts=2,v=2], [key=1,ts=3,v=3]
-    //   batch2: [key=1,ts=4,v=4], [key=2,ts=4,v=40], [key=2,ts=5,v=50]
+    //   [key=1,ts=1,v=1], [key=1,ts=2,v=2], [key=1,ts=3,v=3]
+    //   [key=1,ts=4,v=4], [key=2,ts=4,v=40], [key=2,ts=5,v=50]
     // Reader2 (SEQUENCE_NUMBER 6..11):
-    //   batch1: [key=1,ts=5,v=5], [key=1,ts=6,v=6], [key=2,ts=1,v=10]
-    //   batch2: [key=2,ts=2,v=20], [key=2,ts=3,v=30], [key=2,ts=6,v=60]
+    //   [key=1,ts=5,v=5], [key=1,ts=6,v=6], [key=2,ts=1,v=10]
+    //   [key=2,ts=2,v=20], [key=2,ts=3,v=30], [key=2,ts=6,v=60]
+    //
+    // After sort:
     // With user_defined_seq_comparator on ts field, sort by key asc, then ts asc within same key
     // key=1: ts=1(seq0,v=1), ts=2(seq1,v=2), ts=3(seq2,v=3), ts=4(seq3,v=4),
     //        ts=5(seq6,v=5), ts=6(seq7,v=6)

src/paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h

@@ -18,7 +18,6 @@
 
 #include <memory>
 #include <optional>
-#include <queue>
 #include <utility>
 #include <vector>
 
@@ -37,7 +36,7 @@ class KeyValueRecordReader;
 class Metrics;
 
 /// `SortMergeReader` implemented with loser tree. Merge the KeyValue parsed by
-/// KeyValueRecordReader and return the iterator of KeyValue
+/// `KeyValueRecordReader` and return the iterator of KeyValue
 class SortMergeReaderWithLoserTree : public SortMergeReader {
  public:
     SortMergeReaderWithLoserTree(

src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h

@@ -39,7 +39,7 @@ template <typename T>
 class MergeFunctionWrapper;
 
 /// `SortMergeReader` implemented with min-heap. Merge the KeyValue or only sort the KeyValue parsed
-/// by KeyValueRecordReader and return the iterator of KeyValue
+/// by `KeyValueRecordReader` and return the iterator of KeyValue
 class SortMergeReaderWithMinHeap : public SortMergeReader {
  public:
     SortMergeReaderWithMinHeap(

src/paimon/core/mergetree/external_sort_buffer.cpp

@@ -92,16 +92,6 @@ bool ExternalSortBuffer::HasSpilledData() const {
     return !spill_channel_manager_->GetChannels().empty();
 }
 
-std::vector<FileIOChannel::ID> ExternalSortBuffer::GetSpillChannelIdsSnapshot() const {
-    const auto& channels = spill_channel_manager_->GetChannels();
-    std::vector<FileIOChannel::ID> spill_channel_ids;
-    spill_channel_ids.reserve(channels.size());
-    for (const auto& spill_channel_id : channels) {
-        spill_channel_ids.push_back(spill_channel_id);
-    }
-    return spill_channel_ids;
-}
-
 void ExternalSortBuffer::Clear() {
     in_memory_buffer_->Clear();
     CleanupSpillFiles();
@@ -171,7 +161,7 @@ Result<int64_t> 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(
-        auto spill_writer,
+        std::unique_ptr<SpillWriter> spill_writer,
         SpillWriter::Create(options_.GetFileSystem(), write_schema_, spill_channel_enumerator_,
                             spill_channel_manager_, spill_compress_options.compress,
                             spill_compress_options.zstd_level));
@@ -229,8 +219,7 @@ Status ExternalSortBuffer::MergeSpilledFiles() {
     if (spill_channel_manager_->GetChannels().size() < 2) {
         return Status::OK();
     }
-
-    auto spill_channel_ids_before_merge = GetSpillChannelIdsSnapshot();
+    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);

src/paimon/core/mergetree/external_sort_buffer.h

@@ -53,6 +53,7 @@ class ExternalSortBuffer : public SortBuffer {
         const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
         const CoreOptions& options, const std::shared_ptr<IOManager>& io_manager,
         const std::shared_ptr<MemoryPool>& pool);
+    ~ExternalSortBuffer();
 
     void Clear() override;
     uint64_t GetMemorySize() const override;
@@ -79,7 +80,6 @@ class ExternalSortBuffer : public SortBuffer {
                        const CoreOptions& options,
                        const std::shared_ptr<FileIOChannel::Enumerator>& spill_channel_enumerator,
                        const std::shared_ptr<MemoryPool>& pool);
-    ~ExternalSortBuffer();
 
     std::unique_ptr<InMemorySortBuffer> in_memory_buffer_;