Skip to content

Commit f5841d1

Browse files
zjw1111claude
zjw1111
and
claude
committed
feat: add SetWriteBufferSpillThreadNumber to control spill thread usage
Add WriteContextBuilder::SetWriteBufferSpillThreadNumber(int32_t) to control Arrow IPC thread usage during spill. When > 0, sets Arrow CPU thread pool capacity and enables use_threads in SpillReader/SpillWriter. The bool is passed through the full constructor chain: KeyValueFileStoreWrite -> MergeTreeWriter -> WriteBuffer -> ExternalSortBuffer -> SpillReader/SpillWriter. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 244096f commit f5841d1

22 files changed

Lines changed: 129 additions & 76 deletions

‎include/paimon/write_context.h‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,8 @@ class PAIMON_EXPORT WriteContext {
4040
public:
4141
WriteContext(const std::string& root_path, const std::string& commit_user,
4242
bool is_streaming_mode, bool ignore_num_bucket_check, bool ignore_previous_files,
43-
const std::optional<int32_t>& write_id, const std::string& branch,
44-
const std::vector<std::string>& write_schema,
43+
bool enable_multi_thread_spill, const std::optional<int32_t>& write_id,
44+
const std::string& branch, const std::vector<std::string>& write_schema,
4545
const std::shared_ptr<MemoryPool>& memory_pool,
4646
const std::shared_ptr<Executor>& executor, const std::string& temp_directory,
4747
const std::shared_ptr<FileSystem>& specific_file_system,
@@ -106,13 +106,18 @@ class PAIMON_EXPORT WriteContext {
106106
return specific_file_system_;
107107
}
108108

109+
bool EnableMultiThreadSpill() const {
110+
return enable_multi_thread_spill_;
111+
}
112+
109113
private:
110114
std::string root_path_;
111115
std::string commit_user_;
112116
std::string branch_;
113117
bool is_streaming_mode_;
114118
bool ignore_num_bucket_check_;
115119
bool ignore_previous_files_;
120+
bool enable_multi_thread_spill_;
116121
std::optional<int32_t> write_id_;
117122
std::vector<std::string> write_schema_;
118123
std::shared_ptr<MemoryPool> memory_pool_;
@@ -222,6 +227,13 @@ class PAIMON_EXPORT WriteContextBuilder {
222227
WriteContextBuilder& WithFileSystemSchemeToIdentifierMap(
223228
const std::map<std::string, std::string>& fs_scheme_to_identifier_map);
224229

230+
/// Set the thread number for write buffer spill operations. (default is 0)
231+
/// If <= 0, threading is disabled for spill IPC read/write.
232+
/// If > 0, sets arrow CPU thread pool capacity for spill operations.
233+
/// @param thread_number The thread number to use for spill operations.
234+
/// @return Reference to this builder for method chaining.
235+
WriteContextBuilder& SetWriteBufferSpillThreadNumber(int32_t thread_number);
236+
225237
/// Build and return a `WriteContext` instance with input validation.
226238
/// @return Result containing the constructed `WriteContext` or an error status.
227239
Result<std::unique_ptr<WriteContext>> Finish();

‎src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,8 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam<std::
101101
/*last_sequence_number=*/last_sequence_number, std::vector<std::string>({"key"}),
102102
data_path_factory, key_comparator, /*user_defined_seq_comparator=*/nullptr,
103103
merge_function_wrapper, /*schema_id=*/latest_schema.value()->Id(), arrow_schema_,
104-
options, std::make_shared<NoopCompactManager>(), /*io_manager=*/nullptr, pool_));
104+
options, std::make_shared<NoopCompactManager>(), /*io_manager=*/nullptr,
105+
/*enable_multi_thread_spill=*/false, pool_));
105106

106107
// write data
107108
ArrowArray c_src_array;

