KIKIMR-19521 BTreeIndex Seek RowId

kunga · kunga · commit 402da9580402 · 2023-11-10T20:08:34.000+03:00
diff --git a/ydb/core/tablet_flat/flat_page_btree_index.h b/ydb/core/tablet_flat/flat_page_btree_index.h
@@ -49,7 +49,6 @@ namespace NKikimr::NTable::NPage {
         using TColumns = TArrayRef<const TPartScheme::TColumn>;
 
 #pragma pack(push,1)
-    public:
         struct THeader {
             TRecIdx KeysCount;
             TPgSize KeysSize;
@@ -191,6 +190,11 @@ namespace NKikimr::NTable::NPage {
             return Header->KeysCount;
         }
 
+        TRecIdx GetChildrenCount() const noexcept
+        {
+            return GetKeysCount() + 1;
+        }
+
         TCellsIter GetKeyCells(TRecIdx pos, TColumns columns) const noexcept
         {
             if (IsFixedFormat()) {
@@ -207,7 +211,55 @@ namespace NKikimr::NTable::NPage {
             return Children[pos];
         }
 
-        // TODO: Seek methods will go here
+        TRecIdx Seek(TRowId rowId, std::optional<TRecIdx> on = { }) const noexcept
+        {
+            const TRecIdx childrenCount = GetChildrenCount();
+            if (on >= childrenCount) {
+                Y_DEBUG_ABORT_UNLESS(false, "Should point to some child");
+                on = { };
+            }
+
+            const auto cmp = [](TRowId rowId, const TChild& child) {
+                return rowId < child.Count;
+            };
+
+            TRecIdx result;
+            if (!on) {
+                // Use a full binary search
+                result = std::upper_bound(Children, Children + childrenCount, rowId, cmp) - Children;
+            } else if (Children[*on].Count <= rowId) {
+                // Try a short linear search first
+                result = *on;
+                for (int linear = 0; linear < 4; ++linear) {
+                    result++;
+                    Y_ABORT_UNLESS(result < childrenCount, "Should always seek some child");
+                    if (Children[result].Count > rowId) {
+                        return result;
+                    }
+                }
+
+                // Binary search from the next record
+                result = std::upper_bound(Children + result + 1, Children + childrenCount, rowId, cmp) - Children;
+            } else { // Children[*on].Count > rowId
+                // Try a short linear search first
+                result = *on;
+                for (int linear = 0; linear < 4; ++linear) {
+                    if (result == 0) {
+                        return 0;
+                    }
+                    if (Children[result - 1].Count <= rowId) {
+                        return result;
+                    }
+                    result--;
+                }
+
+                // Binary search up to current record
+                result = std::upper_bound(Children, Children + result, rowId, cmp) - Children;
+            }
+
+            Y_ABORT_UNLESS(result < childrenCount, "Should always seek some child");
+            return result;
+        }
 
     private:
         TSharedData Raw;
diff --git a/ydb/core/tablet_flat/flat_part_btree_index_iter.h b/ydb/core/tablet_flat/flat_part_btree_index_iter.h
@@ -0,0 +1,161 @@
+#pragma once
+
+#include "flat_part_iface.h"
+#include "flat_page_index.h"
+#include "flat_table_part.h"
+
+
+namespace NKikimr::NTable {
+
+class TPartBtreeIndexIt {
+    using TCells = NPage::TCells;
+    using TBtreeIndexNode = NPage::TBtreeIndexNode;
+    using TGroupId = NPage::TGroupId;
+    using TRecIdx = NPage::TRecIdx;
+    using TChild = TBtreeIndexNode::TChild;
+    using TBtreeIndexMeta = NPage::TBtreeIndexMeta;
+
+    struct TNodeState {
+        TChild Meta;
+        TRowId BeginRowId;
+        TRowId EndRowId;
+        // TCells BeginKey;
+        // TCells EndKey;
+        std::optional<TBtreeIndexNode> Node;
+        std::optional<TRecIdx> Pos;
+
+        TNodeState(TChild meta, TRowId beginRowId, TRowId endRowId)
+            : Meta(meta)
+            , BeginRowId(beginRowId)
+            , EndRowId(endRowId)
+        { 
+        }
+
+        bool HasRow(TRowId rowId) const {
+            return BeginRowId <= rowId && rowId < EndRowId;
+        }
+    };
+
+public:
+    TPartBtreeIndexIt(const TPart* part, IPages* env, TGroupId groupId)
+        : Part(part)
+        , Env(env)
+        , GroupId(groupId)
+        , Meta(groupId.IsMain() ? part->IndexPages.BTreeGroups[groupId.Index] : part->IndexPages.BTreeHistoric[groupId.Index])
+    {
+        State.emplace_back(Meta, 0, Meta.Count);
+    }
+    
+    EReady Seek(TRowId rowId) {
+        if (rowId >= Meta.Count) {
+            return Exhaust();
+        }
+
+        while (State.size() > 1 && !State.back().HasRow(rowId)) {
+            State.pop_back();
+        }
+
+        if (IsExhausted()) {
+            // don't use exhausted state as an initial one
+            State[0].Pos = { };
+        }
+
+        for (size_t level : xrange(State.size() - 1, Meta.LevelsCount)) {
+            auto &state = State[level];
+            Y_ABORT_UNLESS(state.HasRow(rowId));
+            if (!TryLoad(state)) {
+                // exiting with an intermediate state
+                Y_DEBUG_ABORT_UNLESS(!IsLeaf() && !IsExhausted());
+                return EReady::Page;
+            }
+            auto pos = state.Node->Seek(rowId, state.Pos);
+            state.Pos.emplace(pos);
+            
+            auto child = state.Node->GetChild(pos);
+            TRowId firstRowId = pos ? state.Node->GetChild(pos - 1).Count : state.BeginRowId;
+            TRowId lastRowId = child.Count;
+            State.emplace_back(child, firstRowId, lastRowId);
+        }
+
+        // State.back() points to the target data page
+        Y_ABORT_UNLESS(IsLeaf());
+        Y_ABORT_UNLESS(State.back().HasRow(rowId));
+        return EReady::Data;
+    }
+
+    // EReady Next() {
+    //     Y_DEBUG_ABORT_UNLESS(IsLeaf());
+    //     return Exhaust();
+    // }
+
+    // EReady Prev() {
+    //     Y_DEBUG_ABORT_UNLESS(IsLeaf());
+    //     return Exhaust();
+    // }
+
+public:
+    bool IsValid() const {
+        Y_DEBUG_ABORT_UNLESS(IsLeaf() || IsExhausted());
+        return IsLeaf();
+    }
+
+    TPageId GetPageId() const {
+        Y_ABORT_UNLESS(IsLeaf());
+        return State.back().Meta.PageId;
+    }
+
+    TRowId GetRowId() const {
+        Y_ABORT_UNLESS(IsLeaf());
+        return State.back().BeginRowId;
+    }
+
+    TRowId GetNextRowId() const {
+        Y_ABORT_UNLESS(IsLeaf());
+        return State.back().EndRowId;
+    }
+
+private:
+    bool IsRoot() const noexcept {
+        return State.size() == 1;
+    }
+    
+    bool IsExhausted() const noexcept {
+        return State[0].Pos == Max<TRecIdx>();
+    }
+
+    bool IsLeaf() const noexcept {
+        // Note: it is possible to have 0 levels in B-Tree
+        // so we may have exhausted state with leaf (data) node
+        return State.size() == Meta.LevelsCount + 1 && !IsExhausted();
+    }
+
+    EReady Exhaust() {
+        while (State.size() > 1) {
+            State.pop_back();
+        }
+        State[0].Pos = Max<TRecIdx>();
+        return EReady::Gone;
+    }
+
+    bool TryLoad(TNodeState& state) {
+        if (state.Node) {
+            return true;
+        }
+
+        auto page = Env->TryGetPage(Part, state.Meta.PageId);
+        if (page) {
+            state.Node.emplace(*page);
+            return true;
+        }
+        return false;
+    }
+
+private:
+    const TPart* const Part;
+    IPages* const Env;
+    const TGroupId GroupId;
+    const TBtreeIndexMeta Meta;
+    TVector<TNodeState> State;
+};
+
+}
diff --git a/ydb/core/tablet_flat/flat_part_index_iter.h b/ydb/core/tablet_flat/flat_part_index_iter.h
@@ -24,13 +24,13 @@ class TPartIndexIt {
         , EndRowId(groupId.IsMain() && part->Stat.Rows ? part->Stat.Rows : Max<TRowId>())
     { }
     
-    EReady Seek(TRowId rowId, bool restart = false) {
+    EReady Seek(TRowId rowId) {
         auto index = TryGetIndex();
         if (!index) {
             return EReady::Page;
         }
 
-        Iter = index->LookupRow(rowId, restart ? TIter() : Iter);
+        Iter = index->LookupRow(rowId, Iter);
         return DataOrGone();
     }
 
@@ -54,42 +54,39 @@ class TPartIndexIt {
         return DataOrGone();
     }
 
-    EReady Next() {
-        auto index = TryGetIndex();
-        if (!index) {
-            return EReady::Page;
-        }
-        Iter++;
-        return DataOrGone();
-    }
-
-    EReady Prev() {
+    EReady SeekLast() {
         auto index = TryGetIndex();
         if (!index) {
             return EReady::Page;
         }
+        Iter = (*index)->End();
         if (Iter.Off() == 0) {
-            Iter = { };
             return EReady::Gone;
         }
         Iter--;
         return DataOrGone();
     }
 
-    EReady SeekLast() {
-        auto index = TryGetIndex();
-        if (!index) {
-            return EReady::Page;
-        }
-        Iter = (*index)->End();
+    EReady Next() {
+        Y_DEBUG_ABORT_UNLESS(Index);
+        Y_DEBUG_ABORT_UNLESS(Iter);
+        Iter++;
+        return DataOrGone();
+    }
+
+    EReady Prev() {
+        Y_DEBUG_ABORT_UNLESS(Index);
+        Y_DEBUG_ABORT_UNLESS(Iter);
         if (Iter.Off() == 0) {
+            Iter = { };
             return EReady::Gone;
         }
         Iter--;
         return DataOrGone();
     }
 
     bool IsValid() const {
+        Y_DEBUG_ABORT_UNLESS(Index);
         return bool(Iter);
     }
 
@@ -125,10 +122,11 @@ class TPartIndexIt {
 
     TRowId GetNextRowId() const {
         Y_ABORT_UNLESS(Index);
+        Y_ABORT_UNLESS(Iter);
         auto next = Iter + 1;
         return next
             ? next->GetRowId()
-            : Max<TRowId>();
+            : EndRowId;
     }
 
     const TRecord * GetRecord() const {
diff --git a/ydb/core/tablet_flat/flat_stat_part.h b/ydb/core/tablet_flat/flat_stat_part.h
@@ -59,7 +59,9 @@ class TScreenedPartIndexIterator {
 
     EReady Start() {
         auto ready = Pos.Seek(0);
-        FillKey();
+        if (ready != EReady::Page) {
+            FillKey();
+        }
 
         for (auto& g : AltGroups) {
             if (g.Pos.Seek(0) == EReady::Page) {
diff --git a/ydb/core/tablet_flat/ut/ut_btree_index.cpp b/ydb/core/tablet_flat/ut/ut_btree_index.cpp
diff --git a/ydb/core/tablet_flat/ut/ut_stat.cpp b/ydb/core/tablet_flat/ut/ut_stat.cpp
