KIKIMR-19139 Load index in TKeysLoader

4 files changed, 93 insertions, 45 deletions

diff --git a/ydb/core/tablet_flat/flat_part_index_iter.h b/ydb/core/tablet_flat/flat_part_index_iter.h
index 4d38837426..886f31cc7a 100644
--- a/ydb/core/tablet_flat/flat_part_index_iter.h
+++ b/ydb/core/tablet_flat/flat_part_index_iter.h
@@ -76,6 +76,19 @@ public:
         return DataOrGone();
     }
 
+    EReady SeekLast() {
+        auto index = TryGetIndex();
+        if (!index) {
+            return EReady::Page;
+        }
+        Iter = (*index)->End();
+        if (Iter.Off() == 0) {
+            return EReady::Gone;
+        }
+        Iter--;
+        return DataOrGone();
+    }
+
     bool IsValid() {
         return bool(Iter);
     }
@@ -92,11 +105,13 @@ public:
 
     TPageId GetPageId() {
         Y_VERIFY(Index);
+        Y_VERIFY(Iter);
         return Iter->GetPageId();
     }
 
     TRowId GetRowId() {
         Y_VERIFY(Index);
+        Y_VERIFY(Iter);
         return Iter->GetRowId();
     }
 
@@ -107,11 +122,18 @@ public:
             ? next->GetRowId()
             : Max<TRowId>();
     }
+
     const TRecord * GetRecord() {
         Y_VERIFY(Index);
+        Y_VERIFY(Iter);
         return Iter.GetRecord();
     }
 