‎src/paimon/core/mergetree/external_sort_buffer.cpp‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ Result<std::unique_ptr<ExternalSortBuffer>> ExternalSortBuffer::Create(
4747
const std::shared_ptr<FieldsComparator>& key_comparator,
4848
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
4949
const CoreOptions& options, const std::shared_ptr<IOManager>& io_manager,
50-
const std::shared_ptr<MemoryPool>& pool) {
50+
bool enable_multi_thread_spill, const std::shared_ptr<MemoryPool>& pool) {
5151
if (options.GetLocalSortMaxNumFileHandles() < kSpillMinFanIn) {
5252
return Status::Invalid(fmt::format(
5353
"invalid '{}': {}, must be at least {}", Options::LOCAL_SORT_MAX_NUM_FILE_HANDLES,
@@ -64,9 +64,10 @@ Result<std::unique_ptr<ExternalSortBuffer>> ExternalSortBuffer::Create(
6464

6565
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<FileIOChannel::Enumerator> spill_channel_enumerator,
6666
io_manager->CreateChannelEnumerator());
67-
return std::unique_ptr<ExternalSortBuffer>(new ExternalSortBuffer(
68-
std::move(in_memory_buffer), key_schema, value_schema, key_comparator,
69-
user_defined_seq_comparator, options, spill_channel_enumerator, pool));
67+
return std::unique_ptr<ExternalSortBuffer>(
68+
new ExternalSortBuffer(std::move(in_memory_buffer), key_schema, value_schema,
69+
key_comparator, user_defined_seq_comparator, options,
70+
spill_channel_enumerator, enable_multi_thread_spill, pool));
7071
}
7172

7273
ExternalSortBuffer::ExternalSortBuffer(
@@ -77,7 +78,7 @@ ExternalSortBuffer::ExternalSortBuffer(
7778
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
7879
const CoreOptions& options,
7980
const std::shared_ptr<FileIOChannel::Enumerator>& spill_channel_enumerator,
80-
const std::shared_ptr<MemoryPool>& pool)
81+
bool enable_multi_thread_spill, const std::shared_ptr<MemoryPool>& pool)
8182
: in_memory_buffer_(std::move(in_memory_buffer)),
8283
pool_(pool),
8384
key_schema_(key_schema),
@@ -87,6 +88,7 @@ ExternalSortBuffer::ExternalSortBuffer(
8788
write_schema_(SpecialFields::CompleteSequenceAndValueKindField(value_schema)),
8889
options_(options),
8990
max_fan_in_(options.GetLocalSortMaxNumFileHandles()),
91+
enable_multi_thread_spill_(enable_multi_thread_spill),
9092
spill_channel_manager_(
9193
std::make_shared<SpillChannelManager>(options_.GetFileSystem(), max_fan_in_)),
9294
spill_merger_(std::make_unique<SpillFileMerger>(max_fan_in_)),
@@ -189,9 +191,10 @@ Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> ExternalSortBuffer::C
189191
std::vector<std::unique_ptr<KeyValueRecordReader>> readers;
190192
readers.reserve(files.size());
191193
for (const auto& file : files) {
192-
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<SpillReader> reader,
193-
SpillReader::Create(options_.GetFileSystem(), key_schema_,
194-
value_schema_, pool_, file.channel_id));
194+
PAIMON_ASSIGN_OR_RAISE(
195+
std::unique_ptr<SpillReader> reader,
196+
SpillReader::Create(options_.GetFileSystem(), key_schema_, value_schema_,
197+
enable_multi_thread_spill_, file.channel_id, pool_));
195198
readers.push_back(std::move(reader));
196199
}
197200
return readers;
@@ -204,7 +207,7 @@ Result<FileChannelInfo> ExternalSortBuffer::SpillToDisk(
204207
std::unique_ptr<SpillWriter> spill_writer,
205208
SpillWriter::Create(options_.GetFileSystem(), write_schema_, spill_channel_enumerator_,
206209
spill_channel_manager_, spill_compress_options.compress,
207-
spill_compress_options.zstd_level, pool_));
210+
spill_compress_options.zstd_level, enable_multi_thread_spill_, pool_));
208211
auto cleanup_guard = ScopeGuard([&]() {
209212
[[maybe_unused]] auto status =
210213
spill_channel_manager_->DeleteChannel(spill_writer->GetChannelId());

‎src/paimon/core/mergetree/external_sort_buffer.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ class ExternalSortBuffer : public SortBuffer {
5353
const std::shared_ptr<FieldsComparator>& key_comparator,
5454
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
5555
const CoreOptions& options, const std::shared_ptr<IOManager>& io_manager,
56-
const std::shared_ptr<MemoryPool>& pool);
56+
bool enable_multi_thread_spill, const std::shared_ptr<MemoryPool>& pool);
5757
~ExternalSortBuffer() override;
5858

5959
void Clear() override;
@@ -85,6 +85,7 @@ class ExternalSortBuffer : public SortBuffer {
8585
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
8686
const CoreOptions& options,
8787
const std::shared_ptr<FileIOChannel::Enumerator>& spill_channel_enumerator,
88+
bool enable_multi_thread_spill,
8889
const std::shared_ptr<MemoryPool>& pool);
8990

9091
std::unique_ptr<InMemorySortBuffer> in_memory_buffer_;
@@ -97,6 +98,7 @@ class ExternalSortBuffer : public SortBuffer {
9798
const std::shared_ptr<arrow::Schema> write_schema_;
9899
const CoreOptions options_;
99100
const int32_t max_fan_in_;
101+
const bool enable_multi_thread_spill_;
100102
const std::shared_ptr<SpillChannelManager> spill_channel_manager_;
101103

102104
std::unique_ptr<SpillFileMerger> spill_merger_;

‎src/paimon/core/mergetree/lookup/remote_lookup_file_manager_test.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -77,12 +77,12 @@ class RemoteLookupFileManagerTest : public testing::Test {
7777
std::make_shared<ReducerMergeFunctionWrapper>(std::move(mfunc));
7878

7979
PAIMON_ASSIGN_OR_RAISE(
80-
auto writer,
81-
MergeTreeWriter::Create(/*last_sequence_number=*/last_sequence_number,
82-
std::vector<std::string>({"key"}), data_path_factory,
83-
key_comparator, /*user_defined_seq_comparator=*/nullptr,
84-
merge_function_wrapper, /*schema_id=*/0, arrow_schema_, options,
85-
noop_compact_manager_, /*io_manager=*/nullptr, pool_));
80+
auto writer, MergeTreeWriter::Create(
81+
/*last_sequence_number=*/last_sequence_number,
82+
std::vector<std::string>({"key"}), data_path_factory, key_comparator,
83+
/*user_defined_seq_comparator=*/nullptr, merge_function_wrapper,
84+
/*schema_id=*/0, arrow_schema_, options, noop_compact_manager_,
85+
/*io_manager=*/nullptr, /*enable_multi_thread_spill=*/false, pool_));
8686

8787
ArrowArray c_src_array;
8888
PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*src_array, &c_src_array));

‎src/paimon/core/mergetree/lookup_levels_test.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -81,12 +81,12 @@ class LookupLevelsTest : public testing::Test {
8181
std::make_shared<ReducerMergeFunctionWrapper>(std::move(mfunc));
8282

8383
PAIMON_ASSIGN_OR_RAISE(
84-
auto writer,
85-
MergeTreeWriter::Create(/*last_sequence_number=*/last_sequence_number,
86-
std::vector<std::string>({"key"}), data_path_factory,
87-
key_comparator, /*user_defined_seq_comparator=*/nullptr,
88-
merge_function_wrapper, /*schema_id=*/0, arrow_schema_, options,
89-
noop_compact_manager_, /*io_manager=*/nullptr, pool_));
84+
auto writer, MergeTreeWriter::Create(
85+
/*last_sequence_number=*/last_sequence_number,
86+
std::vector<std::string>({"key"}), data_path_factory, key_comparator,
87+
/*user_defined_seq_comparator=*/nullptr, merge_function_wrapper,
88+
/*schema_id=*/0, arrow_schema_, options, noop_compact_manager_,
89+
/*io_manager=*/nullptr, /*enable_multi_thread_spill=*/false, pool_));
9090

9191
// write data
9292
ArrowArray c_src_array;

‎src/paimon/core/mergetree/merge_tree_writer.cpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,13 +56,15 @@ Result<std::shared_ptr<MergeTreeWriter>> MergeTreeWriter::Create(
5656
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
5757
int64_t schema_id, const std::shared_ptr<arrow::Schema>& value_schema,
5858
const CoreOptions& options, const std::shared_ptr<CompactManager>& compact_manager,
59-
const std::shared_ptr<IOManager>& io_manager, const std::shared_ptr<MemoryPool>& pool) {
59+
const std::shared_ptr<IOManager>& io_manager, bool enable_multi_thread_spill,
60+
const std::shared_ptr<MemoryPool>& pool) {
6061
auto write_schema = SpecialFields::CompleteSequenceAndValueKindField(value_schema);
6162
PAIMON_ASSIGN_OR_RAISE(
6263
std::unique_ptr<WriteBuffer> write_buffer,
6364
WriteBuffer::Create(last_sequence_number, value_schema, trimmed_primary_keys,
6465
options.GetSequenceField(), key_comparator, user_defined_seq_comparator,
65-
merge_function_wrapper, options, io_manager, pool));
66+
merge_function_wrapper, options, io_manager, enable_multi_thread_spill,
67+
pool));
6668
return std::shared_ptr<MergeTreeWriter>(
6769
new MergeTreeWriter(pool, trimmed_primary_keys, options, path_factory, key_comparator,
6870
user_defined_seq_comparator, merge_function_wrapper, schema_id,

‎src/paimon/core/mergetree/merge_tree_writer.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,8 @@ class MergeTreeWriter : public BatchWriter {
6262
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
6363
int64_t schema_id, const std::shared_ptr<arrow::Schema>& value_schema,
6464
const CoreOptions& options, const std::shared_ptr<CompactManager>& compact_manager,
65-
const std::shared_ptr<IOManager>& io_manager, const std::shared_ptr<MemoryPool>& pool);
65+
const std::shared_ptr<IOManager>& io_manager, bool enable_multi_thread_spill,
66+
const std::shared_ptr<MemoryPool>& pool);
6667

6768
Status Write(std::unique_ptr<RecordBatch>&& batch) override;
6869

‎src/paimon/core/mergetree/merge_tree_writer_test.cpp‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -176,10 +176,10 @@ class MergeTreeWriterTest : public ::testing::TestWithParam<bool> {
176176
compact_manager ? compact_manager : noop_compact_manager_;
177177
std::shared_ptr<IOManager> io_manager =
178178
GetParam() ? std::make_shared<IOManager>(temp_dir + "/tmp", file_system_) : nullptr;
179-
return MergeTreeWriter::Create(last_sequence_number, primary_keys_, path_factory,
180-
key_comparator_, user_defined_seq_comparator,
181-
merge_function_wrapper_, schema_id, value_schema_, options,
182-
writer_compact_manager, io_manager, pool_);
179+
return MergeTreeWriter::Create(
180+
last_sequence_number, primary_keys_, path_factory, key_comparator_,
181+
user_defined_seq_comparator, merge_function_wrapper_, schema_id, value_schema_, options,
182+
writer_compact_manager, io_manager, /*enable_multi_thread_spill=*/false, pool_);
183183
}
184184

185185
private:
@@ -1094,7 +1094,8 @@ TEST_F(MergeTreeWriterTest, TestSpillWithSameKeyDeduplicate) {
10941094
MergeTreeWriter::Create(/*last_sequence_number=*/-1, primary_keys_, path_factory,
10951095
key_comparator_, /*user_defined_seq_comparator=*/nullptr,
10961096
merge_function_wrapper_, /*schema_id=*/0, value_schema_, options,
1097-
noop_compact_manager_, io_manager, pool_));
1097+
noop_compact_manager_, io_manager,
1098+
/*enable_multi_thread_spill=*/false, pool_));
10981099

