KIKIMR-18845 Cut keys Precharge

5 files changed, 157 insertions, 38 deletions

diff --git a/ydb/core/tablet_flat/flat_part_charge.h b/ydb/core/tablet_flat/flat_part_charge.h
index 68b81e60843..05b42219300 100644
--- a/ydb/core/tablet_flat/flat_part_charge.h
+++ b/ydb/core/tablet_flat/flat_part_charge.h
@@ -337,14 +337,10 @@ namespace NTable {
                 key2Page = Index.LookupKey(key2, Scheme.Groups[0], ESeek::Lower, &keyDefaults);
                 auto key2PageExt = key2Page + 1;
                 if (key2PageExt && key2PageExt <= last) {
-                    if (key2PageExt >= firstExt) {
-                        last = key2PageExt; // precharge up to key2PageExt
-                    } else {
-                        last = firstExt; // precharge up to firstExt
-                    }
+                    last = Max(key2PageExt, firstExt);
                     endRow = Min(endRow, last->GetRowId()); // may load the first row of key2PageExt
                 } else {
-                    overshot = true; // may find key > key2 on row > row2
+                    overshot = true; // may find first key > key2 on row > row2
                 }
             }
 
@@ -375,6 +371,9 @@ namespace NTable {
                 }
             }
 
+            // First extra page to precharge (when key placement is uncertain)
+            auto firstExt = first;
+
             // Last page to precharge (contains row2)
             auto last = Index.LookupRow(row2, first);
             if (Y_UNLIKELY(last > first)) {
@@ -391,7 +390,12 @@ namespace NTable {
                 }
                 if (first >= key1Page) {
                     first = key1Page; // use the minimum
+                    firstExt = key1Page - 1; // first key <= key1 might be on the next page
                     startRow = Min(startRow, Index.GetLastRowId(first));
+                    if (key1Page.Off() == 0 || last > firstExt) {
+                        firstExt = last; // never precharge past the last page
+                        overshot = true; // may have to touch the next slice
+                    }
                 } else {
                     key1Page = {};
                 }
@@ -400,16 +404,14 @@ namespace NTable {
             TIter key2Page;
             if (key2) {
                 // Last page to precharge (may contain key <= key2)
+                // We actually use the next page since lookup is not exact
                 key2Page = Index.LookupKeyReverse(key2, Scheme.Groups[0], ESeek::Lower, &keyDefaults);
-                if (key2Page && key2Page >= last) {
-                    if (key2Page <= first) {
-                        last = key2Page; // precharge up to keyPage
-                    } else {
-                        last = first; // precharge up to first
-                    }
-                    endRow = Max(endRow, last->GetRowId());
+                auto key2PageExt = key2Page - 1;
+                if (key2Page && key2Page.Off() != 0 && key2PageExt >= last) {
+                    last = Min(key2PageExt, firstExt);
+                    endRow = Max(endRow, Index.GetLastRowId(last)); // may load the last row of key2PageExt
                 } else {
-                    overshot = true; // may find key <= key2 on row < row2
+                    overshot = true; // may find first key < key2 on row < row2
                 }
             }
 
diff --git a/ydb/core/tablet_flat/test/libs/table/test_iter.h b/ydb/core/tablet_flat/test/libs/table/test_iter.h
index 47bd9fdb1fa..2f842d97d86 100644
--- a/ydb/core/tablet_flat/test/libs/table/test_iter.h
+++ b/ydb/core/tablet_flat/test/libs/table/test_iter.h
@@ -71,6 +71,15 @@ namespace NTest {
             return Displace<IPages>(env), *this;
         }
 
+        template<typename TEnv>
+        TEnv* GetEnv() {
+            auto *casted = dynamic_cast<TEnv*>(Env.Get());
+
+            Y_VERIFY(!Env || casted, "Cannot cast IPages to given env");
+
+            return casted;
+        }
+
         template<typename ...TArgs>
         inline TChecker& IsN(TArgs&&...args)
         {
@@ -326,6 +335,11 @@ namespace NTest {
             return *this;
         }
 
+        EReady GetReady() const noexcept
+        {
+            return Ready;
+        }
+
     private:
         const ui64 Retries = 1;
         const ui64 Erased = true;
diff --git a/ydb/core/tablet_flat/test/libs/table/wrap_part.h b/ydb/core/tablet_flat/test/libs/table/wrap_part.h
index 5ee6ec1cd3a..845f8cf4ebd 100644
--- a/ydb/core/tablet_flat/test/libs/table/wrap_part.h
+++ b/ydb/core/tablet_flat/test/libs/table/wrap_part.h
@@ -148,7 +148,7 @@ namespace NTest {
 
             if (StopKey) {
                 auto cmp = CompareTypedCellVectors(key.Cells().data(), StopKey.data(), Scheme->Keys->Types.data(), StopKey.size());
-                if (cmp > 0) {
+                if (Direction == EDirection::Forward && cmp > 0 || Direction == EDirection::Reverse && cmp < 0) {
                    return EReady::Gone;
                 }
             }
diff --git a/ydb/core/tablet_flat/ut/ut_charge.cpp b/ydb/core/tablet_flat/ut/ut_charge.cpp
index ca8da5f1031..79c2f8116e3 100644
--- a/ydb/core/tablet_flat/ut/ut_charge.cpp
+++ b/ydb/core/tablet_flat/ut/ut_charge.cpp
@@ -145,6 +145,7 @@ namespace {
         {
             CheckPrechargeByKeys(lower, upper, items, false, shouldPrecharge, true);
             CheckPrechargeByKeys(lower, upper, items, true, shouldPrecharge, true);
+            CheckIterByKeysReverse(lower, upper, items ? items : Max<ui32>(), shouldPrecharge);
         }
 
         void CheckByRows(TPageId row1, TPageId row2, ui64 items, TMap<TGroupId, TArr> shouldPrecharge) const
@@ -245,6 +246,41 @@ namespace {
             AssertEqual(env->Touched, precharged, TPageIdFlags::IfIter);
         }
 
+        void CheckIterByKeysReverse(ui32 lower, ui32 upper, ui64 items, const TMap<TGroupId, TArr>& precharged) const
+        {
+            Y_VERIFY(lower < Mass.Saved.Size() && upper < Mass.Saved.Size());
+
+            NTest::TCheckReverseIt wrap(Eggs, { new TTouchEnv(false) });
+
+            wrap.To(CurrentStep());
+            wrap.StopAfter(Tool.KeyCells(Mass.Saved[upper]));
+
+            auto seek = true;
+            for (ui32 key = lower; items-- + 1 > 0; key--) {
+                if (key % 4 == 0) {
+                    --key;
+                }
+
+                if (seek) {
+                    wrap.Seek(Mass.Saved[lower], ESeek::Lower);
+                    seek = false;
+                } else {
+                    wrap.Next().Is(key < upper || key == (ui32)-1 ? EReady::Gone : EReady::Data);
+                }
+
+                // forcebly touch the next stop element that is greater than upper
+                // because instead of having |1 2 3| and stopping as soon as we see 2
+                // we may have |1*2 2*2 3*2| = |2 4 6| and be requested with upper = 2 (not 4)
+                if (key < upper || key == (ui32)-1) {
+                    break;
+                }
+            }
+
+            auto env = wrap.Displace<TTouchEnv>(nullptr);
+
+            AssertEqual(env->Touched, precharged, TPageIdFlags::IfIter);
+        }
+
         const NTest::TMass Mass;
         const NTest::TPartEggs Eggs;
         const NTest::TRowTool Tool;
@@ -747,12 +783,12 @@ Y_UNIT_TEST_SUITE(Charge) {
         */
 
         me.To(100).CheckByKeysReverse(15, 15, 0, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {3}},
+            {TGroupId{0}, {3, 2_I}},
             {TGroupId{2}, {11_g}}
         });
 
         me.To(101).CheckByKeysReverse(15, 5, 0, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {3, 2, 1}},
+            {TGroupId{0}, {3, 2, 1, 0}},
             {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5_g, 4_g, 3_g}}
         });
 
@@ -777,27 +813,27 @@ Y_UNIT_TEST_SUITE(Charge) {
         });
 
         me.To(106).CheckByKeysReverse(17, 17, 0, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {4}},
+            {TGroupId{0}, {4, 3}},
             {TGroupId{2}, {12_g}}
         });
         
         me.To(107).CheckByKeysReverse(16, 16, 0, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {3}},
