Skip to content

Commit 16950cc

Browse files
zjw1111zjw1111
authored
feat(spill): support to create ExternalSortBuffer in WriteBuffer (#265)
1 parent d184ecf commit 16950cc

14 files changed

Lines changed: 385 additions & 184 deletions

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,10 @@
4444
#include "paimon/core/mergetree/lookup_levels.h"
4545
#include "paimon/core/mergetree/merge_tree_writer.h"
4646
#include "paimon/core/schema/schema_manager.h"
47-
#include "paimon/core/table/source/data_split_impl.h"
4847
#include "paimon/format/file_format_factory.h"
4948
#include "paimon/memory/memory_pool.h"
5049
#include "paimon/record_batch.h"
5150
#include "paimon/scan_context.h"
52-
#include "paimon/table/source/table_scan.h"
5351
#include "paimon/testing/mock/mock_index_path_factory.h"
5452
#include "paimon/testing/utils/binary_row_generator.h"
5553
#include "paimon/testing/utils/io_exception_helper.h"
@@ -96,12 +94,14 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam<std::
9694
if (!latest_schema) {
9795
return Status::Invalid("cannot find latest schema");
9896
}
99-
auto writer = std::make_shared<MergeTreeWriter>(
100-
/*last_sequence_number=*/last_sequence_number, std::vector<std::string>({"key"}),
101-
data_path_factory, key_comparator,
102-
/*user_defined_seq_comparator=*/nullptr, merge_function_wrapper,
103-
/*schema_id=*/latest_schema.value()->Id(), arrow_schema_, options,
104-
std::make_shared<NoopCompactManager>(), pool_);
97+
98+
PAIMON_ASSIGN_OR_RAISE(
99+
auto writer,
100+
MergeTreeWriter::Create(
101+
/*last_sequence_number=*/last_sequence_number, std::vector<std::string>({"key"}),
102+
data_path_factory, key_comparator, /*user_defined_seq_comparator=*/nullptr,
103+
merge_function_wrapper, /*schema_id=*/latest_schema.value()->Id(), arrow_schema_,
104+
options, std::make_shared<NoopCompactManager>(), /*io_manager=*/nullptr, pool_));
105105

106106
// write data
107107
ArrowArray c_src_array;

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

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@
3838
#include "paimon/core/schema/schema_manager.h"
3939
#include "paimon/memory/memory_pool.h"
4040
#include "paimon/record_batch.h"
41-
#include "paimon/testing/utils/binary_row_generator.h"
4241
#include "paimon/testing/utils/testharness.h"
4342

4443
namespace paimon::test {
@@ -77,11 +76,13 @@ class RemoteLookupFileManagerTest : public testing::Test {
7776
auto merge_function_wrapper =
7877
std::make_shared<ReducerMergeFunctionWrapper>(std::move(mfunc));
7978

80-
auto writer = std::make_shared<MergeTreeWriter>(
81-
/*last_sequence_number=*/last_sequence_number, std::vector<std::string>({"key"}),
82-
data_path_factory, key_comparator,
83-
/*user_defined_seq_comparator=*/nullptr, merge_function_wrapper, /*schema_id=*/0,
84-
arrow_schema_, options, noop_compact_manager_, pool_);
79+
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_));
8586

8687
ArrowArray c_src_array;
8788
PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*src_array, &c_src_array));

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include "paimon/common/utils/path_util.h"
3030
#include "paimon/core/compact/noop_compact_manager.h"
3131
#include "paimon/core/core_options.h"
32+
#include "paimon/core/disk/io_manager.h"
3233
#include "paimon/core/io/data_file_path_factory.h"
3334
#include "paimon/core/mergetree/compact/deduplicate_merge_function.h"
3435
#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h"
@@ -79,11 +80,13 @@ class LookupLevelsTest : public testing::Test {
7980
auto merge_function_wrapper =
8081
std::make_shared<ReducerMergeFunctionWrapper>(std::move(mfunc));
8182

82-
auto writer = std::make_shared<MergeTreeWriter>(
83-
/*last_sequence_number=*/last_sequence_number, std::vector<std::string>({"key"}),
84-
data_path_factory, key_comparator,
85-
/*user_defined_seq_comparator=*/nullptr, merge_function_wrapper, /*schema_id=*/0,
86-
arrow_schema_, options, noop_compact_manager_, pool_);
83+
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_));
8790