10991100
std::shared_ptr<arrow::Array> batch1 =
11001101
arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([
@@ -1161,7 +1162,8 @@ TEST_F(MergeTreeWriterTest, TestIntermediateMergeSpillFileBound) {
11611162
MergeTreeWriter::Create(/*last_sequence_number=*/-1, primary_keys_, path_factory,
11621163
key_comparator_, /*user_defined_seq_comparator=*/nullptr,
11631164
merge_function_wrapper_, /*schema_id=*/0, value_schema_, options,
1164-
noop_compact_manager_, io_manager, pool_));
1165+
noop_compact_manager_, io_manager,
1166+
/*enable_multi_thread_spill=*/false, pool_));
11651167

11661168
std::shared_ptr<arrow::Array> batch1 =
11671169
arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([
@@ -1226,7 +1228,8 @@ TEST_F(MergeTreeWriterTest, TestDiskQuotaExhaustedFallsBackToFlushWriteBuffer) {
12261228
MergeTreeWriter::Create(/*last_sequence_number=*/-1, primary_keys_, path_factory,
12271229
key_comparator_, /*user_defined_seq_comparator=*/nullptr,
12281230
merge_function_wrapper_, /*schema_id=*/0, value_schema_, options,
1229-
noop_compact_manager_, io_manager, pool_));
1231+
noop_compact_manager_, io_manager,
1232+
/*enable_multi_thread_spill=*/false, pool_));
12301233