+            {TGroupId{0}, {3, 2_I}},
             {TGroupId{2}, {}}
         });
 
         me.To(108).CheckByKeysReverse(35, 35, 0, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {8}},
+            {TGroupId{0}, {8, 7_I}},
             {TGroupId{2}, {26_g}}
         });
 
         me.To(109).CheckByKeysReverse(35, 33, 0, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {8}},
+            {TGroupId{0}, {8, 7}},
             {TGroupId{2}, {26_g, 25_g, 24_g}}
         });
 
         me.To(110).CheckByKeysReverse(35, 32, 0, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {8, 7}},
+            {TGroupId{0}, {8, 7, 6_I}},
             {TGroupId{2}, {26_g, 25_g, 24_g}}
         });
 
@@ -832,7 +868,7 @@ Y_UNIT_TEST_SUITE(Charge) {
         });
 
         me.To(202).CheckByKeysReverse(15, 5, 5, TMap<TGroupId, TArr>{
-            {TGroupId{0}, {3, 2, 1_f}},
+            {TGroupId{0}, {3, 2, 1_f, 0_f}},
             {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6}}
         });
 
diff --git a/ydb/core/tablet_flat/ut/ut_part.cpp b/ydb/core/tablet_flat/ut/ut_part.cpp
index d6fbd3268e9..3bbf4705a6d 100644
--- a/ydb/core/tablet_flat/ut/ut_part.cpp
+++ b/ydb/core/tablet_flat/ut/ut_part.cpp
@@ -1,3 +1,4 @@
+#include "flat_part_charge.h"
 #include <ydb/core/tablet_flat/flat_part_dump.h>
 #include <ydb/core/tablet_flat/test/libs/rows/cook.h>
 #include <ydb/core/tablet_flat/test/libs/rows/layout.h>