8891
// write data
8992
ArrowArray c_src_array;
@@ -429,7 +432,7 @@ TEST_F(LookupLevelsTest, TestLookupLevel0WithMultipleFiles) {
429432
ASSERT_FALSE(positioned_kv);
430433
}
431434

432-
TEST_F(LookupLevelsTest, TestWithPosistion) {
435+
TEST_F(LookupLevelsTest, TestWithPosition) {
433436
std::map<std::string, std::string> options = {};
434437
ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options));
435438
ASSERT_OK_AND_ASSIGN(auto table_path, CreateTable(options));

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

Lines changed: 46 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818

1919
#include <algorithm>
2020
#include <cassert>
21-
#include <cstddef>
2221
#include <unordered_set>
2322
#include <utility>
2423

@@ -29,6 +28,7 @@
2928
#include "paimon/common/table/special_fields.h"
3029
#include "paimon/common/utils/arrow/status_utils.h"
3130
#include "paimon/common/utils/scope_guard.h"
31+
#include "paimon/core/disk/io_manager.h"
3232
#include "paimon/core/io/async_key_value_producer_and_consumer.h"
3333
#include "paimon/core/io/compact_increment.h"
3434
#include "paimon/core/io/data_file_path_factory.h"
@@ -44,20 +44,40 @@
4444
#include "paimon/core/utils/commit_increment.h"
4545
#include "paimon/format/file_format.h"
4646
#include "paimon/format/writer_builder.h"
47-
#include "paimon/metrics.h"
4847