12311234
// Phase 1: Manual FlushMemory path — disk quota exhausted causes fallback.
12321235
std::shared_ptr<arrow::Array> array1 =
@@ -1303,7 +1306,8 @@ TEST_F(MergeTreeWriterTest, TestFlushMemoryQuotaExhaustedFallsBackToFlushWriteBu
13031306
MergeTreeWriter::Create(/*last_sequence_number=*/-1, primary_keys_, path_factory,
13041307
key_comparator_, /*user_defined_seq_comparator=*/nullptr,
13051308
merge_function_wrapper_, /*schema_id=*/0, value_schema_, options,
1306-
noop_compact_manager_, io_manager, pool_));
1309+
noop_compact_manager_, io_manager,
1310+
/*enable_multi_thread_spill=*/false, pool_));
13071311

13081312
std::shared_ptr<arrow::Array> array =
13091313
arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([
@@ -1346,7 +1350,8 @@ TEST_F(MergeTreeWriterTest, TestCloseDeletesSpillTempFiles) {
13461350
MergeTreeWriter::Create(/*last_sequence_number=*/-1, primary_keys_, path_factory,
13471351
key_comparator_, /*user_defined_seq_comparator=*/nullptr,
13481352
merge_function_wrapper_, /*schema_id=*/0, value_schema_, options,
1349-
noop_compact_manager_, io_manager, pool_));
1353+
noop_compact_manager_, io_manager,
1354+
/*enable_multi_thread_spill=*/false, pool_));
13501355

