fix review

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

Author: zjw1111

src/paimon/core/mergetree/merge_tree_writer_test.cpp

@@ -46,8 +46,8 @@ template <typename T>
 class MergeFunctionWrapper;
 
 /// WriteBuffer manages the batch buffer for MergeTreeWriter.
-/// It delegates to a SortBuffer implementation (BinaryInMemorySortBuffer or
-/// BinaryExternalSortBuffer) based on the spillable configuration.
+/// It delegates to a SortBuffer implementation (InMemorySortBuffer or ExternalSortBuffer) based on
+/// the spillable configuration.
 class WriteBuffer {
  public:
     static Result<std::unique_ptr<WriteBuffer>> Create(

src/paimon/core/mergetree/write_buffer.h

@@ -24,7 +24,6 @@
 #include "fmt/format.h"
 #include "paimon/common/data/binary_row.h"
 #include "paimon/common/metrics/metrics_impl.h"
-#include "paimon/core/memory/writer_memory_manager.h"
 #include "paimon/core/operation/file_store_scan.h"
 #include "paimon/core/operation/file_system_write_restore.h"
 #include "paimon/core/operation/metrics/compaction_metrics.h"
@@ -83,7 +82,7 @@ AbstractFileStoreWrite::AbstractFileStoreWrite(
       metrics_(std::make_shared<MetricsImpl>()),
       logger_(Logger::GetLogger("AbstractFileStoreWrite")) {
     writer_memory_manager_ =
-        std::make_shared<WriterMemoryManager>(static_cast<uint64_t>(options.GetWriteBufferSize()));
+        std::make_unique<WriterMemoryManager>(static_cast<uint64_t>(options.GetWriteBufferSize()));
     cache_manager_ = std::make_shared<CacheManager>(options.GetLookupCacheMaxMemory(),
                                                     options.GetLookupCacheHighPrioPoolRatio());
 }

src/paimon/core/operation/abstract_file_store_write.cpp

@@ -29,6 +29,7 @@
 #include "paimon/common/io/cache/cache_manager.h"
 #include "paimon/core/core_options.h"
 #include "paimon/core/deletionvectors/bucketed_dv_maintainer.h"
+#include "paimon/core/memory/writer_memory_manager.h"
 #include "paimon/file_store_write.h"
 #include "paimon/logging.h"
 #include "paimon/metrics.h"
@@ -61,7 +62,6 @@ class MemoryPool;
 class RecordBatch;
 class RestoreFiles;
 class IOManager;
-class WriterMemoryManager;
 
 class AbstractFileStoreWrite : public FileStoreWrite {
  public:
@@ -132,7 +132,7 @@ class AbstractFileStoreWrite : public FileStoreWrite {
     std::shared_ptr<BucketedDvMaintainer::Factory> dv_maintainer_factory_;
     std::shared_ptr<IOManager> io_manager_;
     std::shared_ptr<CacheManager> cache_manager_;
-    std::shared_ptr<WriterMemoryManager> writer_memory_manager_;
+    std::unique_ptr<WriterMemoryManager> writer_memory_manager_;
 
     CoreOptions options_;
     std::shared_ptr<Executor> compact_executor_;

src/paimon/core/operation/abstract_file_store_write.h

@@ -117,7 +117,7 @@ Result<std::shared_ptr<BatchWriter>> KeyValueFileStoreWrite::CreateWriter(
             restore_max_seq_number, trimmed_primary_keys, data_file_path_factory, key_comparator_,
             user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_,
             options_, compact_manager, io_manager_, pool_));
-    return std::shared_ptr<BatchWriter>(writer);
+    return writer;
 }
 
 Status KeyValueFileStoreWrite::Close() {

src/paimon/core/operation/key_value_file_store_write.cpp

@@ -218,8 +218,7 @@ TEST_F(KeyValueFileStoreWriteTest,
     ASSERT_EQ(commit_messages.size(), 1);
 }
 
-TEST_F(KeyValueFileStoreWriteTest,
-       TestWriteShouldSpillLargestWriterWhenGlobalMemoryLimitIsExceeded) {
+TEST_F(KeyValueFileStoreWriteTest, TestSpillSimple) {
     auto fields = {arrow::field("f0", arrow::utf8(), /*nullable=*/false)};
     arrow::Schema typed_schema(fields);
     ::ArrowSchema schema;
@@ -241,17 +240,38 @@ TEST_F(KeyValueFileStoreWriteTest,
 
     ASSERT_OK_AND_ASSIGN(std::unique_ptr<WriteContext> write_context, context_builder.Finish());
     ASSERT_OK_AND_ASSIGN(auto file_store_write, FileStoreWrite::Create(std::move(write_context)));
+    auto key_value_file_store_write = dynamic_cast<KeyValueFileStoreWrite*>(file_store_write.get());
+    auto get_writer = [&](int32_t bucket) -> std::shared_ptr<paimon::BatchWriter> {
+        auto partition_iter = key_value_file_store_write->writers_.find(BinaryRow::EmptyRow());
+        if (partition_iter != key_value_file_store_write->writers_.end()) {
+            auto& buckets = partition_iter->second;
+            auto bucket_iter = buckets.find(bucket);
+            if (PAIMON_LIKELY(bucket_iter != buckets.end())) {
+                return bucket_iter->second.writer;
+            }
+        }
+        assert(false);
+        return nullptr;
+    };
 
+    // write bucket 0, not trigger spill
     ASSERT_OK(WriteSingleStringRow(file_store_write.get(), /*bucket=*/0, std::string(48, 'a')));
     ASSERT_EQ(TestHelper::CountChannelFiles(dir->GetFileSystem(), dir->Str()), 0);
+    ASSERT_GT(get_writer(0)->GetMemoryUsage(), 0);
 
-    ASSERT_OK(WriteSingleStringRow(file_store_write.get(), /*bucket=*/1, std::string(48, 'b')));
+    // write bucket 1, spill bucket 0 (pick largest writer)
+    ASSERT_OK(WriteSingleStringRow(file_store_write.get(), /*bucket=*/1, std::string(32, 'b')));
     ASSERT_GT(TestHelper::CountChannelFiles(dir->GetFileSystem(), dir->Str()), 0);
+    ASSERT_EQ(get_writer(0)->GetMemoryUsage(), 0);
+    ASSERT_GT(get_writer(1)->GetMemoryUsage(), 0);
 
+    // prepare commit, clean all spill files and memory buffers
     ASSERT_OK_AND_ASSIGN(auto commit_messages,
                          file_store_write->PrepareCommit(/*wait_compaction=*/true));
     ASSERT_EQ(commit_messages.size(), 2);
     ASSERT_EQ(TestHelper::CountChannelFiles(dir->GetFileSystem(), dir->Str()), 0);
+    ASSERT_EQ(get_writer(0)->GetMemoryUsage(), 0);
+    ASSERT_EQ(get_writer(1)->GetMemoryUsage(), 0);
 }
 
 }  // namespace paimon::test

src/paimon/core/operation/key_value_file_store_write_test.cpp

@@ -381,18 +381,21 @@ class TestHelper {
 
     static int64_t CountChannelFiles(const std::shared_ptr<FileSystem>& file_system,
                                      const std::string& tmp_path) {
-        std::vector<std::unique_ptr<BasicFileStatus>> file_statuses;
-        EXPECT_OK(file_system->ListDir(tmp_path, &file_statuses));
+        std::vector<std::unique_ptr<BasicFileStatus>> dir_statuses;
+        EXPECT_OK(file_system->ListDir(tmp_path, &dir_statuses));
 
         int64_t channel_file_count = 0;
-        for (const auto& file_status : file_statuses) {
-            const std::string path = file_status->GetPath();
-            if (file_status->IsDir() && path.find("paimon-io-") != std::string::npos) {
-                channel_file_count += CountChannelFiles(file_system, path);
-                continue;
-            }
-            if (StringUtils::EndsWith(path, ".channel")) {
-                ++channel_file_count;
+        for (const auto& dir_status : dir_statuses) {
+            const std::string dir_path = dir_status->GetPath();
+            if (dir_status->IsDir() && dir_path.find("paimon-io-") != std::string::npos) {
+                std::vector<std::unique_ptr<BasicFileStatus>> file_statuses;
+                EXPECT_OK(file_system->ListDir(dir_path, &file_statuses));
+
+                for (const auto& file_status : file_statuses) {
+                    if (StringUtils::EndsWith(file_status->GetPath(), ".channel")) {
+                        ++channel_file_count;
+                    }
+                }
             }
         }
         return channel_file_count;