4948
namespace paimon {
5049
class FormatStatsExtractor;
5150

52-
MergeTreeWriter::MergeTreeWriter(
51+
Result<std::shared_ptr<MergeTreeWriter>> MergeTreeWriter::Create(
5352
int64_t last_sequence_number, const std::vector<std::string>& trimmed_primary_keys,
5453
const std::shared_ptr<DataFilePathFactory>& path_factory,
5554
const std::shared_ptr<FieldsComparator>& key_comparator,
5655
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
5756
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
5857
int64_t schema_id, const std::shared_ptr<arrow::Schema>& value_schema,
5958
const CoreOptions& options, const std::shared_ptr<CompactManager>& compact_manager,
60-
const std::shared_ptr<MemoryPool>& pool)
59+
const std::shared_ptr<IOManager>& io_manager, const std::shared_ptr<MemoryPool>& pool) {
60+
auto write_schema = SpecialFields::CompleteSequenceAndValueKindField(value_schema);
61+
PAIMON_ASSIGN_OR_RAISE(
62+
std::unique_ptr<WriteBuffer> write_buffer,
63+
WriteBuffer::Create(last_sequence_number, value_schema, trimmed_primary_keys,
64+
options.GetSequenceField(), key_comparator, user_defined_seq_comparator,
65+
merge_function_wrapper, options, io_manager, pool));
66+
return std::shared_ptr<MergeTreeWriter>(
67+
new MergeTreeWriter(pool, trimmed_primary_keys, options, path_factory, key_comparator,
68+
user_defined_seq_comparator, merge_function_wrapper, schema_id,
69+
write_schema, compact_manager, std::move(write_buffer)));
70+
}
71+
72+
MergeTreeWriter::MergeTreeWriter(
73+
const std::shared_ptr<MemoryPool>& pool, const std::vector<std::string>& trimmed_primary_keys,
74+
const CoreOptions& options, const std::shared_ptr<DataFilePathFactory>& path_factory,
75+
const std::shared_ptr<FieldsComparator>& key_comparator,
76+
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
77+
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
78+
int64_t schema_id, const std::shared_ptr<arrow::Schema>& write_schema,
79+
const std::shared_ptr<CompactManager>& compact_manager,
80+
std::unique_ptr<WriteBuffer>&& write_buffer)
6181
: pool_(pool),
6282
trimmed_primary_keys_(trimmed_primary_keys),
6383
options_(options),
@@ -66,13 +86,10 @@ MergeTreeWriter::MergeTreeWriter(
6686
user_defined_seq_comparator_(user_defined_seq_comparator),
6787
merge_function_wrapper_(merge_function_wrapper),
6888
schema_id_(schema_id),
89+
write_schema_(write_schema),
6990
compact_manager_(compact_manager),
70-
metrics_(std::make_shared<MetricsImpl>()) {
71-
write_schema_ = SpecialFields::CompleteSequenceAndValueKindField(value_schema);
72-
write_buffer_ = std::make_unique<WriteBuffer>(
73-
last_sequence_number, arrow::struct_(value_schema->fields()), trimmed_primary_keys_,
74-
options_.GetSequenceField(), key_comparator_, merge_function_wrapper_, pool_);
75-
}
91+
write_buffer_(std::move(write_buffer)),
92+
metrics_(std::make_shared<MetricsImpl>()) {}
7693

7794
Status MergeTreeWriter::DoClose() {
7895
// Request cancellation and wait for running compaction to exit.
@@ -115,16 +132,26 @@ Status MergeTreeWriter::DoClose() {
115132
return Status::OK();
116133
}
117134

135+
Status MergeTreeWriter::FlushMemory() {
136+
PAIMON_ASSIGN_OR_RAISE(bool has_remaining_quota, write_buffer_->FlushMemory());
137+
if (!has_remaining_quota) {
138+
PAIMON_RETURN_NOT_OK(FlushWriteBuffer(/*wait_for_latest_compaction=*/false,
139+
/*forced_full_compaction=*/false));
140+
}
141+
return Status::OK();
142+
}
143+
118144
Status MergeTreeWriter::Write(std::unique_ptr<RecordBatch>&& moved_batch) {
119-
PAIMON_RETURN_NOT_OK(write_buffer_->Write(std::move(moved_batch)));
120-
if (write_buffer_->GetMemoryUsage() >= static_cast<uint64_t>(options_.GetWriteBufferSize())) {
121-
return Flush(/*wait_for_latest_compaction=*/false, /*forced_full_compaction=*/false);
145+
PAIMON_ASSIGN_OR_RAISE(bool has_remaining_quota, write_buffer_->Write(std::move(moved_batch)));
146+
if (!has_remaining_quota) {
147+
return FlushWriteBuffer(/*wait_for_latest_compaction=*/false,
148+
/*forced_full_compaction=*/false);
122149
}
123150
return Status::OK();
124151
}
125152

126153
Status MergeTreeWriter::Compact(bool full_compaction) {
127-
return Flush(/*wait_for_latest_compaction=*/true, full_compaction);
154+
return FlushWriteBuffer(/*wait_for_latest_compaction=*/true, full_compaction);
128155
}
129156

130157
Status MergeTreeWriter::Sync() {
@@ -197,7 +224,7 @@ Status MergeTreeWriter::UpdateCompactDeletionFile(
197224
}
198225

199226
Result<CommitIncrement> MergeTreeWriter::PrepareCommit(bool wait_compaction) {
200-
PAIMON_RETURN_NOT_OK(Flush(wait_compaction, /*forced_full_compaction=*/false));
227+
PAIMON_RETURN_NOT_OK(FlushWriteBuffer(wait_compaction, /*forced_full_compaction=*/false));
201228
if (options_.CommitForceCompact()) {
202229
wait_compaction = true;
203230
}
@@ -218,13 +245,14 @@ Result<bool> MergeTreeWriter::CompactNotCompleted() {
218245
return compact_manager_->CompactNotCompleted();
219246
}
220247

221-
Status MergeTreeWriter::Flush(bool wait_for_latest_compaction, bool forced_full_compaction) {
248+
Status MergeTreeWriter::FlushWriteBuffer(bool wait_for_latest_compaction,
249+
bool forced_full_compaction) {
222250
if (!write_buffer_->IsEmpty()) {
223251
if (compact_manager_->ShouldWaitForLatestCompaction()) {
224252
wait_for_latest_compaction = true;
225253
}
226254
auto cleanup_guard = ScopeGuard([&]() { write_buffer_->Clear(); });
227-
// 1. flush write buffer to get in-memory readers
255+
// 1. flush write buffer to get sorted readers
228256
PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> readers,
229257
write_buffer_->CreateReaders());
230258
// 2. prepare loser tree sort merge reader
@@ -257,6 +285,7 @@ Status MergeTreeWriter::Flush(bool wait_for_latest_compaction, bool forced_full_
257285
PAIMON_RETURN_NOT_OK(rolling_writer->Close());
258286
PAIMON_ASSIGN_OR_RAISE(std::vector<std::shared_ptr<DataFileMeta>> flushed_files,
259287
rolling_writer->GetResult());
288+
async_key_value_producer_consumer->Close();
260289
write_guard.Release();
261290

262291
for (const auto& flushed_file : flushed_files) {

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

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@
3232
#include "paimon/core/mergetree/write_buffer.h"
3333
#include "paimon/core/utils/batch_writer.h"
3434
#include "paimon/core/utils/commit_increment.h"
35-
#include "paimon/core/utils/path_factory.h"
3635
#include "paimon/record_batch.h"
3736
#include "paimon/result.h"
3837
#include "paimon/status.h"
@@ -46,6 +45,7 @@ class StructArray;
4645

4746
namespace paimon {
4847
class DataFilePathFactory;
48+
class IOManager;
4949
class FieldsComparator;
5050
class MemoryPool;
5151
class Metrics;
@@ -54,16 +54,15 @@ class MergeFunctionWrapper;
5454

5555
class MergeTreeWriter : public BatchWriter {
5656
public:
57-
MergeTreeWriter(int64_t last_sequence_number,
58-
const std::vector<std::string>& trimmed_primary_keys,
59-
const std::shared_ptr<DataFilePathFactory>& path_factory,
60-
const std::shared_ptr<FieldsComparator>& key_comparator,
61-
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
62-
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
63-
int64_t schema_id, const std::shared_ptr<arrow::Schema>& value_schema,
64-
const CoreOptions& options,
65-
const std::shared_ptr<CompactManager>& compact_manager,
66-
const std::shared_ptr<MemoryPool>& pool);
57+
static Result<std::shared_ptr<MergeTreeWriter>> Create(
58+
int64_t last_sequence_number, const std::vector<std::string>& trimmed_primary_keys,
59+
const std::shared_ptr<DataFilePathFactory>& path_factory,
60+
const std::shared_ptr<FieldsComparator>& key_comparator,
61+
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
62+
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
63+
int64_t schema_id, const std::shared_ptr<arrow::Schema>& value_schema,
64+
const CoreOptions& options, const std::shared_ptr<CompactManager>& compact_manager,
65+
const std::shared_ptr<IOManager>& io_manager, const std::shared_ptr<MemoryPool>& pool);
6766

6867
Status Write(std::unique_ptr<RecordBatch>&& batch) override;
6968

@@ -76,12 +75,10 @@ class MergeTreeWriter : public BatchWriter {
7675
Result<CommitIncrement> PrepareCommit(bool wait_compaction) override;
7776

7877
uint64_t GetMemoryUsage() const override {
79-
return 0;
78+
return write_buffer_->GetMemoryUsage();
8079
}
8180

82-
Status FlushMemory() override {
83-
return Flush(/*wait_for_latest_compaction=*/false, /*forced_full_compaction=*/false);
84-
}
81+
Status FlushMemory() override;
8582

8683
Status Close() override {
8784
return DoClose();
@@ -94,7 +91,7 @@ class MergeTreeWriter : public BatchWriter {
9491
private:
9592
Status DoClose();
9693

97-
Status Flush(bool wait_for_latest_compaction, bool forced_full_compaction);
94+
Status FlushWriteBuffer(bool wait_for_latest_compaction, bool forced_full_compaction);
9895
Result<CommitIncrement> DrainIncrement();
9996

10097
std::unique_ptr<RollingFileWriter<KeyValueBatch, std::shared_ptr<DataFileMeta>>>
@@ -105,6 +102,17 @@ class MergeTreeWriter : public BatchWriter {
105102
Status UpdateCompactDeletionFile(const std::shared_ptr<CompactDeletionFile>& new_deletion_file);
106103

107104
private:
105+
MergeTreeWriter(const std::shared_ptr<MemoryPool>& pool,
106+
const std::vector<std::string>& trimmed_primary_keys,
107+
const CoreOptions& options,
108+
const std::shared_ptr<DataFilePathFactory>& path_factory,
109+
const std::shared_ptr<FieldsComparator>& key_comparator,
110+
const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
111+
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
112+
int64_t schema_id, const std::shared_ptr<arrow::Schema>& write_schema,
113+
const std::shared_ptr<CompactManager>& compact_manager,
114+
std::unique_ptr<WriteBuffer>&& write_buffer);
115+
108116
std::shared_ptr<MemoryPool> pool_;
109117
std::vector<std::string> trimmed_primary_keys_;
110118
CoreOptions options_;

0 commit comments

Comments
 (0)