From ffc8d034b90e068850d6157e56bf9264301c8ad7 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:26:34 +0800 Subject: [PATCH] refactor(mergetree): write sorted key-value readers --- .../core/mergetree/merge_tree_writer.cpp | 95 +++++--- src/paimon/core/mergetree/merge_tree_writer.h | 5 + .../core/mergetree/merge_tree_writer_test.cpp | 216 ++++++++++++++++++ 3 files changed, 282 insertions(+), 34 deletions(-) diff --git a/src/paimon/core/mergetree/merge_tree_writer.cpp b/src/paimon/core/mergetree/merge_tree_writer.cpp index bb7a01054..748740f88 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer.cpp @@ -161,6 +161,65 @@ Status MergeTreeWriter::Write(std::unique_ptr&& moved_batch) { return Status::OK(); } +Status MergeTreeWriter::WriteSortedReadersToFiles( + std::vector>&& readers) { + auto raw_readers_guard = ScopeGuard([&]() -> void { + for (std::unique_ptr& reader : readers) { + if (reader != nullptr) { + reader->Close(); + } + } + }); + if (readers.empty()) { + return Status::Invalid("sorted readers must not be empty"); + } + for (const std::unique_ptr& reader : readers) { + if (reader == nullptr) { + return Status::Invalid("sorted readers must not contain null reader"); + } + } + + // prepare loser tree sort merge reader + auto sort_merge_reader = std::make_unique( + std::move(readers), key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_); + raw_readers_guard.Release(); + // project key value to arrow array + auto create_consumer = [target_schema = write_schema_, pool = pool_]() + -> Result>> { + return KeyValueMetaProjectionConsumer::Create(target_schema, pool); + }; + // consumer batch size is WriteBatchSize + std::unique_ptr producer = + std::make_unique(std::move(sort_merge_reader), + options_.GetWriteBatchSize()); + auto async_key_value_producer_consumer = + std::make_unique>( + std::move(producer), create_consumer, /*projection_thread_num=*/1); + ScopeGuard async_readers_guard([&]() -> void { async_key_value_producer_consumer->Close(); }); + std::unique_ptr>> rolling_writer; + PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); + ScopeGuard abort_writer_guard([&]() -> void { rolling_writer->Abort(); }); + while (true) { + PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, + async_key_value_producer_consumer->NextBatch()); + if (key_value_batch.batch == nullptr) { + break; + } + PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); + } + PAIMON_RETURN_NOT_OK(rolling_writer->Close()); + PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, + rolling_writer->GetResult()); + abort_writer_guard.Release(); + + for (const std::shared_ptr& flushed_file : flushed_files) { + new_files_.emplace_back(flushed_file); + PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); + } + metrics_->Merge(rolling_writer->GetMetrics()); + return Status::OK(); +} + Status MergeTreeWriter::Compact(bool full_compaction) { return FlushWriteBuffer(/*wait_for_latest_compaction=*/true, full_compaction); } @@ -309,52 +368,20 @@ Status MergeTreeWriter::FlushWriteBuffer(bool wait_for_latest_compaction, PAIMON_ASSIGN_OR_RAISE(flushed_changelog_files, changelog_writer->GetResult()); } + // 1. flush write buffer to get sorted readers // Flush write buffer to get sorted and merged data readers. PAIMON_ASSIGN_OR_RAISE(std::vector> readers, write_buffer_->CreateReaders()); - auto sort_merge_reader = std::make_unique( - std::move(readers), key_comparator_, user_defined_seq_comparator_, - merge_function_wrapper_); - std::unique_ptr producer = - std::make_unique(std::move(sort_merge_reader), - options_.GetWriteBatchSize()); - auto async_key_value_producer_consumer = - std::make_unique>( - std::move(producer), create_consumer, /*projection_thread_num=*/1); - std::unique_ptr>> - rolling_writer; - PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); - ScopeGuard write_guard([&]() -> void { - rolling_writer->Abort(); - async_key_value_producer_consumer->Close(); - }); - while (true) { - PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, - async_key_value_producer_consumer->NextBatch()); - if (key_value_batch.batch == nullptr) { - break; - } - PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); - } - PAIMON_RETURN_NOT_OK(rolling_writer->Close()); - PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, - rolling_writer->GetResult()); - async_key_value_producer_consumer->Close(); + PAIMON_RETURN_NOT_OK(WriteSortedReadersToFiles(std::move(readers))); if (async_changelog_producer_consumer) { async_changelog_producer_consumer->Close(); } new_changelog_files_.insert(new_changelog_files_.end(), flushed_changelog_files.begin(), flushed_changelog_files.end()); - new_files_.insert(new_files_.end(), flushed_files.begin(), flushed_files.end()); - write_guard.Release(); changelog_write_guard.Release(); - for (const auto& flushed_file : flushed_files) { - PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); - } - metrics_->Merge(rolling_writer->GetMetrics()); if (changelog_writer) { metrics_->Merge(changelog_writer->GetMetrics()); } diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index 10f35fbbe..575bd76f6 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.h +++ b/src/paimon/core/mergetree/merge_tree_writer.h @@ -51,6 +51,7 @@ class IOManager; class FieldsComparator; class MemoryPool; class Metrics; +class KeyValueRecordReader; template class MergeFunctionWrapper; @@ -69,6 +70,10 @@ class MergeTreeWriter : public BatchWriter { Status Write(std::unique_ptr&& batch) override; + /// Consumes readers whose complete streams are individually sorted by primary key and sequence + /// number. Readers are closed on success or failure. + Status WriteSortedReadersToFiles(std::vector>&& readers); + Status Compact(bool full_compaction) override; Result CompactNotCompleted() override; diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index b3f53f487..c1c804f9f 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include "arrow/api.h" #include "arrow/array/array_base.h" @@ -43,6 +44,7 @@ #include "paimon/core/io/compact_increment.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/data_increment.h" +#include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" @@ -53,6 +55,8 @@ #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/mock/mock_key_value_data_file_record_reader.h" #include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/io_exception_helper.h" #include "paimon/testing/utils/read_result_collector.h" @@ -65,6 +69,36 @@ class MergeFunctionWrapper; } // namespace paimon namespace paimon::test { +namespace { + +class TrackingKeyValueRecordReader : public KeyValueRecordReader { + public: + TrackingKeyValueRecordReader(std::unique_ptr&& inner_reader, + bool* closed_flag) + : inner_reader_(std::move(inner_reader)), closed_flag_(closed_flag) {} + + Result> NextBatch() override { + return inner_reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return inner_reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_flag_ != nullptr) { + *closed_flag_ = true; + } + inner_reader_->Close(); + } + + private: + std::unique_ptr inner_reader_; + bool* closed_flag_; +}; + +} // namespace + class MergeTreeWriterTest : public ::testing::TestWithParam { public: class FakeCompactManager : public paimon::CompactManager { @@ -217,6 +251,23 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { writer_compact_manager, io_manager, /*enable_multi_thread_spill=*/false, pool_); } + std::unique_ptr CreateSingleReader( + const std::shared_ptr& array, int32_t batch_size = 16, + const Status& next_batch_status = Status::OK()) const { + std::vector write_fields = {SpecialFields::SequenceNumber(), + SpecialFields::ValueKind()}; + write_fields.insert(write_fields.end(), value_fields_.begin(), value_fields_.end()); + std::shared_ptr write_schema = + DataField::ConvertDataFieldsToArrowSchema(write_fields); + std::shared_ptr key_schema = + arrow::schema(arrow::FieldVector({write_schema->field(2)})); + auto file_batch_reader = + std::make_unique(array, array->type(), batch_size); + file_batch_reader->SetNextBatchStatus(next_batch_status); + return std::make_unique( + std::move(file_batch_reader), key_schema, value_schema_, 0, pool_); + } + private: std::shared_ptr pool_; std::shared_ptr file_system_; @@ -558,6 +609,171 @@ TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); } +TEST_P(MergeTreeWriterTest, TestSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + std::string uuid = path_factory->uuid_; + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(4, dir->Str(), path_factory, 7, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])") + .ValueOrDie()); + + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + + ASSERT_OK(merge_writer->WriteSortedReadersToFiles(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + std::string expected_data_file_name = "data-" + uuid + "-0.orc"; + std::string expected_data_file_path = dir->Str() + "/" + expected_data_file_name; + ASSERT_OK_AND_ASSIGN(FileStatus data_file_status, + options.GetFileSystem()->GetFileStatus(expected_data_file_path)); + + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])"}, + &expected_array) + .ok()); + CheckFileContent(expected_data_file_path, expected_array); + + ASSERT_TRUE(commit_increment.GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(expected_data_file_name, new_file->file_name); + ASSERT_EQ(data_file_status.GetLen(), new_file->file_size); + ASSERT_EQ(3, new_file->row_count); + ASSERT_EQ(7, new_file->min_sequence_number); + ASSERT_EQ(9, new_file->max_sequence_number); + ASSERT_EQ(7, new_file->schema_id); + ASSERT_EQ(1, new_file->delete_row_count); +} + +TEST_P(MergeTreeWriterTest, TestMergeSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(/*last_sequence_number=*/4, dir->Str(), path_factory, + /*schema_id=*/7, options)); + + auto first_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [5, 0, "Alice", 10, 0, 15.1], + [7, 0, "Carol", 20, 1, 17.1], + [10, 0, "Eve", 30, 2, 20.1] + ])") + .ValueOrDie()); + auto second_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1] + ])") + .ValueOrDie()); + bool first_closed = false; + bool second_closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(first_array), &first_closed)); + sorted_readers.push_back(std::make_unique( + CreateSingleReader(second_array), &second_closed)); + + ASSERT_OK(merge_writer->WriteSortedReadersToFiles(std::move(sorted_readers))); + ASSERT_TRUE(first_closed); + ASSERT_TRUE(second_closed); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(5, new_file->row_count); + ASSERT_EQ(1, new_file->delete_row_count); + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [5, 0, "Alice", 10, 0, 15.1], + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1], + [10, 0, "Eve", 30, 2, 20.1] + ])"}, + &expected_array) + .ok()); + CheckFileContent(path_factory->ToPath(new_file), expected_array); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + std::vector> empty_readers; + Status empty_status = merge_writer->WriteSortedReadersToFiles(std::move(empty_readers)); + ASSERT_TRUE(empty_status.IsInvalid()); + + std::vector> null_readers; + null_readers.push_back(nullptr); + Status null_status = merge_writer->WriteSortedReadersToFiles(std::move(null_readers)); + ASSERT_TRUE(null_status.IsInvalid()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + bool first_mixed_reader_closed = false; + bool second_mixed_reader_closed = false; + std::vector> mixed_readers; + mixed_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &first_mixed_reader_closed)); + mixed_readers.push_back(nullptr); + mixed_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &second_mixed_reader_closed)); + Status mixed_status = merge_writer->WriteSortedReadersToFiles(std::move(mixed_readers)); + ASSERT_TRUE(mixed_status.IsInvalid()); + ASSERT_TRUE(first_mixed_reader_closed); + ASSERT_TRUE(second_mixed_reader_closed); + + Status expected_status = Status::IOError("sorted reader failure"); + bool failing_reader_closed = false; + std::vector> failing_readers; + failing_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array, /*batch_size=*/16, expected_status), + &failing_reader_closed)); + Status failing_status = merge_writer->WriteSortedReadersToFiles(std::move(failing_readers)); + ASSERT_EQ(expected_status, failing_status); + ASSERT_TRUE(failing_reader_closed); + ASSERT_OK(merge_writer->Close()); +} + TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({