@@ -13,6 +14,8 @@ namespace NKikimr {
 namespace NTable {
 
 namespace {
+    using namespace NTest;
+
     NPage::TConf PageConf(size_t groups = 1) noexcept
     {
         NPage::TConf conf{ true, 2 * 1024 };
@@ -52,6 +55,54 @@ namespace {
         static const auto eggs1 = NTest::TPartCook::Make(mass1, PageConf(mass1.Model->Scheme->Families.size()));
         return eggs1;
     }
+
+    struct TTouchEnv : public NTest::TTestEnv {
+        const TSharedData* TryGetPage(const TPart *part, TPageId id, TGroupId groupId) override
+        {
+            if (PrechargePhase) {
+                Precharged[groupId].insert(id);
+                return NTest::TTestEnv::TryGetPage(part, id, groupId);
+            } else {
+                Y_VERIFY_S(Precharged[groupId].count(id), "Requested page " << id << " should be precharged");
+                return NTest::TTestEnv::TryGetPage(part, id, groupId);
+            }
+        }
+
+        bool PrechargePhase = true;
+        TMap<TGroupId, TSet<TPageId>> Precharged;
+    };
+
+    template<EDirection Direction>
+    void Precharge(TChecker<TWrapPartImpl<Direction>, TPartEggs>& wrap, const TRow &row) {
+        auto env = wrap.template GetEnv<TTouchEnv>();
+        env->Precharged.clear();
+        env->PrechargePhase = true;
+
+        const auto part = (*wrap).Eggs.Lone();
+        const auto &keyDefaults = (*wrap).Eggs.Scheme->Keys;
+
+        TRun run(*keyDefaults);
+        for (auto& slice : *part->Slices) {
+            run.Insert(part, slice);
+        }
+
+        TRowTool tool(*(*wrap).Eggs.Scheme);
+        const auto from = tool.KeyCells(row);
+        const auto to = tool.KeyCells(row);
+
+        auto tags = TVector<TTag>();
+        for (auto c : (*wrap).Eggs.Scheme->Cols) {
+            tags.push_back(c.Tag);
+        }
+
+        if constexpr (Direction == EDirection::Forward) {
+            TCharge::Range(env, from, to, run, *keyDefaults, tags, 0, 0, true);
+        } else {
+            TCharge::RangeReverse(env, from, to, run, *keyDefaults, tags, 0, 0, true);
+        }
+
+        env->PrechargePhase = false;
+    }
 }
 
 Y_UNIT_TEST_SUITE(TPart) {
@@ -647,8 +698,8 @@ Y_UNIT_TEST_SUITE(TPart) {
             fullCookR.Add(*TSchemedCookRow(*lay).Col(r.first, r.second));
         }
 
-        TCheckIt cutWrap(cutCook.Finish(), { }), fullWrap(fullCook.Finish(), { });
-        TCheckReverseIt cutWrapR(cutCookR.Finish(), { }), fullWrapR(fullCookR.Finish(), { });
+        TCheckIt cutWrap(cutCook.Finish(), { new TTouchEnv() }), fullWrap(fullCook.Finish(), { new TTouchEnv() });
+        TCheckReverseIt cutWrapR(cutCookR.Finish(), { new TTouchEnv() }), fullWrapR(fullCookR.Finish(), { new TTouchEnv() });
 
         const auto cutPart = (*cutWrap).Eggs.Lone();
         const auto fullPart = (*fullWrap).Eggs.Lone();
@@ -673,7 +724,7 @@ Y_UNIT_TEST_SUITE(TPart) {
         }
 
         for (size_t rowId = 0; rowId < fullRows.size(); rowId++)
-        for (auto seekMode : {ESeek::Exact, ESeek::Lower, ESeek::Upper})
+        for (auto seekMode : {ESeek::Exact, ESeek::Lower, ESeek::Upper })
         for (auto transformMode : {ESeek::Exact, ESeek::Lower, ESeek::Upper}) {
             auto str = fullRows[rowId].second;
 
@@ -690,13 +741,21 @@ Y_UNIT_TEST_SUITE(TPart) {
                 break;
             }
 
-            cutWrap.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
-            fullWrap.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
-            UNIT_ASSERT_VALUES_EQUAL(cutWrap->GetRowId(), fullWrap->GetRowId());
+            auto seekRow = *TSchemedCookRow(*lay).Col(fullRows[rowId].first, str);
 
-            cutWrapR.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
-            fullWrapR.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
+            Precharge(cutWrap, seekRow);
+            Precharge(fullWrap, seekRow);
+            cutWrap.Seek(seekRow, seekMode);
+            fullWrap.Seek(seekRow, seekMode);
+            UNIT_ASSERT_VALUES_EQUAL(cutWrap.GetReady(), fullWrap.GetReady());
             UNIT_ASSERT_VALUES_EQUAL(cutWrap->GetRowId(), fullWrap->GetRowId());
+
+            Precharge(cutWrapR, seekRow);
+            Precharge(fullWrapR, seekRow);
+            cutWrapR.Seek(seekRow, seekMode);
+            fullWrapR.Seek(seekRow, seekMode);
+            UNIT_ASSERT_VALUES_EQUAL(cutWrapR.GetReady(), fullWrapR.GetReady());
+            UNIT_ASSERT_VALUES_EQUAL(cutWrapR->GetRowId(), fullWrapR->GetRowId());
         }
     }
 
@@ -735,8 +794,8 @@ Y_UNIT_TEST_SUITE(TPart) {
             fullCookR.Add(*TSchemedCookRow(*lay).Col(r.first, r.second));
         }
 
-        TCheckIt cutWrap(cutCook.Finish(), { }), fullWrap(fullCook.Finish(), { });
-        TCheckReverseIt cutWrapR(cutCookR.Finish(), { }), fullWrapR(fullCookR.Finish(), { });
+        TCheckIt cutWrap(cutCook.Finish(), { new TTouchEnv() }), fullWrap(fullCook.Finish(), { new TTouchEnv() });
+        TCheckReverseIt cutWrapR(cutCookR.Finish(), { new TTouchEnv() }), fullWrapR(fullCookR.Finish(), { new TTouchEnv() });
 
         const auto cutPart = (*cutWrap).Eggs.Lone();
         const auto fullPart = (*fullWrap).Eggs.Lone();
@@ -771,13 +830,21 @@ Y_UNIT_TEST_SUITE(TPart) {
                 break;
             }
 
-            cutWrap.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
-            fullWrap.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
-            UNIT_ASSERT_VALUES_EQUAL(cutWrap->GetRowId(), fullWrap->GetRowId());
+            auto seekRow = *TSchemedCookRow(*lay).Col(fullRows[rowId].first, str);
 
-            cutWrapR.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
-            fullWrapR.Seek(*TSchemedCookRow(*lay).Col(fullRows[rowId].first, str), seekMode);
+            Precharge(cutWrap, seekRow);
+            Precharge(fullWrap, seekRow);
+            cutWrap.Seek(seekRow, seekMode);
+            fullWrap.Seek(seekRow, seekMode);
+            UNIT_ASSERT_VALUES_EQUAL(cutWrap.GetReady(), fullWrap.GetReady());
             UNIT_ASSERT_VALUES_EQUAL(cutWrap->GetRowId(), fullWrap->GetRowId());
+
+            Precharge(cutWrapR, seekRow);
+            Precharge(fullWrapR, seekRow);
+            cutWrapR.Seek(seekRow, seekMode);
+            fullWrapR.Seek(seekRow, seekMode);
+            UNIT_ASSERT_VALUES_EQUAL(cutWrapR.GetReady(), fullWrapR.GetReady());
+            UNIT_ASSERT_VALUES_EQUAL(cutWrapR->GetRowId(), fullWrapR->GetRowId());
         }
     }