13511356
std::shared_ptr<arrow::Array> array =
13521357
arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([
@@ -1378,7 +1383,8 @@ TEST_F(MergeTreeWriterTest, TestMultiplePrepareCommitWithSpill) {
13781383
MergeTreeWriter::Create(/*last_sequence_number=*/-1, primary_keys_, path_factory,
13791384
key_comparator_, /*user_defined_seq_comparator=*/nullptr,
13801385
merge_function_wrapper_, /*schema_id=*/0, value_schema_, options,
1381-
noop_compact_manager_, io_manager, pool_));
1386+
noop_compact_manager_, io_manager,
1387+
/*enable_multi_thread_spill=*/false, pool_));
13821388

13831389
std::shared_ptr<arrow::Array> array1 =
13841390
arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([
@@ -1455,7 +1461,8 @@ TEST_F(MergeTreeWriterTest, TestSpillWithIOException) {
14551461
MergeTreeWriter::Create(/*last_sequence_number=*/-1, primary_keys_, path_factory,
14561462
key_comparator_, /*user_defined_seq_comparator=*/nullptr,
14571463
merge_function_wrapper_, /*schema_id=*/0, value_schema_,
1458-
options, noop_compact_manager_, io_manager, pool_));
1464+
options, noop_compact_manager_, io_manager,
1465+
/*enable_multi_thread_spill=*/false, pool_));
14591466

14601467
ScopeGuard guard([&io_hook]() { io_hook->Clear(); });
14611468
io_hook->Reset(i, IOHook::Mode::RETURN_ERROR);

‎src/paimon/core/mergetree/sort_buffer_test.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ class SortBufferTest : public ::testing::Test {
140140
/*sequence_fields_ascending=*/true, key_comparator_, write_buffer_size, pool_);
141141
return ExternalSortBuffer::Create(std::move(in_memory_buffer), value_schema_, primary_keys_,
142142
key_comparator_, sequence_comparator_, options,
143-
io_manager_, pool_);
143+
io_manager_, /*enable_multi_thread_spill=*/false, pool_);
144144
}
145145

146146
void AssertRows(const std::vector<ReaderResult>& actual,

0 commit comments

Comments
 (0)