This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-cpp.git
The following commit(s) were added to refs/heads/main by this push:
new 4d1c320 feat: Migrate sort merge reader (#85)
4d1c320 is described below
commit 4d1c320093ff32408269e4fda11c1b3edae1c64c
Author: lxy <[email protected]>
AuthorDate: Thu Jun 18 10:16:49 2026 +0800
feat: Migrate sort merge reader (#85)
---
src/paimon/core/mergetree/compact/loser_tree.cpp | 196 +++++
src/paimon/core/mergetree/compact/loser_tree.h | 171 ++++
src/paimon/core/mergetree/compact/merge_function.h | 30 +
.../mergetree/compact/merge_function_wrapper.h | 32 +
.../core/mergetree/compact/sort_merge_reader.h | 41 +
.../mergetree/compact/sort_merge_reader_test.cpp | 903 +++++++++++++++++++++
.../compact/sort_merge_reader_with_loser_tree.cpp | 88 ++
.../compact/sort_merge_reader_with_loser_tree.h | 97 +++
.../compact/sort_merge_reader_with_min_heap.cpp | 137 ++++
.../compact/sort_merge_reader_with_min_heap.h | 160 ++++
src/paimon/core/mergetree/drop_delete_reader.h | 85 ++
.../core/mergetree/drop_delete_reader_test.cpp | 111 +++
12 files changed, 2051 insertions(+)
diff --git a/src/paimon/core/mergetree/compact/loser_tree.cpp
b/src/paimon/core/mergetree/compact/loser_tree.cpp
new file mode 100644
index 0000000..e932119
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/loser_tree.cpp
@@ -0,0 +1,196 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/mergetree/compact/loser_tree.h"
+
+#include <algorithm>
+#include <cassert>
+
+namespace paimon {
+LoserTree::LoserTree(std::vector<std::unique_ptr<KeyValueRecordReader>>&&
readers,
+ const CompareFunc& first_comparator, const CompareFunc&
second_comparator)
+ : size_(readers.size()),
+ initialized_(false),
+ readers_holder_(std::move(readers)),
+ tree_(size_),
+ first_comparator_(first_comparator),
+ second_comparator_(second_comparator) {
+ leaves_.reserve(size_);
+ for (const auto& reader : readers_holder_) {
+ leaves_.emplace_back(reader.get());
+ }
+}
+
+Status LoserTree::InitializeIfNeeded() {
+ if (!initialized_) {
+ std::fill(tree_.begin(), tree_.end(), -1);
+ for (int32_t i = size_ - 1; i >= 0; i--) {
+ PAIMON_RETURN_NOT_OK(leaves_[i].AdvanceIfAvailable());
+ Adjust(i);
+ }
+ initialized_ = true;
+ }
+ return Status::OK();
+}
+
+Status LoserTree::AdjustForNextLoop() {
+ LeafIterator* winner = &leaves_[tree_[0]];
+ while (winner->state == State::WINNER_POPPED) {
+ PAIMON_RETURN_NOT_OK(winner->AdvanceIfAvailable());
+ Adjust(tree_[0]);
+ winner = &leaves_[tree_[0]];
+ }
+ return Status::OK();
+}
+
+std::optional<KeyValue> LoserTree::PopWinner() {
+ LeafIterator* winner = &leaves_[tree_[0]];
+ if (winner->state == State::WINNER_POPPED) {
+ // if the winner has already been popped, it means that all the same
key has been
+ // processed.
+ return std::nullopt;
+ }
+ std::optional<KeyValue> result = std::move(winner->Pop());
+ Adjust(tree_[0]);
+ return result;
+}
+
+const std::optional<KeyValue>& LoserTree::PeekWinner() const {
+ static const std::optional<KeyValue> empty_kv = std::nullopt;
+ return leaves_[tree_[0]].state != State::WINNER_POPPED ?
leaves_[tree_[0]].Peek() : empty_kv;
+}
+
+void LoserTree::Adjust(int32_t winner) {
+ for (int32_t parent = (winner + size_) / 2; parent > 0 && winner >= 0;
parent /= 2) {
+ LeafIterator* winner_node = &leaves_[winner];
+ LeafIterator* parent_node = nullptr;
+
+ if (tree_[parent] == -1) {
+ // initialize the tree.
+ winner_node->state = State::LOSER_WITH_NEW_KEY;
+ } else {
+ parent_node = &leaves_[tree_[parent]];
+ switch (winner_node->state) {
+ case State::WINNER_WITH_NEW_KEY: {
+ AdjustWithNewWinnerKey(parent, parent_node, winner_node);
+ break;
+ }
+ case State::WINNER_WITH_SAME_KEY: {
+ AdjustWithSameWinnerKey(parent, parent_node, winner_node);
+ break;
+ }
+ case State::WINNER_POPPED: {
+ if (winner_node->first_same_key_index < 0) {
+ // fast path, which means that the same key is not yet
processed in the
+ // current tree.
+ parent = -1;
+ } else {
+ // fast path. Directly exchange positions with the
same key that has not
+ // yet been processed, no need to compare level by
level.
+ parent = winner_node->first_same_key_index;
+ parent_node = &leaves_[tree_[parent]];
+ winner_node->state = State::LOSER_POPPED;
+ parent_node->state = State::WINNER_WITH_SAME_KEY;
+ }
+ break;
+ }
+ default:
+ assert(false);
+ }
+ }
+
+ // if the winner loses, exchange nodes.
+ if (!IsWinner(winner_node->state)) {
+ std::swap(winner, tree_[parent]);
+ }
+ }
+ tree_[0] = winner;
+}
+
+void LoserTree::AdjustWithSameWinnerKey(int32_t index, LeafIterator*
parent_node,
+ LeafIterator* winner_node) {
+ switch (parent_node->state) {
+ case State::LOSER_WITH_SAME_KEY: {
+ // the key of the previous loser is the same as the key of the
current winner,
+ // only the sequence needs to be compared.
+ const auto& parent_key = parent_node->Peek();
+ const auto& child_key = winner_node->Peek();
+ int32_t second_result = second_comparator_(parent_key, child_key);
+ if (second_result > 0) {
+ parent_node->state = State::WINNER_WITH_SAME_KEY;
+ winner_node->state = State::LOSER_WITH_SAME_KEY;
+ parent_node->SetFirstSameKeyIndex(index);
+ } else {
+ winner_node->SetFirstSameKeyIndex(index);
+ }
+ return;
+ }
+ case State::LOSER_WITH_NEW_KEY:
+ case State::LOSER_POPPED:
+ return;
+ default:
+ assert(false);
+ }
+}
+
+void LoserTree::AdjustWithNewWinnerKey(int32_t index, LeafIterator*
parent_node,
+ LeafIterator* winner_node) {
+ switch (parent_node->state) {
+ case State::LOSER_WITH_NEW_KEY: {
+ // when the new winner is also a new key, it needs to be compared.
+ const auto& parent_key = parent_node->Peek();
+ const auto& child_key = winner_node->Peek();
+ int32_t first_result = first_comparator_(parent_key, child_key);
+ if (first_result == 0) {
+ // if the compared keys are the same, we need to update the
state of the node
+ // and record the index of the same key for the winner.
+ int32_t second_result = second_comparator_(parent_key,
child_key);
+ if (second_result < 0) {
+ parent_node->state = State::LOSER_WITH_SAME_KEY;
+ winner_node->SetFirstSameKeyIndex(index);
+ } else {
+ winner_node->state = State::LOSER_WITH_SAME_KEY;
+ parent_node->state = State::WINNER_WITH_NEW_KEY;
+ parent_node->SetFirstSameKeyIndex(index);
+ }
+ } else if (first_result > 0) {
+ // the two keys are completely different and just need to
update the state.
+ parent_node->state = State::WINNER_WITH_NEW_KEY;
+ winner_node->state = State::LOSER_WITH_NEW_KEY;
+ }
+ return;
+ }
+ case State::LOSER_WITH_SAME_KEY: {
+ // A node in the WINNER_WITH_NEW_KEY state cannot encounter a node
in the
+ // LOSER_WITH_SAME_KEY state.
+ assert(false);
+ break;
+ }
+ case State::LOSER_POPPED: {
+ // this case will only happen during adjustForNextLoop.
+ parent_node->state = State::WINNER_POPPED;
+ parent_node->first_same_key_index = -1;
+ winner_node->state = State::LOSER_WITH_NEW_KEY;
+ return;
+ }
+ default:
+ assert(false);
+ }
+}
+
+} // namespace paimon
diff --git a/src/paimon/core/mergetree/compact/loser_tree.h
b/src/paimon/core/mergetree/compact/loser_tree.h
new file mode 100644
index 0000000..1417d8d
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/loser_tree.h
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "paimon/common/metrics/metrics_impl.h"
+#include "paimon/core/io/key_value_record_reader.h"
+#include "paimon/core/key_value.h"
+#include "paimon/result.h"
+#include "paimon/status.h"
+
+namespace paimon {
+class Metrics;
+
+class LoserTree {
+ public:
+ using CompareFunc =
+ std::function<int32_t(const std::optional<KeyValue>&, const
std::optional<KeyValue>&)>;
+ LoserTree(std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers,
+ const CompareFunc& first_comparator, const CompareFunc&
second_comparator);
+
+ /// Initialize the loser tree in the same way as the regular loser tree.
+ Status InitializeIfNeeded();
+
+ /// Adjust the Key that needs to be returned in the next round.
+ Status AdjustForNextLoop();
+
+ /// Pop the current winner and update its state to `State#WINNER_POPPED`.
+ std::optional<KeyValue> PopWinner();
+
+ /// Peek the current winner, mainly for key comparisons.
+ const std::optional<KeyValue>& PeekWinner() const;
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const {
+ return MetricsImpl::CollectReadMetrics(readers_holder_);
+ }
+
+ void Close() {
+ for (const auto& reader : readers_holder_) {
+ reader->Close();
+ }
+ }
+
+ private:
+ struct LeafIterator;
+
+ /// Adjust the winner from bottom to top. Using different `State`, we can
quickly compare
+ /// whether all the current same keys have been processed.
+ void Adjust(int32_t winner);
+
+ /// The winner node has the same userKey as the global winner.
+ void AdjustWithSameWinnerKey(int32_t index, LeafIterator* parent_node,
+ LeafIterator* winner_node);
+
+ /// The userKey of the new local winner node is different from that of the
previous global
+ /// winner.
+ void AdjustWithNewWinnerKey(int32_t index, LeafIterator* parent_node,
+ LeafIterator* winner_node);
+
+ private:
+ enum class State {
+ LOSER_WITH_NEW_KEY = 1,
+ LOSER_WITH_SAME_KEY = 2,
+ LOSER_POPPED = 3,
+ WINNER_WITH_NEW_KEY = 4,
+ WINNER_WITH_SAME_KEY = 5,
+ WINNER_POPPED = 6
+ };
+
+ static bool IsWinner(State state) {
+ if (state == State::LOSER_WITH_NEW_KEY || state ==
State::LOSER_WITH_SAME_KEY ||
+ state == State::LOSER_POPPED) {
+ return false;
+ }
+ return true;
+ }
+
+ struct LeafIterator {
+ explicit LeafIterator(KeyValueRecordReader* reader) : reader(reader) {}
+
+ const std::optional<KeyValue>& Peek() const {
+ return kv;
+ }
+
+ std::optional<KeyValue>&& Pop() {
+ state = State::WINNER_POPPED;
+ return std::move(kv);
+ }
+
+ void SetFirstSameKeyIndex(int32_t index) {
+ if (first_same_key_index == -1) {
+ first_same_key_index = index;
+ }
+ }
+
+ /// Reads the next kv if any, otherwise returns null.
+ Status AdvanceIfAvailable() {
+ first_same_key_index = -1;
+ state = State::WINNER_WITH_NEW_KEY;
+ bool has_next = false;
+ if (iterator != nullptr) {
+ PAIMON_ASSIGN_OR_RAISE(has_next, iterator->HasNext());
+ }
+ if (iterator == nullptr || !has_next) {
+ while (!end_of_input) {
+ PAIMON_ASSIGN_OR_RAISE(iterator, reader->NextBatch());
+ if (!iterator) {
+ // read eof
+ reader->Close();
+ end_of_input = true;
+ kv = std::nullopt;
+ } else {
+ PAIMON_ASSIGN_OR_RAISE(has_next, iterator->HasNext());
+ if (!has_next) {
+ continue;
+ }
+ PAIMON_ASSIGN_OR_RAISE(kv, iterator->Next());
+ break;
+ }
+ }
+ } else {
+ PAIMON_ASSIGN_OR_RAISE(kv, iterator->Next());
+ }
+ return Status::OK();
+ }
+
+ bool end_of_input = false;
+ int32_t first_same_key_index = -1;
+ State state = State::WINNER_WITH_NEW_KEY;
+ KeyValueRecordReader* reader;
+ std::unique_ptr<KeyValueRecordReader::Iterator> iterator;
+ std::optional<KeyValue> kv;
+ };
+
+ private:
+ int32_t size_;
+ bool initialized_;
+ // must hold all readers, as data array is allocated by the pool of data
file
+ // reader
+ std::vector<std::unique_ptr<KeyValueRecordReader>> readers_holder_;
+
+ std::vector<int32_t> tree_;
+ std::vector<LeafIterator> leaves_;
+ /// if comparator.compare('a', 'b') > 0, then 'a' is the winner. In the
following
+ /// implementation, we always let 'a' represent the parent node.
+ CompareFunc first_comparator_;
+ /// same as first_comparator, but mainly used to compare sequenceNumber.
+ CompareFunc second_comparator_;
+};
+} // namespace paimon
diff --git a/src/paimon/core/mergetree/compact/merge_function.h
b/src/paimon/core/mergetree/compact/merge_function.h
new file mode 100644
index 0000000..851e4df
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/merge_function.h
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include "paimon/core/key_value.h"
+namespace paimon {
+/// Merge function to merge multiple `KeyValue`s.
+class MergeFunction {
+ public:
+ virtual ~MergeFunction() = default;
+ virtual void Reset() = 0;
+ virtual Status Add(KeyValue&& kv) = 0;
+ virtual Result<std::optional<KeyValue>> GetResult() = 0;
+};
+} // namespace paimon
diff --git a/src/paimon/core/mergetree/compact/merge_function_wrapper.h
b/src/paimon/core/mergetree/compact/merge_function_wrapper.h
new file mode 100644
index 0000000..215e9f0
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/merge_function_wrapper.h
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include "paimon/core/key_value.h"
+
+namespace paimon {
+/// A wrapper for `MergeFunction`, which adds new functionalities or
optimizations.
+template <typename T>
+class MergeFunctionWrapper {
+ public:
+ virtual ~MergeFunctionWrapper() = default;
+ virtual void Reset() = 0;
+ virtual Status Add(KeyValue&& kv) = 0;
+ virtual Result<std::optional<T>> GetResult() = 0;
+};
+} // namespace paimon
diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader.h
b/src/paimon/core/mergetree/compact/sort_merge_reader.h
new file mode 100644
index 0000000..c7754e3
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/sort_merge_reader.h
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include <memory>
+
+#include "paimon/core/key_value.h"
+#include "paimon/metrics.h"
+#include "paimon/result.h"
+
+namespace paimon {
+class SortMergeReader {
+ public:
+ virtual ~SortMergeReader() = default;
+ class Iterator {
+ public:
+ virtual ~Iterator() = default;
+ virtual Result<bool> HasNext() = 0;
+ virtual KeyValue&& Next() = 0;
+ };
+
+ virtual Result<std::unique_ptr<Iterator>> NextBatch() = 0;
+ virtual void Close() = 0;
+ virtual std::shared_ptr<Metrics> GetReaderMetrics() const = 0;
+};
+} // namespace paimon
diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp
b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp
new file mode 100644
index 0000000..3c217b5
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp
@@ -0,0 +1,903 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/mergetree/compact/sort_merge_reader.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "arrow/api.h"
+#include "arrow/array/array_base.h"
+#include "arrow/array/array_nested.h"
+#include "arrow/ipc/json_simple.h"
+#include "gtest/gtest.h"
+#include "paimon/common/types/data_field.h"
+#include "paimon/common/utils/fields_comparator.h"
+#include "paimon/core/core_options.h"
+#include "paimon/core/io/concat_key_value_record_reader.h"
+#include "paimon/core/io/key_value_record_reader.h"
+#include "paimon/core/io/merged_key_value_record_reader.h"
+#include "paimon/core/key_value.h"
+#include "paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h"
+#include "paimon/core/mergetree/compact/deduplicate_merge_function.h"
+#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h"
+#include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h"
+#include "paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h"
+#include "paimon/memory/memory_pool.h"
+#include "paimon/result.h"
+#include "paimon/status.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/key_value_checker.h"
+#include "paimon/testing/utils/read_result_collector.h"
+#include "paimon/testing/utils/testharness.h"
+
+namespace paimon::test {
+class SortMergeReaderTest : public testing::Test {
+ public:
+ void SetUp() override {
+ pool_ = GetDefaultPool();
+ }
+
+ std::vector<DataField> CreateDataField(const arrow::FieldVector&
arrow_fields) {
+ // create DataField with fake field id
+ std::vector<DataField> data_fields;
+ data_fields.reserve(arrow_fields.size());
+ for (const auto& field : arrow_fields) {
+ data_fields.emplace_back(/*id=*/0, field);
+ }
+ return data_fields;
+ }
+
+ void CheckResult(const std::vector<std::shared_ptr<arrow::StructArray>>&
src_array_vec,
+ const std::shared_ptr<FieldsComparator>&
user_key_comparator,
+ const std::shared_ptr<FieldsComparator>&
user_defined_seq_comparator,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::shared_ptr<arrow::Schema>& value_schema,
+ const std::vector<KeyValue>& expected, bool ignore_delete
= false) const {
+ CheckSortMergeResult<SortMergeReaderWithLoserTree>(src_array_vec,
user_key_comparator,
+
user_defined_seq_comparator, key_schema,
+ value_schema,
expected,
+
/*need_merge=*/true, ignore_delete);
+ CheckSortMergeResult<SortMergeReaderWithMinHeap>(src_array_vec,
user_key_comparator,
+
user_defined_seq_comparator, key_schema,
+ value_schema,
expected,
+ /*need_merge=*/true,
ignore_delete);
+ }
+
+ private:
+ template <typename SortMergeReaderType>
+ std::unique_ptr<SortMergeReader> CreateSortMergeReader(
+ const std::vector<std::shared_ptr<arrow::StructArray>>& src_array_vec,
+ const std::shared_ptr<FieldsComparator>& user_key_comparator,
+ const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::shared_ptr<arrow::Schema>& value_schema, int32_t
batch_size, bool need_merge,
+ bool ignore_delete) const {
+ auto mfunc = std::make_unique<DeduplicateMergeFunction>(ignore_delete);
+ auto merge_function_wrapper =
+ std::make_shared<ReducerMergeFunctionWrapper>(std::move(mfunc));
+ if (!need_merge) {
+ if constexpr (std::is_same_v<SortMergeReaderType,
SortMergeReaderWithMinHeap>) {
+ merge_function_wrapper = nullptr;
+ } else {
+ ADD_FAILURE() << "Only SortMergeReaderWithMinHeap supports no
merge";
+ }
+ }
+ std::vector<std::unique_ptr<KeyValueRecordReader>> concat_readers;
+ for (const auto& src_array : src_array_vec) {
+ auto file_batch_reader = std::make_unique<MockFileBatchReader>(
+ src_array, src_array->type(), /*batch_size=*/batch_size);
+ auto record_reader =
std::make_unique<MockKeyValueDataFileRecordReader>(
+ std::move(file_batch_reader), key_schema, value_schema,
/*level=*/0, pool_);
+ std::vector<std::unique_ptr<KeyValueRecordReader>> readers;
+ readers.push_back(std::move(record_reader));
+ concat_readers.push_back(
+
std::make_unique<ConcatKeyValueRecordReader>(std::move(readers)));
+ }
+
+ return
std::make_unique<SortMergeReaderType>(std::move(concat_readers),
user_key_comparator,
+
user_defined_seq_comparator,
+ merge_function_wrapper);
+ }
+
+ template <typename SortMergeReaderType>
+ void CheckSortMergeResult(const
std::vector<std::shared_ptr<arrow::StructArray>>& src_array_vec,
+ const std::shared_ptr<FieldsComparator>&
user_key_comparator,
+ const std::shared_ptr<FieldsComparator>&
user_defined_seq_comparator,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::shared_ptr<arrow::Schema>&
value_schema,
+ const std::vector<KeyValue>& expected, bool
need_merge,
+ bool ignore_delete = false) const {
+ for (auto batch_size : {1, 2, 3, 4, 100}) {
+ auto sort_merge_reader =
CreateSortMergeReader<SortMergeReaderType>(
+ src_array_vec, user_key_comparator,
user_defined_seq_comparator, key_schema,
+ value_schema, batch_size, need_merge, ignore_delete);
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<KeyValue> results,
+ (ReadResultCollector::CollectKeyValueResult<
+ SortMergeReader,
SortMergeReader::Iterator>(sort_merge_reader.get())));
+ KeyValueChecker::CheckResult(expected, results,
key_schema->num_fields(),
+ value_schema->num_fields());
+ }
+ }
+
+ template <typename SortMergeReaderType>
+ void CheckSortMergeResultForAggregate(
+ const std::vector<std::shared_ptr<arrow::StructArray>>& src_array_vec,
+ const std::shared_ptr<FieldsComparator>& user_key_comparator,
+ const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::shared_ptr<arrow::Schema>& value_schema,
+ const std::vector<std::string>& user_defined_sequence_fields,
+ const std::vector<std::string>& primary_keys, const CoreOptions&
core_options,
+ const std::vector<KeyValue>& expected) const {
+ for (auto batch_size : {1, 2, 3, 4, 100}) {
+ ASSERT_OK_AND_ASSIGN(
+ std::unique_ptr<AggregateMergeFunction> mfunc,
+ AggregateMergeFunction::Create(value_schema, primary_keys,
core_options));
+ auto merge_function_wrapper =
+
std::make_shared<ReducerMergeFunctionWrapper>(std::move(mfunc));
+ std::vector<std::unique_ptr<KeyValueRecordReader>> merged_readers;
+
+ std::vector<std::unique_ptr<KeyValueRecordReader>> readers;
+ for (const auto& src_array : src_array_vec) {
+ auto file_batch_reader = std::make_unique<MockFileBatchReader>(
+ src_array, src_array->type(), /*batch_size=*/batch_size);
+ auto record_reader =
std::make_unique<MockKeyValueDataFileRecordReader>(
+ std::move(file_batch_reader), key_schema, value_schema,
/*level=*/0, pool_);
+
merged_readers.push_back(std::make_unique<MergedKeyValueRecordReader>(
+ std::move(record_reader), user_key_comparator,
merge_function_wrapper));
+ }
+
+ auto sort_merge_reader = std::make_unique<SortMergeReaderType>(
+ std::move(merged_readers), user_key_comparator,
user_defined_seq_comparator,
+ merge_function_wrapper);
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<KeyValue> results,
+ (ReadResultCollector::CollectKeyValueResult<
+ SortMergeReader,
SortMergeReader::Iterator>(sort_merge_reader.get())));
+ KeyValueChecker::CheckResult(expected, results,
key_schema->num_fields(),
+ value_schema->num_fields());
+ }
+ }
+
+ private:
+ std::shared_ptr<MemoryPool> pool_;
+};
+
+TEST_F(SortMergeReaderTest, TestSimpleWithTwoSameKeys) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 2, 10]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 3, 30]
+ ])")
+ .ValueOrDie());
+
+ auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 5, 50]
+ ])")
+ .ValueOrDie());
+
+ auto src_array4 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 2, 30]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {1, 1, 1}, {{2}, {3}, {5}}, {{2, 30}, {3, 30}, {5, 50}}, pool_);
+ CheckSortMergeResult<SortMergeReaderWithLoserTree>(
+ {src_array1, src_array2, src_array3, src_array4}, user_key_comparator,
nullptr, key_schema,
+ value_schema, expected, /*need_merge=*/true);
+}
+
+TEST_F(SortMergeReaderTest, TestSimpleWithThreeSameKeys) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 2, 10]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [2, 0, 2, 30]
+ ])")
+ .ValueOrDie());
+
+ auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 5, 50]
+ ])")
+ .ValueOrDie());
+
+ auto src_array4 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 2, 30]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected =
+ KeyValueChecker::GenerateKeyValues({2, 1}, {{2}, {5}}, {{2, 30}, {5,
50}}, pool_);
+ CheckSortMergeResult<SortMergeReaderWithLoserTree>(
+ {src_array1, src_array2, src_array3, src_array4}, user_key_comparator,
nullptr, key_schema,
+ value_schema, expected, /*need_merge=*/true);
+}
+
+TEST_F(SortMergeReaderTest, TestSimpleWithThreeSameKeys2) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 2, 10]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 2, 30]
+ ])")
+ .ValueOrDie());
+
+ auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 5, 50]
+ ])")
+ .ValueOrDie());
+
+ auto src_array4 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [2, 0, 2, 30]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected =
+ KeyValueChecker::GenerateKeyValues({2, 1}, {{2}, {5}}, {{2, 30}, {5,
50}}, pool_);
+ CheckSortMergeResult<SortMergeReaderWithLoserTree>(
+ {src_array1, src_array2, src_array3, src_array4}, user_key_comparator,
nullptr, key_schema,
+ value_schema, expected, /*need_merge=*/true);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeIn2Ways) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("k1", arrow::int32()),
+ arrow::field("v0", arrow::int32()),
+ arrow::field("v1", arrow::int32()),
+ arrow::field("v2", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3], fields[4],
fields[5], fields[6]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 1, 10, 20, 30],
+ [0, 0, 1, 3, 11, 21, 31],
+ [1, 0, 2, 2, 12, 22, 32],
+ [2, 0, 2, 3, 13, 23, 33]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 1, 1, 14, 24, 34],
+ [0, 0, 1, 2, 15, 25, 35],
+ [2, 0, 2, 2, 16, 26, 36],
+ [2, 0, 2, 5, 17, 27, 37]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(
+ std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2], data_fields[3]},
std::vector<int32_t>({0, 1}),
+ /*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {1, 0, 0, 2, 2, 2}, {{1, 1}, {1, 2}, {1, 3}, {2, 2}, {2, 3}, {2, 5}},
+ {{1, 1, 14, 24, 34},
+ {1, 2, 15, 25, 35},
+ {1, 3, 11, 21, 31},
+ {2, 2, 16, 26, 36},
+ {2, 3, 13, 23, 33},
+ {2, 5, 17, 27, 37}},
+ pool_);
+ CheckResult({src_array1, src_array2}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema, expected);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeIn3Ways) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 10],
+ [0, 0, 2, 11],
+ [1, 0, 3, 12],
+ [2, 0, 4, 13]
+ ])")
+ .ValueOrDie());
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 2, 14],
+ [2, 0, 3, 15],
+ [3, 0, 4, 16],
+ [2, 0, 5, 17]
+ ])")
+ .ValueOrDie());
+ auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 1, 17],
+ [2, 0, 2, 18],
+ [4, 0, 4, 19],
+ [4, 0, 5, 20]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected =
+ KeyValueChecker::GenerateKeyValues({1, 2, 2, 4, 4}, {{1}, {2}, {3},
{4}, {5}},
+ {{1, 17}, {2, 18}, {3, 15}, {4,
19}, {5, 20}}, pool_);
+
+ CheckResult({src_array1, src_array2, src_array3}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema, expected);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeWithDeleteMessages) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 3, 1, 10],
+ [2, 3, 2, 200],
+ [1, 0, 3, 300]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [3, 3, 1, 11],
+ [4, 3, 2, 240],
+ [5, 3, 3, 30]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+
+ std::vector<RowKind*> delete_row_kinds =
{const_cast<RowKind*>(RowKind::Delete()),
+
const_cast<RowKind*>(RowKind::Delete()),
+
const_cast<RowKind*>(RowKind::Delete())};
+ std::vector<KeyValue> expected_delete = KeyValueChecker::GenerateKeyValues(
+ delete_row_kinds, /*seq_vec=*/{3, 4, 5}, /*level_vec=*/{0, 0, 0},
+ /*key_vec=*/{{1}, {2}, {3}}, /*value_vec=*/{{1, 11}, {2, 240}, {3,
30}}, pool_);
+ CheckResult({src_array1, src_array2}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema, expected_delete,
+ /*ignore_delete=*/false);
+
+ std::vector<KeyValue> expected_ignore_delete =
KeyValueChecker::GenerateKeyValues(
+ /*seq_vec=*/{1}, /*key_vec=*/{{3}}, /*value_vec=*/{{3, 300}}, pool_);
+ CheckResult({src_array1, src_array2}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema,
+ expected_ignore_delete, /*ignore_delete=*/true);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeIn2WaysWithEmptyArray) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+ auto data_fields = CreateDataField(fields);
+
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 10],
+ [0, 0, 2, 11],
+ [1, 0, 3, 12],
+ [2, 0, 4, 13]
+ ])")
+ .ValueOrDie());
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {0, 0, 1, 2}, {{1}, {2}, {3}, {4}}, {{1, 10}, {2, 11}, {3, 12}, {4,
13}}, pool_);
+
+ CheckResult({src_array1, src_array2}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema, expected);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeIn2WaysWithNoOverlap) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 10],
+ [0, 0, 2, 11],
+ [1, 0, 3, 12],
+ [2, 0, 4, 13]
+ ])")
+ .ValueOrDie());
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 102, 14],
+ [2, 0, 103, 15],
+ [3, 0, 104, 16],
+ [2, 0, 105, 17]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {0, 0, 1, 2, 1, 2, 3, 2}, {{1}, {2}, {3}, {4}, {102}, {103}, {104},
{105}},
+ {{1, 10}, {2, 11}, {3, 12}, {4, 13}, {102, 14}, {103, 15}, {104, 16},
{105, 17}}, pool_);
+
+ CheckResult({src_array1, src_array2}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema, expected);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeIn2WaysWithFullOverlap) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 10],
+ [0, 0, 2, 11],
+ [1, 0, 3, 12],
+ [2, 0, 4, 13]
+ ])")
+ .ValueOrDie());
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 1, 14],
+ [2, 0, 2, 15],
+ [3, 0, 3, 16],
+ [3, 0, 4, 17]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {1, 2, 3, 3}, {{1}, {2}, {3}, {4}}, {{1, 14}, {2, 15}, {3, 16}, {4,
17}}, pool_);
+
+ CheckResult({src_array1, src_array2}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema, expected);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeIn2WaysWithPartialOverlap) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 10],
+ [0, 0, 2, 11],
+ [1, 0, 3, 12],
+ [2, 0, 4, 13]
+ ])")
+ .ValueOrDie());
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 1, 14],
+ [2, 0, 2, 15],
+ [3, 0, 3, 16],
+ [3, 0, 4, 17],
+ [0, 0, 5, 18],
+ [0, 0, 6, 19]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {1, 2, 3, 3, 0, 0}, {{1}, {2}, {3}, {4}, {5}, {6}},
+ {{1, 14}, {2, 15}, {3, 16}, {4, 17}, {5, 18}, {6, 19}}, pool_);
+
+ CheckResult({src_array1, src_array2}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema,
value_schema, expected);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeIn3WaysWithUserDefinedSeq) {
+ // key: k0, k1
+ // user defined sequence field: v0, v1
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("k1", arrow::int32()),
+ arrow::field("v0", arrow::int32()),
+ arrow::field("v1", arrow::int32()),
+ arrow::field("v2", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3], fields[4],
fields[5], fields[6]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [3, 0, 1, 1, 10, 20, 30],
+ [0, 0, 1, 3, 11, 21, 31],
+ [3, 0, 2, 2, 12, 22, 32],
+ [2, 0, 2, 3, 13, 23, 33]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 1, 1, 14, 24, 34],
+ [1, 0, 1, 2, 15, 25, 35],
+ [2, 0, 2, 2, 18, 28, 38],
+ [2, 0, 2, 5, 17, 27, 37]
+ ])")
+ .ValueOrDie());
+
+ auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [2, 0, 1, 1, 14, 24, 34],
+ [0, 0, 1, 2, 15, 25, 35],
+ [5, 0, 2, 2, 16, 26, 36],
+ [3, 0, 2, 5, 17, 28, 37]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(
+ std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2], data_fields[3]},
std::vector<int32_t>({0, 1}),
+ /*is_ascending_order=*/true));
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator>
user_defined_seq_comparator,
+ FieldsComparator::Create({data_fields[2],
data_fields[3], data_fields[4],
+ data_fields[5],
data_fields[6]},
+ std::vector<int32_t>({2, 3}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {2, 1, 0, 2, 2, 3}, {{1, 1}, {1, 2}, {1, 3}, {2, 2}, {2, 3}, {2, 5}},
+ {{1, 1, 14, 24, 34},
+ {1, 2, 15, 25, 35},
+ {1, 3, 11, 21, 31},
+ {2, 2, 18, 28, 38},
+ {2, 3, 13, 23, 33},
+ {2, 5, 17, 28, 37}},
+ pool_);
+ CheckResult({src_array1, src_array2, src_array3}, user_key_comparator,
+ user_defined_seq_comparator, key_schema, value_schema,
expected);
+}
+
+TEST_F(SortMergeReaderTest, TestSortMergeWithAggMergeFunction) {
+ // key: k0, user defined sequence field: ts, value: v0
+ // Format: [_SEQUENCE_NUMBER, _VALUE_KIND, k0, ts, v0]
+ // Using sum aggregation: k0 uses primary-key agg, ts use last_value agg
and v0 use sum agg.
+ //
+ // Reader1 (SEQUENCE_NUMBER 0..5):
+ // [key=1,ts=1,v=1], [key=1,ts=2,v=2], [key=1,ts=3,v=3]
+ // [key=1,ts=4,v=4], [key=2,ts=4,v=40], [key=2,ts=5,v=50]
+ // Reader2 (SEQUENCE_NUMBER 6..11):
+ // [key=1,ts=5,v=5], [key=1,ts=6,v=6], [key=2,ts=1,v=10]
+ // [key=2,ts=2,v=20], [key=2,ts=3,v=30], [key=2,ts=6,v=60]
+ //
+ // With user_defined_seq_comparator on ts field, sort by key asc, then ts
asc within same key:
+ // key=1: ts=1(v=1), ts=2(v=2), ts=3(v=3), ts=4(v=4), ts=5(v=5), ts=6(v=6)
+ // key=2: ts=1(v=10), ts=2(v=20), ts=3(v=30), ts=4(v=40), ts=5(v=50),
ts=6(v=60)
+ //
+ // After sum aggregation:
+ // key=1: k0=1, ts=last_value(1,2,3,4,5,6)=6, v0=sum(1,2,3,4,5,6)=21, seq=7
+ // key=2: k0=2, ts=last_value(1,2,3,4,5,6)=6,
v0=sum(10,20,30,40,50,60)=210, seq=11
+
+ arrow::FieldVector fields = {
+ arrow::field("_SEQUENCE_NUMBER", arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("k0",
arrow::int32()),
+ arrow::field("ts", arrow::int32()), arrow::field("v0",
arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3], fields[4]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 1, 1],
+ [1, 0, 1, 2, 2],
+ [2, 0, 1, 3, 3],
+ [3, 0, 1, 4, 4],
+ [4, 0, 2, 4, 40],
+ [5, 0, 2, 5, 50]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [6, 0, 1, 5, 5],
+ [7, 0, 1, 6, 6],
+ [8, 0, 2, 1, 10],
+ [9, 0, 2, 2, 20],
+ [10, 0, 2, 3, 30],
+ [11, 0, 2, 6, 60]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ // user_defined_seq_comparator based on ts field (index 1 in value schema
{k0, ts, v0})
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator>
user_defined_seq_comparator,
+ FieldsComparator::Create({data_fields[2],
data_fields[3], data_fields[4]},
+ std::vector<int32_t>({1}),
+
/*is_ascending_order=*/true));
+ // Configure sum aggregation for all non-primary-key fields
+ std::string user_defined_sequence_field = "ts";
+ ASSERT_OK_AND_ASSIGN(
+ CoreOptions core_options,
+ CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"},
+ {Options::SEQUENCE_FIELD,
user_defined_sequence_field}}));
+
+ // After sum aggregation, same-key rows are merged:
+ // key=1: seq=7, k0=1, ts=6, v0=21
+ // key=2: seq=11, k0=2, ts=6, v0=210
+ std::vector<KeyValue> expected =
+ KeyValueChecker::GenerateKeyValues({7, 11}, {{1}, {2}}, {{1, 6, 21},
{2, 6, 210}}, pool_);
+ for (auto& kv : expected) {
+ kv.level = KeyValue::UNKNOWN_LEVEL;
+ }
+ CheckSortMergeResultForAggregate<SortMergeReaderWithLoserTree>(
+ {src_array1, src_array2}, user_key_comparator,
user_defined_seq_comparator, key_schema,
+ value_schema, {user_defined_sequence_field}, {"k0"}, core_options,
expected);
+ CheckSortMergeResultForAggregate<SortMergeReaderWithMinHeap>(
+ {src_array1, src_array2}, user_key_comparator,
user_defined_seq_comparator, key_schema,
+ value_schema, {user_defined_sequence_field}, {"k0"}, core_options,
expected);
+}
+
+TEST_F(SortMergeReaderTest, TestRawSortNoMergeKeepsDuplicateKeys) {
+ arrow::FieldVector fields = {arrow::field("_SEQUENCE_NUMBER",
arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()),
+ arrow::field("k0", arrow::int32()),
+ arrow::field("v0", arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 2, 10]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [2, 0, 2, 30]
+ ])")
+ .ValueOrDie());
+
+ auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 5, 50]
+ ])")
+ .ValueOrDie());
+
+ auto src_array4 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [1, 0, 2, 30]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {0, 1, 2, 1}, {{2}, {2}, {2}, {5}}, {{2, 10}, {2, 30}, {2, 30}, {5,
50}}, pool_);
+ CheckSortMergeResult<SortMergeReaderWithMinHeap>(
+ {src_array1, src_array2, src_array3, src_array4}, user_key_comparator,
+ /*user_defined_seq_comparator=*/nullptr, key_schema, value_schema,
expected,
+ /*need_merge=*/false);
+}
+
+TEST_F(SortMergeReaderTest, TestRawSortNoMergeWithMinHeap) {
+ // key: k0, user defined sequence field: ts, value: v0
+ // Format: [_SEQUENCE_NUMBER, _VALUE_KIND, k0, ts, v0]
+ // Reader1 (SEQUENCE_NUMBER 0..5):
+ // [key=1,ts=1,v=1], [key=1,ts=2,v=2], [key=1,ts=3,v=3]
+ // [key=1,ts=4,v=4], [key=2,ts=4,v=40], [key=2,ts=5,v=50]
+ // Reader2 (SEQUENCE_NUMBER 6..11):
+ // [key=1,ts=5,v=5], [key=1,ts=6,v=6], [key=2,ts=1,v=10]
+ // [key=2,ts=2,v=20], [key=2,ts=3,v=30], [key=2,ts=6,v=60]
+ //
+ // After sort:
+ // With user_defined_seq_comparator on ts field, sort by key asc, then ts
asc within same key
+ // key=1: ts=1(seq0,v=1), ts=2(seq1,v=2), ts=3(seq2,v=3), ts=4(seq3,v=4),
+ // ts=5(seq6,v=5), ts=6(seq7,v=6)
+ // key=2: ts=1(seq8,v=10), ts=2(seq9,v=20), ts=3(seq10,v=30),
ts=4(seq4,v=40),
+ // ts=5(seq5,v=50), ts=6(seq11,v=60)
+
+ arrow::FieldVector fields = {
+ arrow::field("_SEQUENCE_NUMBER", arrow::int64()),
+ arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("k0",
arrow::int32()),
+ arrow::field("ts", arrow::int32()), arrow::field("v0",
arrow::int32())};
+
+ auto data_fields = CreateDataField(fields);
+ std::shared_ptr<arrow::Schema> key_schema =
arrow::schema(arrow::FieldVector({fields[2]}));
+ std::shared_ptr<arrow::Schema> value_schema =
+ arrow::schema(arrow::FieldVector({fields[2], fields[3], fields[4]}));
+ std::shared_ptr<arrow::DataType> src_type = arrow::struct_(fields);
+
+ auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [0, 0, 1, 1, 1],
+ [1, 0, 1, 2, 2],
+ [2, 0, 1, 3, 3],
+ [3, 0, 1, 4, 4],
+ [4, 0, 2, 4, 40],
+ [5, 0, 2, 5, 50]
+ ])")
+ .ValueOrDie());
+
+ auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([
+ [6, 0, 1, 5, 5],
+ [7, 0, 1, 6, 6],
+ [8, 0, 2, 1, 10],
+ [9, 0, 2, 2, 20],
+ [10, 0, 2, 3, 30],
+ [11, 0, 2, 6, 60]
+ ])")
+ .ValueOrDie());
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator> user_key_comparator,
+ FieldsComparator::Create({data_fields[2]},
std::vector<int32_t>({0}),
+
/*is_ascending_order=*/true));
+ // user_defined_seq_comparator based on ts field (index 1 in value schema
{k0, ts, v0})
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<FieldsComparator>
user_defined_seq_comparator,
+ FieldsComparator::Create({data_fields[2],
data_fields[3], data_fields[4]},
+ std::vector<int32_t>({1}),
+
/*is_ascending_order=*/true));
+
+ std::vector<KeyValue> expected = KeyValueChecker::GenerateKeyValues(
+ {0, 1, 2, 3, 6, 7, 8, 9, 10, 4, 5, 11},
+ {{1}, {1}, {1}, {1}, {1}, {1}, {2}, {2}, {2}, {2}, {2}, {2}},
+ {{1, 1, 1},
+ {1, 2, 2},
+ {1, 3, 3},
+ {1, 4, 4},
+ {1, 5, 5},
+ {1, 6, 6},
+ {2, 1, 10},
+ {2, 2, 20},
+ {2, 3, 30},
+ {2, 4, 40},
+ {2, 5, 50},
+ {2, 6, 60}},
+ pool_);
+ CheckSortMergeResult<SortMergeReaderWithMinHeap>({src_array1, src_array2},
user_key_comparator,
+
user_defined_seq_comparator, key_schema,
+ value_schema, expected,
/*need_merge=*/false);
+}
+
+} // namespace paimon::test
diff --git
a/src/paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.cpp
b/src/paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.cpp
new file mode 100644
index 0000000..17b0009
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.cpp
@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h"
+
+#include <cassert>
+#include <cstdint>
+
+#include "paimon/common/utils/fields_comparator.h"
+#include "paimon/core/io/key_value_record_reader.h"
+
+namespace paimon {
+SortMergeReaderWithLoserTree::SortMergeReaderWithLoserTree(
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers,
+ const std::shared_ptr<FieldsComparator>& user_key_comparator,
+ const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
+ const std::shared_ptr<MergeFunctionWrapper<KeyValue>>&
merge_function_wrapper)
+ : merge_function_wrapper_(merge_function_wrapper) {
+ // if lhs and rhs are both null, it doesn't matter who becomes the new
winner. But if
+ // first_comparator returns 0, it means that second_comparator must be
used to compare
+ // again.
+ // lhs and rhs are swapped when compare to generate loser tree pop
smallest first
+ auto first_comparator = [user_key_comparator](const
std::optional<KeyValue>& lhs,
+ const
std::optional<KeyValue>& rhs) -> int32_t {
+ if (lhs == std::nullopt) {
+ return -1;
+ }
+ if (rhs == std::nullopt) {
+ return 1;
+ }
+ return user_key_comparator->CompareTo(*(rhs.value().key),
*(lhs.value().key));
+ };
+ auto second_comparator = [user_defined_seq_comparator](
+ const std::optional<KeyValue>& lhs,
+ const std::optional<KeyValue>& rhs) ->
int32_t {
+ if (lhs == std::nullopt) {
+ return -1;
+ }
+ if (rhs == std::nullopt) {
+ return 1;
+ }
+ if (user_defined_seq_comparator != nullptr) {
+ int32_t result =
+ user_defined_seq_comparator->CompareTo(*(rhs.value().value),
*(lhs.value().value));
+ if (result != 0) {
+ return result;
+ }
+ }
+ assert(lhs.value().sequence_number != rhs.value().sequence_number);
+ return rhs.value().sequence_number < lhs.value().sequence_number ? -1
: 1;
+ };
+ loser_tree_ =
+ std::make_unique<LoserTree>(std::move(readers), first_comparator,
second_comparator);
+}
+
+Result<bool> SortMergeReaderWithLoserTree::Iterator::HasNext() {
+ while (true) {
+ PAIMON_RETURN_NOT_OK(reader_->loser_tree_->AdjustForNextLoop());
+ std::optional<KeyValue> winner = reader_->loser_tree_->PopWinner();
+ if (winner == std::nullopt) {
+ return false;
+ }
+ reader_->merge_function_wrapper_->Reset();
+
PAIMON_RETURN_NOT_OK(reader_->merge_function_wrapper_->Add(std::move(winner.value())));
+ PAIMON_RETURN_NOT_OK(Merge());
+ if (result_ != std::nullopt) {
+ return true;
+ }
+ }
+ return false;
+}
+
+} // namespace paimon
diff --git
a/src/paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h
b/src/paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h
new file mode 100644
index 0000000..debeee2
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "paimon/common/utils/fields_comparator.h"
+#include "paimon/core/io/concat_key_value_record_reader.h"
+#include "paimon/core/key_value.h"
+#include "paimon/core/mergetree/compact/loser_tree.h"
+#include "paimon/core/mergetree/compact/merge_function_wrapper.h"
+#include "paimon/core/mergetree/compact/sort_merge_reader.h"
+#include "paimon/result.h"
+#include "paimon/status.h"
+
+namespace paimon {
+class FieldsComparator;
+class KeyValueRecordReader;
+class Metrics;
+
+/// `SortMergeReader` implemented with loser tree. Merge the KeyValue parsed by
+/// `KeyValueRecordReader` and return the iterator of KeyValue
+class SortMergeReaderWithLoserTree : public SortMergeReader {
+ public:
+ SortMergeReaderWithLoserTree(
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers,
+ const std::shared_ptr<FieldsComparator>& user_key_comparator,
+ const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
+ const std::shared_ptr<MergeFunctionWrapper<KeyValue>>&
merge_function_wrapper);
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return loser_tree_->GetReaderMetrics();
+ }
+
+ void Close() override {
+ loser_tree_->Close();
+ if (merge_function_wrapper_) {
+ merge_function_wrapper_->Reset();
+ }
+ }
+
+ class Iterator : public SortMergeReader::Iterator {
+ public:
+ explicit Iterator(SortMergeReaderWithLoserTree* reader) :
reader_(reader) {}
+
+ Result<bool> HasNext() override;
+
+ KeyValue&& Next() override {
+ return std::move(result_).value();
+ }
+
+ private:
+ Status Merge() {
+ while (reader_->loser_tree_->PeekWinner() != std::nullopt) {
+ PAIMON_RETURN_NOT_OK(reader_->merge_function_wrapper_->Add(
+ std::move(reader_->loser_tree_->PopWinner().value())));
+ }
+ PAIMON_ASSIGN_OR_RAISE(result_,
reader_->merge_function_wrapper_->GetResult());
+ return Status::OK();
+ }
+
+ private:
+ SortMergeReaderWithLoserTree* reader_;
+ std::optional<KeyValue> result_;
+ };
+
+ Result<std::unique_ptr<SortMergeReader::Iterator>> NextBatch() override {
+ PAIMON_RETURN_NOT_OK(loser_tree_->InitializeIfNeeded());
+ return loser_tree_->PeekWinner() == std::nullopt
+ ? std::unique_ptr<SortMergeReader::Iterator>()
+ : std::make_unique<Iterator>(this);
+ }
+
+ private:
+ std::unique_ptr<LoserTree> loser_tree_;
+ std::shared_ptr<MergeFunctionWrapper<KeyValue>> merge_function_wrapper_;
+};
+} // namespace paimon
diff --git
a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp
b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp
new file mode 100644
index 0000000..08629e6
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.cpp
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h"
+
+#include "paimon/core/mergetree/compact/merge_function_wrapper.h"
+#include "paimon/status.h"
+
+namespace paimon {
+class InternalRow;
+
+SortMergeReaderWithMinHeap::SortMergeReaderWithMinHeap(
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers,
+ const std::shared_ptr<FieldsComparator>& user_key_comparator,
+ const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
+ const std::shared_ptr<MergeFunctionWrapper<KeyValue>>&
merge_function_wrapper)
+ : need_merge_(merge_function_wrapper != nullptr),
+ readers_holder_(std::move(readers)),
+ user_key_comparator_(user_key_comparator),
+ merge_function_wrapper_(merge_function_wrapper),
+ min_heap_(HeapSorter(user_key_comparator, user_defined_seq_comparator)) {
+ next_batch_readers_.reserve(readers_holder_.size());
+ for (auto& reader : readers_holder_) {
+ next_batch_readers_.push_back(reader.get());
+ }
+}
+
+Result<std::unique_ptr<SortMergeReader::Iterator>>
SortMergeReaderWithMinHeap::NextBatch() {
+ for (auto* reader : next_batch_readers_) {
+ while (true) {
+
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<KeyValueRecordReader::Iterator> iterator,
+ reader->NextBatch());
+ if (!iterator) {
+ // no more batches, permanently remove this reader
+ reader->Close();
+ break;
+ }
+ PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext());
+ if (has_next) {
+ PAIMON_ASSIGN_OR_RAISE(KeyValue kv, iterator->Next());
+ min_heap_.emplace(std::move(kv), std::move(iterator), reader);
+ break;
+ }
+ }
+ }
+ next_batch_readers_.clear();
+ if (min_heap_.empty()) {
+ return std::unique_ptr<SortMergeReader::Iterator>();
+ }
+ return std::make_unique<SortMergeReaderWithMinHeap::Iterator>(this);
+}
+
+Result<bool> SortMergeReaderWithMinHeap::Iterator::HasNext() {
+ while (true) {
+ PAIMON_ASSIGN_OR_RAISE(bool has_more, NextImpl());
+ if (!reader_->need_merge_) {
+ // no merge: just return every kv in sorted order, possibly with
duplicate keys
+ return has_more;
+ }
+ if (!has_more) {
+ return false;
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::optional<KeyValue> result,
+ reader_->merge_function_wrapper_->GetResult());
+ if (result) {
+ result_ = std::move(result);
+ return true;
+ }
+ }
+}
+
+Result<bool> SortMergeReaderWithMinHeap::Iterator::NextImpl() {
+ assert(reader_->next_batch_readers_.empty());
+ // add previously polled elements back to priority queue
+ for (auto& element : reader_->polled_) {
+ PAIMON_ASSIGN_OR_RAISE(bool updated, element.Update());
+ if (updated) {
+ // still kvs left, add back to priority queue
+ reader_->min_heap_.push(std::move(element));
+ } else {
+ // reach end of batch, clean up
+ reader_->next_batch_readers_.push_back(std::move(element.reader));
+ }
+ }
+ reader_->polled_.clear();
+ if (!reader_->next_batch_readers_.empty()) {
+ return false;
+ }
+
+ if (reader_->min_heap_.empty()) {
+ return Status::Invalid("Min heap is empty. This is a bug.");
+ }
+
+ if (!reader_->need_merge_) {
+ // no merge: only poll the top element, set it as result directly
+ auto& element = const_cast<Element&>(reader_->min_heap_.top());
+ result_ = std::move(element.kv);
+ reader_->polled_.push_back(std::move(element));
+ reader_->min_heap_.pop();
+ return true;
+ }
+
+ reader_->merge_function_wrapper_->Reset();
+ std::shared_ptr<InternalRow> key = reader_->min_heap_.top().kv.key;
+ bool is_first = true;
+
+ // fetch all elements with the same key
+ // note that the same iterator should not produce the same keys, so this
code is correct
+ while (!reader_->min_heap_.empty()) {
+ auto& element = const_cast<Element&>(reader_->min_heap_.top());
+ if (!is_first && reader_->user_key_comparator_->CompareTo(*key,
*(element.kv.key)) != 0) {
+ break;
+ }
+
PAIMON_RETURN_NOT_OK(reader_->merge_function_wrapper_->Add(std::move(element.kv)));
+ reader_->polled_.push_back(std::move(element));
+ reader_->min_heap_.pop();
+ is_first = false;
+ }
+ return true;
+}
+
+} // namespace paimon
diff --git
a/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h
b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h
new file mode 100644
index 0000000..3d99100
--- /dev/null
+++ b/src/paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <queue>
+#include <utility>
+#include <vector>
+
+#include "paimon/common/metrics/metrics_impl.h"
+#include "paimon/common/utils/fields_comparator.h"
+#include "paimon/core/io/concat_key_value_record_reader.h"
+#include "paimon/core/io/key_value_record_reader.h"
+#include "paimon/core/key_value.h"
+#include "paimon/core/mergetree/compact/merge_function_wrapper.h"
+#include "paimon/core/mergetree/compact/sort_merge_reader.h"
+#include "paimon/result.h"
+
+namespace paimon {
+class Metrics;
+template <typename T>
+class MergeFunctionWrapper;
+
+/// `SortMergeReader` implemented with min-heap. Merge the KeyValue or only
sort the KeyValue parsed
+/// by `KeyValueRecordReader` and return the iterator of KeyValue
+class SortMergeReaderWithMinHeap : public SortMergeReader {
+ public:
+ SortMergeReaderWithMinHeap(
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers,
+ const std::shared_ptr<FieldsComparator>& user_key_comparator,
+ const std::shared_ptr<FieldsComparator>& user_defined_seq_comparator,
+ const std::shared_ptr<MergeFunctionWrapper<KeyValue>>&
merge_function_wrapper);
+
+ class Iterator : public SortMergeReader::Iterator {
+ public:
+ explicit Iterator(SortMergeReaderWithMinHeap* reader) :
reader_(reader) {}
+ Result<bool> HasNext() override;
+ KeyValue&& Next() override {
+ return std::move(result_).value();
+ }
+
+ private:
+ Result<bool> NextImpl();
+
+ private:
+ SortMergeReaderWithMinHeap* reader_;
+ std::optional<KeyValue> result_;
+ };
+
+ Result<std::unique_ptr<SortMergeReader::Iterator>> NextBatch() override;
+
+ void Close() override {
+ for (const auto& reader : readers_holder_) {
+ reader->Close();
+ }
+ if (merge_function_wrapper_) {
+ merge_function_wrapper_->Reset();
+ }
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return MetricsImpl::CollectReadMetrics(readers_holder_);
+ }
+
+ private:
+ struct Element {
+ Element(KeyValue&& _kv,
std::unique_ptr<KeyValueRecordReader::Iterator>&& _iterator,
+ KeyValueRecordReader* _reader)
+ : kv(std::move(_kv)), reader(_reader),
iterator(std::move(_iterator)) {
+ assert(iterator);
+ assert(reader);
+ }
+
+ Element(Element&& other) noexcept : kv(std::move(other.kv)) {
+ iterator = std::move(other.iterator);
+ reader = other.reader;
+ }
+
+ Element& operator=(Element&& other) noexcept {
+ if (&other == this) {
+ return *this;
+ }
+ kv = std::move(other.kv);
+ iterator = std::move(other.iterator);
+ reader = other.reader;
+ return *this;
+ }
+
+ Result<bool> Update() {
+ PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext());
+ if (!has_next) {
+ return false;
+ }
+ PAIMON_ASSIGN_OR_RAISE(KeyValue tmp_kv, iterator->Next());
+ kv = std::move(tmp_kv);
+ return true;
+ }
+
+ public:
+ KeyValue kv;
+ KeyValueRecordReader* reader;
+ std::unique_ptr<KeyValueRecordReader::Iterator> iterator;
+ };
+
+ class HeapSorter {
+ public:
+ HeapSorter(const std::shared_ptr<FieldsComparator>& key_comparator,
+ const std::shared_ptr<FieldsComparator>& seq_comparator)
+ : key_comparator_(key_comparator), seq_comparator_(seq_comparator)
{
+ assert(key_comparator_);
+ }
+ bool operator()(const Element& lhs, const Element& rhs) const {
+ int32_t result = key_comparator_->CompareTo(*(lhs.kv.key),
*(rhs.kv.key));
+ if (result != 0) {
+ return result > 0;
+ }
+ if (seq_comparator_ != nullptr) {
+ int32_t seq_result =
seq_comparator_->CompareTo(*(lhs.kv.value), *(rhs.kv.value));
+ if (seq_result != 0) {
+ return seq_result > 0;
+ }
+ }
+ return lhs.kv.sequence_number > rhs.kv.sequence_number;
+ }
+
+ private:
+ std::shared_ptr<FieldsComparator> key_comparator_;
+ std::shared_ptr<FieldsComparator> seq_comparator_;
+ };
+
+ private:
+ const bool need_merge_;
+ // must hold all readers, as data array is allocated by the pool of data
file reader
+ std::vector<std::unique_ptr<KeyValueRecordReader>> readers_holder_;
+ std::vector<KeyValueRecordReader*> next_batch_readers_;
+ std::shared_ptr<FieldsComparator> user_key_comparator_;
+ std::shared_ptr<MergeFunctionWrapper<KeyValue>> merge_function_wrapper_;
+ std::priority_queue<Element, std::vector<Element>, HeapSorter> min_heap_;
+ std::vector<Element> polled_;
+};
+} // namespace paimon
diff --git a/src/paimon/core/mergetree/drop_delete_reader.h
b/src/paimon/core/mergetree/drop_delete_reader.h
new file mode 100644
index 0000000..0d67caf
--- /dev/null
+++ b/src/paimon/core/mergetree/drop_delete_reader.h
@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include <memory>
+#include <optional>
+#include <utility>
+
+#include "paimon/common/types/row_kind.h"
+#include "paimon/core/key_value.h"
+#include "paimon/core/mergetree/compact/sort_merge_reader.h"
+#include "paimon/result.h"
+
+namespace paimon {
+class Metrics;
+
+/// A `RecordReader` which drops `KeyValue` that does not meet `RowKind#isAdd`
from
+/// the wrapped reader.
+class DropDeleteReader : public SortMergeReader {
+ public:
+ explicit DropDeleteReader(std::unique_ptr<SortMergeReader>&& reader)
+ : reader_(std::move(reader)) {}
+
+ class Iterator : public SortMergeReader::Iterator {
+ public:
+ explicit Iterator(std::unique_ptr<SortMergeReader::Iterator>&&
iterator)
+ : iterator_(std::move(iterator)) {}
+ Result<bool> HasNext() override {
+ while (true) {
+ PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator_->HasNext());
+ if (!has_next) {
+ return false;
+ }
+ result_ = std::move(iterator_->Next());
+ if (result_.value().value_kind->IsAdd()) {
+ break;
+ }
+ }
+ return true;
+ }
+ KeyValue&& Next() override {
+ return std::move(result_).value();
+ }
+
+ private:
+ std::optional<KeyValue> result_;
+ std::unique_ptr<SortMergeReader::Iterator> iterator_;
+ };
+
+ Result<std::unique_ptr<SortMergeReader::Iterator>> NextBatch() override {
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<SortMergeReader::Iterator> iter,
+ reader_->NextBatch());
+ if (iter == nullptr) {
+ return iter;
+ }
+ return std::make_unique<Iterator>(std::move(iter));
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return reader_->GetReaderMetrics();
+ }
+
+ void Close() override {
+ reader_->Close();
+ }
+
+ private:
+ std::unique_ptr<SortMergeReader> reader_;
+};
+} // namespace paimon
diff --git a/src/paimon/core/mergetree/drop_delete_reader_test.cpp
b/src/paimon/core/mergetree/drop_delete_reader_test.cpp
new file mode 100644
index 0000000..af09286
--- /dev/null
+++ b/src/paimon/core/mergetree/drop_delete_reader_test.cpp
@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/mergetree/drop_delete_reader.h"
+
+#include <cstddef>
+#include <variant>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paimon/common/data/internal_row.h"
+#include "paimon/memory/memory_pool.h"
+#include "paimon/status.h"
+#include "paimon/testing/utils/binary_row_generator.h"
+#include "paimon/testing/utils/read_result_collector.h"
+#include "paimon/testing/utils/testharness.h"
+
+namespace paimon {
+class Metrics;
+} // namespace paimon
+
+namespace paimon::test {
+class DropDeleteReaderTest : public testing::Test {
+ public:
+ class FakeSortMergeReader : public SortMergeReader {
+ public:
+ explicit FakeSortMergeReader(std::vector<KeyValue>&& data) :
data_(std::move(data)) {}
+
+ class Iterator : public SortMergeReader::Iterator {
+ public:
+ explicit Iterator(FakeSortMergeReader* reader) : reader_(reader) {}
+ Result<bool> HasNext() override {
+ return reader_->iter_ < reader_->data_.size();
+ }
+ KeyValue&& Next() override {
+ return std::move(reader_->data_[reader_->iter_++]);
+ }
+
+ private:
+ FakeSortMergeReader* reader_;
+ };
+
+ Result<std::unique_ptr<SortMergeReader::Iterator>> NextBatch()
override {
+ if (iter_ < data_.size()) {
+ return std::make_unique<Iterator>(this);
+ }
+ return std::unique_ptr<SortMergeReader::Iterator>();
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return nullptr;
+ }
+
+ void Close() override {}
+
+ private:
+ std::vector<KeyValue> data_;
+ size_t iter_ = 0;
+ };
+};
+
+TEST_F(DropDeleteReaderTest, TestSimple) {
+ auto pool = GetDefaultPool();
+ KeyValue kv1(RowKind::UpdateAfter(), /*sequence_number=*/1, /*level=*/0,
/*key=*/
+ BinaryRowGenerator::GenerateRowPtr({10}, pool.get()),
+ /*value=*/BinaryRowGenerator::GenerateRowPtr({10, 100},
pool.get()));
+ KeyValue kv2(RowKind::Delete(), /*sequence_number=*/2, /*level=*/2,
+ /*key=*/BinaryRowGenerator::GenerateRowPtr({20}, pool.get()),
+ /*value=*/BinaryRowGenerator::GenerateRowPtr({20, 200},
pool.get()));
+ KeyValue kv3(RowKind::Insert(), /*sequence_number=*/3, /*level=*/1,
/*key=*/
+ BinaryRowGenerator::GenerateRowPtr({30}, pool.get()),
+ /*value=*/BinaryRowGenerator::GenerateRowPtr({30, 300},
pool.get()));
+ KeyValue kv4(RowKind::UpdateBefore(), /*sequence_number=*/1, /*level=*/0,
/*key=*/
+ BinaryRowGenerator::GenerateRowPtr({40}, pool.get()),
+ /*value=*/BinaryRowGenerator::GenerateRowPtr({40, 100},
pool.get()));
+ std::vector<KeyValue> kvs;
+ kvs.reserve(4);
+ kvs.push_back(std::move(kv1));
+ kvs.push_back(std::move(kv2));
+ kvs.push_back(std::move(kv3));
+ kvs.push_back(std::move(kv4));
+
+ auto sort_merge_reader =
std::make_unique<FakeSortMergeReader>(std::move(kvs));
+ auto drop_delete_reader =
std::make_unique<DropDeleteReader>(std::move(sort_merge_reader));
+
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<KeyValue> results,
+ (ReadResultCollector::CollectKeyValueResult<SortMergeReader,
SortMergeReader::Iterator>(
+ drop_delete_reader.get())));
+ ASSERT_EQ(results.size(), 2);
+ ASSERT_EQ(results[0].value->GetInt(0), 10);
+ ASSERT_EQ(results[0].value->GetInt(1), 100);
+ ASSERT_EQ(results[1].value->GetInt(0), 30);
+ ASSERT_EQ(results[1].value->GetInt(1), 300);
+}
+} // namespace paimon::test
