Merge branch 'main' into feature/spill-to-disk

Author: zjw1111

include/paimon/write_context.h

@@ -154,6 +154,10 @@ class PAIMON_EXPORT WriteContextBuilder {
     /// @return Reference to this builder for method chaining.
     WriteContextBuilder& WithStreamingMode(bool is_streaming_mode);
 
+    /// Set whether to skip num-bucket consistency check (default is false)
+    /// @return Reference to this builder for method chaining.
+    WriteContextBuilder& WithIgnoreNumBucketCheck(bool ignore_num_bucket_check);
+
     /// Set whether the write operation should ignore previously stored files. (default is false)
     /// @return Reference to this builder for method chaining.
     WriteContextBuilder& WithIgnorePreviousFiles(bool ignore_previous_files);

src/paimon/core/operation/write_context.cpp

@@ -116,6 +116,11 @@ WriteContextBuilder& WriteContextBuilder::WithStreamingMode(bool is_streaming_mo
     return *this;
 }
 
+WriteContextBuilder& WriteContextBuilder::WithIgnoreNumBucketCheck(bool ignore_num_bucket_check) {
+    impl_->ignore_num_bucket_check_ = ignore_num_bucket_check;
+    return *this;
+}
+
 WriteContextBuilder& WriteContextBuilder::WithMemoryPool(
     const std::shared_ptr<MemoryPool>& memory_pool) {
     impl_->memory_pool_ = memory_pool;

src/paimon/core/operation/write_context_test.cpp

@@ -17,8 +17,11 @@
 #include "paimon/write_context.h"
 
 #include "gtest/gtest.h"
+#include "paimon/executor.h"
+#include "paimon/memory/memory_pool.h"
 #include "paimon/result.h"
 #include "paimon/status.h"
+#include "paimon/testing/mock/mock_file_system.h"
 #include "paimon/testing/utils/testharness.h"
 
 namespace paimon::test {
@@ -41,11 +44,41 @@ TEST(WriteContextTest, TestSimple) {
     ASSERT_TRUE(ctx->GetFileSystemSchemeToIdentifierMap().empty());
 }
 
-TEST(WriteContextTest, TestWithTempDirectory) {
+TEST(WriteContextTest, TestAllWithMethods) {
     WriteContextBuilder builder("table_root_path", "commit_user_1");
 
-    ASSERT_OK_AND_ASSIGN(auto ctx, builder.WithTempDirectory("/tmp").Finish());
-    ASSERT_EQ(ctx->GetTempDirectory(), "/tmp");
+    auto memory_pool = GetDefaultPool();
+    std::shared_ptr<Executor> executor = CreateDefaultExecutor();
+    auto file_system = std::make_shared<MockFileSystem>();
+    std::vector<std::string> write_schema = {"f0", "f1"};
+    std::map<std::string, std::string> fs_scheme_to_identifier_map = {{"file", "local"},
+                                                                      {"oss", "jindo"}};
+
+    ASSERT_OK_AND_ASSIGN(auto ctx,
+                         builder.WithStreamingMode(true)
+                             .WithIgnoreNumBucketCheck(true)
+                             .WithIgnorePreviousFiles(true)
+                             .WithMemoryPool(memory_pool)
+                             .WithExecutor(executor)
+                             .WithTempDirectory("/tmp/with-all")
+                             .WithWriteId(123)
+                             .WithBranch("test_branch")
+                             .WithWriteSchema(write_schema)
+                             .WithFileSystemSchemeToIdentifierMap(fs_scheme_to_identifier_map)
+                             .WithFileSystem(file_system)
+                             .Finish());
+
+    ASSERT_TRUE(ctx->IsStreamingMode());
+    ASSERT_TRUE(ctx->IgnoreNumBucketCheck());
+    ASSERT_TRUE(ctx->IgnorePreviousFiles());
+    ASSERT_EQ(ctx->GetMemoryPool(), memory_pool);
+    ASSERT_EQ(ctx->GetExecutor(), executor);
+    ASSERT_EQ(ctx->GetTempDirectory(), "/tmp/with-all");
+    ASSERT_EQ(ctx->GetWriteId(), 123);
+    ASSERT_EQ(ctx->GetBranch(), "test_branch");
+    ASSERT_EQ(ctx->GetWriteSchema(), write_schema);
+    ASSERT_EQ(ctx->GetFileSystemSchemeToIdentifierMap(), fs_scheme_to_identifier_map);
+    ASSERT_EQ(ctx->GetSpecificFileSystem(), file_system);
 }
 
 }  // namespace paimon::test