+    const TRecord * TryGetLastRecord() {
+        Y_VERIFY(Index);
+        return Index->GetLastKeyRecord();
+    }
+
 private:
     EReady DataOrGone() {
         return Iter ? EReady::Data : EReady::Gone;
diff --git a/ydb/core/tablet_flat/flat_part_keys.h b/ydb/core/tablet_flat/flat_part_keys.h
index dde566d092..deaf8ce51a 100644
--- a/ydb/core/tablet_flat/flat_part_keys.h
+++ b/ydb/core/tablet_flat/flat_part_keys.h
@@ -1,5 +1,6 @@
 #pragma once
 #include "flat_part_iface.h"
+#include "flat_part_index_iter.h"
 #include "flat_part_slice.h"
 #include "flat_sausage_fetch.h"
 #include "flat_sausagecache.h"
@@ -81,8 +82,8 @@ namespace NTable {
         explicit TKeysLoader(const TPart* part, IPages* env)
             : Part(part)
             , Env(env)
+            , Index(Part, Env, {})
         {
-
         }
 
         TIntrusivePtr<TSlices> Do(TIntrusiveConstPtr<TScreen> screen)
@@ -149,7 +150,14 @@ namespace NTable {
                     Key = { };
                     return true;
                 }
-                if (const auto* lastKey = Part->Index.GetLastKeyRecord()) {
+
+                auto hasLast = Index.SeekLast();
+                if (hasLast == EReady::Page) {
+                    return false;
+                }
+                Y_VERIFY(hasLast != EReady::Gone, "Unexpected failure to find the last index record");
+
+                if (const auto* lastKey = Index.TryGetLastRecord()) {
                     if (lastKey->GetRowId() == rowId) {
                         LoadIndexKey(*lastKey);
                         return true;
@@ -160,16 +168,19 @@ namespace NTable {
                         return true;
                     }
                 }
-                SeekIndex(rowId);
-                if (Index->GetRowId() == rowId) {
-                    LoadIndexKey(*Index);
+
+                if (!SeekIndex(rowId)) {
+                    return false;
+                }
+                if (Index.GetRowId() == rowId) {
+                    LoadIndexKey(*Index.GetRecord());
                     return true;
                 }
-                Y_VERIFY(Index->GetRowId() < rowId, "SeekIndex invariant failure");
-                if (!LoadPage(Index->GetPageId())) {
+                Y_VERIFY(Index.GetRowId() < rowId, "SeekIndex invariant failure");
+                if (!LoadPage(Index.GetPageId())) {
                     return false;
                 }
-                Y_VERIFY(Page.BaseRow() == Index->GetRowId(), "Index and data are out of sync");
+                Y_VERIFY(Page.BaseRow() == Index.GetRowId(), "Index and data are out of sync");
                 auto lastRowId = Page.BaseRow() + (Page->Count - 1);
                 if (lastRowId < rowId) {
                     // Row is out of range for this page
@@ -184,25 +195,31 @@ namespace NTable {
 
         bool SeekLastRow() noexcept
         {
-            if (const auto* lastKey = Part->Index.GetLastKeyRecord()) {
+            auto hasLast = Index.SeekLast();
+            if (hasLast == EReady::Page) {
+                return false;
+            }
+            Y_VERIFY(hasLast != EReady::Gone, "Unexpected failure to find the last index record");
+
+            if (const auto* lastKey = Index.TryGetLastRecord()) {
                 LoadIndexKey(*lastKey);
                 return true;
             }
-            auto it = Part->Index->End();
-            Y_VERIFY(--it, "Unexpected failure to find the last index record");
-            if (!LoadPage(it->GetPageId())) {
+
+            if (!LoadPage(Index.GetPageId())) {
                 return false;
             }
-            Y_VERIFY(Page.BaseRow() == it->GetRowId(), "Index and data are out of sync");
+            Y_VERIFY(Page.BaseRow() == Index.GetRowId(), "Index and data are out of sync");
             auto lastRowId = Page.BaseRow() + (Page->Count - 1);
             LoadRow(lastRowId);
             return true;
         }
 
-        void SeekIndex(TRowId rowId) noexcept
+        bool SeekIndex(TRowId rowId) noexcept
         {
-            Index = Part->Index.LookupRow(rowId, Index);
-            Y_VERIFY(Index, "SeekIndex called with an out of bounds row");
+            auto ready = Index.Seek(rowId);
+            Y_VERIFY(ready != EReady::Gone, "SeekIndex called with an out of bounds row");
+            return ready == EReady::Data;
         }
 
         bool LoadPage(TPageId pageId) noexcept
@@ -248,7 +265,7 @@ namespace NTable {
         IPages* Env;
         TRowId RowId = Max<TRowId>();
         TPageId PageId = Max<TPageId>();
-        NPage::TIndex::TIter Index;
+        TPartIndexIt Index;
         NPage::TDataPage Page;
         TSmallVec<TCell> Key;
     };
diff --git a/ydb/core/tablet_flat/flat_part_loader.h b/ydb/core/tablet_flat/flat_part_loader.h
index ab3d02d1d1..32ee5de7f3 100644
--- a/ydb/core/tablet_flat/flat_part_loader.h
+++ b/ydb/core/tablet_flat/flat_part_loader.h
@@ -69,13 +69,9 @@ namespace NTable {
                     if (!fetch->Pages) {
                         Y_Fail("TLoader is trying to fetch 0 pages");
                     }
-                    if (++FetchAttempts > 1) {
-                        Y_Fail("TLoader needs multiple fetches in " << Stage << " stage");
-                    }
                     return { fetch };
                 }
 
-                FetchAttempts = 0;
                 Stage = EStage(ui8(Stage) + 1);
             }
 
@@ -169,7 +165,6 @@ namespace NTable {
         const TVector<TString> Deltas;
         const TEpoch Epoch;
         EStage Stage = EStage::Meta;
-        ui8 FetchAttempts = 0;
         bool Rooted = false; /* Has full topology metablob */
         TPageId SchemeId = Max<TPageId>();
         TPageId IndexId = Max<TPageId>();
diff --git a/ydb/core/tablet_flat/ut/ut_slice_loader.cpp b/ydb/core/tablet_flat/ut/ut_slice_loader.cpp
index 3b8e3020fb..10f5d3406e 100644
--- a/ydb/core/tablet_flat/ut/ut_slice_loader.cpp
+++ b/ydb/core/tablet_flat/ut/ut_slice_loader.cpp
@@ -103,7 +103,7 @@ namespace {
     };
 
     struct TCheckResult {
-        size_t Pages;
+        size_t Pages = 0;
         TIntrusivePtr<TSlices> Run;
     };
 
@@ -131,28 +131,24 @@ namespace {
         TKeysEnv env(part.Get(), new TCache(pageCollection));
         TKeysLoader loader(part.Get(), &env);
 
-        if (result.Run = loader.Do(screen)) {
-            env.Check(false); /* On success there shouldn't be left loads */
-            result.Pages = 0;
-        } else  if (auto fetch = env.GetFetches()) {
-            UNIT_ASSERT_C(fetch->PageCollection.Get() == pageCollection.Get(),
-                "TLoader wants to fetch from an unexpected pageCollection");
-            UNIT_ASSERT_C(fetch->Pages, "TLoader wants a fetch, but there are no pages");
-            result.Pages = fetch->Pages.size();
+        while (!(result.Run = loader.Do(screen))) {
+            if (auto fetch = env.GetFetches()) {
+                UNIT_ASSERT_C(fetch->PageCollection.Get() == pageCollection.Get(),
+                    "TLoader wants to fetch from an unexpected pageCollection");
+                UNIT_ASSERT_C(fetch->Pages, "TLoader wants a fetch, but there are no pages");
+                result.Pages += fetch->Pages.size();
 
-            for (auto pageId : fetch->Pages) {
-                auto* page = part->Store->GetPage(0, pageId);
-                UNIT_ASSERT_C(page, "TLoader wants a missing page " << pageId);
+                for (auto pageId : fetch->Pages) {
+                    auto* page = part->Store->GetPage(0, pageId);
+                    UNIT_ASSERT_C(page, "TLoader wants a missing page " << pageId);
 
-                env.Save(fetch->Cookie, { pageId, TSharedPageRef::MakePrivate(*page) });
+                    env.Save(fetch->Cookie, { pageId, TSharedPageRef::MakePrivate(*page) });
+                }
+            } else {
+                UNIT_ASSERT_C(false, "TKeysLoader was stalled");
             }
-
-            result.Run = loader.Do(screen);
-            UNIT_ASSERT_C(result.Run, "TKeysLoader wants to do unexpected fetches");
-            env.Check(false); /* On success there shouldn't be left loads */
-        } else {
-            UNIT_ASSERT_C(false, "TKeysLoader was stalled");
         }
+        env.Check(false); /* On success there shouldn't be left loads */
 
         const auto scrSize = screen ? screen->Size() : 1;
 
@@ -170,10 +166,28 @@ Y_UNIT_TEST_SUITE(TPartSliceLoader) {
 
     Y_UNIT_TEST(RestoreMissingSlice) {
         auto result = RunLoaderTest(Part0(), nullptr);
-        UNIT_ASSERT_C(result.Pages == 0,
+        UNIT_ASSERT_C(result.Pages == 1, // index page
             "Restoring slice bounds needed " << result.Pages << " extra pages");
     }
 
+    Y_UNIT_TEST(RestoreOneSlice) {
+        for (int startOff = 0; startOff < 5; startOff++) {
+            for (int endOff = -5; endOff < 5; endOff++) {
+                TVector<TScreen::THole> holes;
+                holes.emplace_back(Part0()->Index->Begin()->GetRowId() + startOff, Part0()->Index.GetEndRowId() + endOff);
+                TIntrusiveConstPtr<TScreen> screen = new TScreen(std::move(holes));
+                auto result = RunLoaderTest(Part0(), screen);
+
+                size_t expected = 1; // index page
+                if (startOff != 0) expected++;
+                if (endOff < -1) expected++;
+                UNIT_ASSERT_C(result.Pages == expected, // index page
+                    "Restoring slice [" << startOff << ", " << Part0()->Index.GetEndRowId() + endOff << "] bounds needed " << 
+                        result.Pages << " extra pages but expected " << expected);
+            }
+        }
+    }
+
     Y_UNIT_TEST(RestoreMissingSliceFullScreen) {
         TIntrusiveConstPtr<TScreen> screen;
         {
@@ -194,7 +208,7 @@ Y_UNIT_TEST_SUITE(TPartSliceLoader) {
             screen = new TScreen(std::move(holes));
         }
         auto result = RunLoaderTest(Part0(), screen);
-        UNIT_ASSERT_C(result.Pages == 0,
+        UNIT_ASSERT_C(result.Pages == 1, // index page
             "Restoring slice bounds needed " << result.Pages << " extra pages");
     }
 
@@ -219,7 +233,7 @@ Y_UNIT_TEST_SUITE(TPartSliceLoader) {
             screen = new TScreen(std::move(holes));
         }
         auto result = RunLoaderTest(Part0(), screen);
-        UNIT_ASSERT_C(result.Pages == 0,
+        UNIT_ASSERT_C(result.Pages == 1, // index page
             "Restoring slice bounds needed " << result.Pages << " extra pages");
     }
 
@@ -247,7 +261,7 @@ Y_UNIT_TEST_SUITE(TPartSliceLoader) {
             screen = new TScreen(std::move(holes));
         }
         auto result = RunLoaderTest(Part0(), screen);
-        UNIT_ASSERT_C(result.Pages == screen->Size(),
+        UNIT_ASSERT_C(result.Pages == screen->Size() + 1,  // with index page
             "Restoring slice bounds needed " << result.Pages <<
             " extra pages, expected " << screen->Size());
     }