feat: Migrate sort merge reader (#85)

Author: lxy-9602

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

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

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