KIKIMR-19521 BTreeIndex Iter Multi

ydb/core/tablet_flat/CMakeLists.darwin-arm64.txt

@@ -110,6 +110,7 @@ target_sources(ydb-core-tablet_flat PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_page_label.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_dump.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_iter_multi.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_index_iter_create.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_loader.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_overlay.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_outset.cpp

ydb/core/tablet_flat/CMakeLists.darwin-x86_64.txt

@@ -110,6 +110,7 @@ target_sources(ydb-core-tablet_flat PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_page_label.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_dump.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_iter_multi.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_index_iter_create.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_loader.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_overlay.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_outset.cpp

ydb/core/tablet_flat/CMakeLists.linux-aarch64.txt

@@ -111,6 +111,7 @@ target_sources(ydb-core-tablet_flat PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_page_label.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_dump.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_iter_multi.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_index_iter_create.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_loader.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_overlay.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_outset.cpp

ydb/core/tablet_flat/CMakeLists.linux-x86_64.txt

@@ -111,6 +111,7 @@ target_sources(ydb-core-tablet_flat PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_page_label.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_dump.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_iter_multi.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_index_iter_create.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_loader.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_overlay.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_outset.cpp

ydb/core/tablet_flat/CMakeLists.windows-x86_64.txt

@@ -110,6 +110,7 @@ target_sources(ydb-core-tablet_flat PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_page_label.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_dump.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_iter_multi.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_index_iter_create.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_loader.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_overlay.cpp
   ${CMAKE_SOURCE_DIR}/ydb/core/tablet_flat/flat_part_outset.cpp

ydb/core/tablet_flat/flat_fwd_env.h

@@ -113,8 +113,8 @@ namespace NFwd {
             ui16 room = (groupId.Historic ? part->GroupsCount + 2 : 0) + groupId.Index;
             TSlot slot = GetQueueSlot(part, room);
 
-            if (part->GetPageType(ref, groupId) == EPage::Index) {
-                return TryGetIndexPage(slot, ref);
+            if (auto type = part->GetPageType(ref, groupId); IsIndexPage(type)) {
+                return TryGetIndexPage(slot, ref, type);
             }
                 
             return Handle(Queues.at(slot), ref).Page;
@@ -162,7 +162,7 @@ namespace NFwd {
             TVector<NPageCollection::TLoadedPage> queuePages(Reserve(pages.size()));
             for (auto& page : pages) {
                 const auto &meta = pageCollection->Page(page.PageId);
-                if (NTable::EPage(meta.Type) == EPage::Index) {
+                if (IsIndexPage(NTable::EPage(meta.Type))) {
                     IndexPages.at(part).Pages[page.PageId] = page;
                 } else {
                     queuePages.push_back(page);
@@ -259,8 +259,14 @@ namespace NFwd {
         }
 
     private:
-        const TSharedData* TryGetIndexPage(TSlot slot, TPageId pageId) noexcept
+        static bool IsIndexPage(EPage type) noexcept {
+            return type == EPage::Index || type == EPage::BTreeIndex;
+        }
+
+        const TSharedData* TryGetIndexPage(TSlot slot, TPageId pageId, EPage type) noexcept
         {
+            Y_DEBUG_ABORT_UNLESS(IsIndexPage(type));
+
             // TODO: count index pages in Stats later
 
             auto &env = IndexPages.at(slot);
@@ -270,7 +276,7 @@ namespace NFwd {
                 return &pageIt->second.Data;
             } else {
                 auto &queue = Queues.at(slot);
-                queue.AddToQueue(pageId, EPage::Index);
+                queue.AddToQueue(pageId, type);
                 Queue.PushBack(&queue);
                 return nullptr;
             }            

ydb/core/tablet_flat/flat_page_btree_index.h

@@ -150,6 +150,20 @@ namespace NKikimr::NTable::NPage {
                 return result;
             }
 
+            TCell At(TPos index)
+            {
+                Y_ABORT_UNLESS(Pos == 0, "Shouldn't be used");
+
+                // We may use offset = Columns[index].Offset - index * sizeof(TIndex::TItem) for fixed format
+                // But it looks too complicated because this method will be used only for history columns 0, 1, 2
+                Y_DEBUG_ABORT_UNLESS(index < 3);
+
+                for (TPos i = 0; i < index; i++) {
+                    Next();
+                }
+                return Next();
+            }
+
             int CompareTo(const TCells key, const TKeyCellDefaults *keyDefaults) noexcept
             {
                 Y_ABORT_UNLESS(Pos == 0, "Shouldn't be used");

ydb/core/tablet_flat/flat_part_btree_index_iter.h

@@ -3,11 +3,11 @@
 #include "flat_part_iface.h"
 #include "flat_page_index.h"
 #include "flat_table_part.h"
-
+#include "flat_part_index_iter_iface.h"
 
 namespace NKikimr::NTable {
 
-class TPartBtreeIndexIt {
+class TPartBtreeIndexIt : public IIndexIter {
     using TCells = NPage::TCells;
     using TBtreeIndexNode = NPage::TBtreeIndexNode;
     using TGroupId = NPage::TGroupId;
@@ -33,7 +33,19 @@ class TPartBtreeIndexIt {
             , EndRowId(endRowId)
             , BeginKey(beginKey)
             , EndKey(endKey)
-        { 
+        {
+        }
+
+        bool IsLastPos() const noexcept {
+            Y_ABORT_UNLESS(Node);
+            Y_ABORT_UNLESS(Pos);
+            return *Pos == Node->GetKeysCount();
+        }
+
+        bool IsFirstPos() const noexcept {
+            Y_ABORT_UNLESS(Node);
+            Y_ABORT_UNLESS(Pos);
+            return *Pos == 0;
         }
     };
 
@@ -103,13 +115,13 @@ public:
         , Env(env)
         , GroupId(groupId)
         , GroupInfo(part->Scheme->GetLayout(groupId))
-        , Meta(groupId.IsMain() ? part->IndexPages.BTreeGroups[groupId.Index] : part->IndexPages.BTreeHistoric[groupId.Index])
+        , Meta(groupId.IsHistoric() ? part->IndexPages.BTreeHistoric[groupId.Index] : part->IndexPages.BTreeGroups[groupId.Index])
     {
         const static TCellsIterable EmptyKey(static_cast<const char*>(nullptr), TColumns());
         State.emplace_back(Meta, 0, GetEndRowId(), EmptyKey, EmptyKey);
     }
     
-    EReady Seek(TRowId rowId) {
+    EReady Seek(TRowId rowId) override {
         if (rowId >= GetEndRowId()) {
             return Exhaust();
         }
@@ -122,7 +134,7 @@ public:
      *
      * Result is approximate and may be off by one page
      */
-    EReady Seek(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) {
+    EReady Seek(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) override {
         if (!key) {
             // Special treatment for an empty key
             switch (seek) {
@@ -142,7 +154,7 @@ public:
      *
      * Result is approximate and may be off by one page
      */
-    EReady SeekReverse(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) {
+    EReady SeekReverse(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) override {
         if (!key) {
             // Special treatment for an empty key
             switch (seek) {
@@ -157,39 +169,101 @@ public:
         return DoSeek<TSeekKeyReverse>({seek, key, GroupInfo.ColsKeyIdx, keyDefaults});
     }
 
-    // EReady Next() {
-    //     Y_DEBUG_ABORT_UNLESS(IsLeaf());
-    //     return Exhaust();
-    // }
+    EReady Next() override {
+        Y_ABORT_UNLESS(!IsExhausted());
 
-    // EReady Prev() {
-    //     Y_DEBUG_ABORT_UNLESS(IsLeaf());
-    //     return Exhaust();
-    // }
+        if (Meta.LevelsCount == 0) {
+            return Exhaust();
+        }
+
+        if (IsLeaf()) {
+            do {
+                State.pop_back();
+            } while (State.size() > 1 && State.back().IsLastPos());
+            if (State.back().IsLastPos()) {
+                return Exhaust();
+            }
+            PushNextState(*State.back().Pos + 1);
+        }
+
+        for (size_t level : xrange(State.size() - 1, Meta.LevelsCount)) {
+            if (!TryLoad(State[level])) {
+                // exiting with an intermediate state
+                Y_DEBUG_ABORT_UNLESS(!IsLeaf() && !IsExhausted());
+                return EReady::Page;
+            }
+            PushNextState(0);
+        }
+
+        // State.back() points to the target data page
+        Y_ABORT_UNLESS(IsLeaf());
+        return EReady::Data;
+    }
+
+    EReady Prev() override {
+        Y_ABORT_UNLESS(!IsExhausted());
+
+        if (Meta.LevelsCount == 0) {
+            return Exhaust();
+        }
+
+        if (IsLeaf()) {
+            do {
+                State.pop_back();
+            } while (State.size() > 1 && State.back().IsFirstPos());
+            if (State.back().IsFirstPos()) {
+                return Exhaust();
+            }
+            PushNextState(*State.back().Pos - 1);
+        }
+
+        for (size_t level : xrange(State.size() - 1, Meta.LevelsCount)) {
+            if (!TryLoad(State[level])) {
+                // exiting with an intermediate state
+                Y_DEBUG_ABORT_UNLESS(!IsLeaf() && !IsExhausted());
+                return EReady::Page;
+            }
+            PushNextState(State[level].Node->GetKeysCount());
+        }
+
+        // State.back() points to the target data page
+        Y_ABORT_UNLESS(IsLeaf());
+        return EReady::Data;
+    }
 
 public:
-    bool IsValid() const {
+    bool IsValid() const override {
         Y_DEBUG_ABORT_UNLESS(IsLeaf() || IsExhausted());
         return IsLeaf();
     }
 
-    TPageId GetPageId() const {
+    TRowId GetEndRowId() const override {
+        return Meta.Count;
+    }
+
+    TPageId GetPageId() const override {
         Y_ABORT_UNLESS(IsLeaf());
         return State.back().Meta.PageId;
     }
 
-    TRowId GetRowId() const {
+    TRowId GetRowId() const override {
         Y_ABORT_UNLESS(IsLeaf());
         return State.back().BeginRowId;
     }
 
-    TRowId GetNextRowId() const {
+    TRowId GetNextRowId() const override {
         Y_ABORT_UNLESS(IsLeaf());
         return State.back().EndRowId;
     }
 
-    TRowId GetEndRowId() const {
-        return Meta.Count;
+    bool HasKeyCells() const override {
+        Y_ABORT_UNLESS(IsLeaf());
+        return State.back().BeginKey.Count();
+    }
+
+    TCell GetKeyCell(TPos index) const override {
+        Y_ABORT_UNLESS(IsLeaf());
+        return State.back().BeginKey.Iter().At(index);
     }
 
 private:
@@ -214,7 +288,7 @@ private:
             }
             auto pos = seek.Do(state);
             
-            PushNextState(state, pos);
+            PushNextState(pos);
         }
 
         // State.back() points to the target data page
@@ -245,7 +319,8 @@ private:
         return EReady::Gone;
     }
 
-    void PushNextState(TNodeState& current, TRecIdx pos) {
+    void PushNextState(TRecIdx pos) {
+        TNodeState& current = State.back();
         Y_ABORT_UNLESS(pos < current.Node->GetChildrenCount(), "Should point to some child");
         current.Pos.emplace(pos);
 

ydb/core/tablet_flat/flat_part_index_iter.h

@@ -2,13 +2,14 @@
 
 #include "flat_part_iface.h"
 #include "flat_page_index.h"
+#include "flat_part_index_iter_iface.h"
 #include "flat_table_part.h"
 #include <ydb/library/yverify_stream/yverify_stream.h>
 
 
 namespace NKikimr::NTable {
 
-class TPartIndexIt {
+class TPartIndexIt : public IIndexIter {
 public:
     using TCells = NPage::TCells;
     using TRecord = NPage::TIndex::TRecord;
@@ -25,7 +26,7 @@ public:
         , EndRowId(groupId.IsMain() && part->Stat.Rows ? part->Stat.Rows : Max<TRowId>())
     { }
     
-    EReady Seek(TRowId rowId) {
+    EReady Seek(TRowId rowId) override {
         auto index = TryGetIndex();
         if (!index) {
             return EReady::Page;
@@ -35,7 +36,7 @@ public:
         return DataOrGone();
     }
 
-    EReady Seek(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) {
+    EReady Seek(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) override {
         auto index = TryGetIndex();
         if (!index) {
             return EReady::Page;
@@ -45,7 +46,7 @@ public:
         return DataOrGone();
     }
 
-    EReady SeekReverse(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) {
+    EReady SeekReverse(ESeek seek, TCells key, const TKeyCellDefaults *keyDefaults) override {
         auto index = TryGetIndex();
         if (!index) {
             return EReady::Page;
@@ -68,14 +69,14 @@ public:
         return DataOrGone();
     }
 
-    EReady Next() {
+    EReady Next() override {
         Y_DEBUG_ABORT_UNLESS(Index);
         Y_DEBUG_ABORT_UNLESS(Iter);
         Iter++;
         return DataOrGone();
     }
 
-    EReady Prev() {
+    EReady Prev() override {
         Y_DEBUG_ABORT_UNLESS(Index);
         Y_DEBUG_ABORT_UNLESS(Iter);
         if (Iter.Off() == 0) {
@@ -86,7 +87,7 @@ public:
         return DataOrGone();
     }
 
-    bool IsValid() const {
+    bool IsValid() const override {
         Y_DEBUG_ABORT_UNLESS(Index);
         return bool(Iter);
     }
@@ -97,23 +98,23 @@ public:
     }
 
 public:
-    TRowId GetEndRowId() const {
+    TRowId GetEndRowId() const override {
         return EndRowId;
     }
 
-    TPageId GetPageId() const {
+    TPageId GetPageId() const override {
         Y_ABORT_UNLESS(Index);
         Y_ABORT_UNLESS(Iter);
         return Iter->GetPageId();
     }
 
-    TRowId GetRowId() const {
+    TRowId GetRowId() const override {
         Y_ABORT_UNLESS(Index);
         Y_ABORT_UNLESS(Iter);
         return Iter->GetRowId();
     }
 
-    TRowId GetNextRowId() const {
+    TRowId GetNextRowId() const override {
         Y_ABORT_UNLESS(Index);
         Y_ABORT_UNLESS(Iter);
         auto next = Iter + 1;
@@ -122,6 +123,18 @@ public:
             : EndRowId;
     }
 
+    bool HasKeyCells() const override {
+        Y_ABORT_UNLESS(Index);
+        Y_ABORT_UNLESS(Iter);
+        return true;
+    }
+
+    TCell GetKeyCell(TPos index) const override {
+        Y_ABORT_UNLESS(Index);
+        Y_ABORT_UNLESS(Iter);
+        return Iter.GetRecord()->Cell(GroupInfo.ColsKeyIdx[index]);
+    }
+
     const TRecord * GetRecord() const {
         Y_ABORT_UNLESS(Index);
         Y_ABORT_UNLESS(Iter);

ydb/core/tablet_flat/flat_part_index_iter_create.cpp

@@ -0,0 +1,15 @@
+#include "flat_part_index_iter.h"
+#include "flat_part_btree_index_iter.h"
+
+namespace NKikimr::NTable {
+
+THolder<IIndexIter> CreateIndexIter(const TPart* part, IPages* env, NPage::TGroupId groupId)
+{
+    if (groupId.Index < (groupId.IsHistoric() ? part->IndexPages.BTreeHistoric : part->IndexPages.BTreeGroups).size()) {
+        return MakeHolder<TPartBtreeIndexIt>(part, env, groupId);
+    } else {
+        return MakeHolder<TPartIndexIt>(part, env, groupId);
+    }
+}
+
+}