(Reverse) Precharge only foolproof needed pages

2 files changed, 243 insertions, 65 deletions

diff --git a/ydb/core/tablet_flat/flat_part_charge.h b/ydb/core/tablet_flat/flat_part_charge.h
index 9f721e9d5f..08f2b0b338 100644
--- a/ydb/core/tablet_flat/flat_part_charge.h
+++ b/ydb/core/tablet_flat/flat_part_charge.h
@@ -151,7 +151,7 @@ namespace NTable {
                         // Unfortunately key <= key2 might be at the end of the previous slice
                         TRowId lastRow = pos->Slice.EndRowId() - 1;
                         // Precharge the last row on the previous slice
-                        TCharge(env, *pos->Part, tags, includeHistory).DoReverse(lastRow, lastRow, items, bytes);
+                        TCharge(env, *pos->Part, tags, includeHistory).DoReverse(lastRow, lastRow, keyDefaults, items, bytes);
                     }
 
                     break;
@@ -256,8 +256,8 @@ namespace NTable {
          *
          * Important caveat: assumes iteration won't touch any row > row2
          */
-        bool DoReverse(const TRowId row1, const TRowId row2, ui64 itemsLimit,
-                ui64 bytesLimit) const noexcept
+        bool DoReverse(const TRowId row1, const TRowId row2, 
+                const TKeyCellDefaults &keyDefaults, ui64 itemsLimit, ui64 bytesLimit) const noexcept
         {
             auto startRow = row1;
             auto endRow = row2;
@@ -280,7 +280,7 @@ namespace NTable {
                 endRow = Max(endRow, last->GetRowId());
             }
 
-            return DoPrechargeReverse(first, last, startRow, endRow, true, itemsLimit, bytesLimit);
+            return DoPrechargeReverse(TCells{}, TCells{}, TIter{}, TIter{}, first, last, startRow, endRow, keyDefaults, itemsLimit, bytesLimit);
         }
 
         /**
@@ -362,7 +362,6 @@ namespace NTable {
         {
             auto startRow = row1;
             auto endRow = row2;
-            bool firstPageLimits = true;
             bool overshot = !key2; // +inf is always on the next slice
 
             // First page to precharge (contains row1)
@@ -383,25 +382,28 @@ namespace NTable {
                 endRow = Max(endRow, last->GetRowId());
             }
 
+            TIter key1Page;
             if (key1) {
                 // First page to precharge (may contain key <= key1)
-                auto keyPage = Index.LookupKeyReverse(key1, Scheme.Groups[0], ESeek::Lower, &keyDefaults);
-                if (!keyPage || keyPage < last) {
+                key1Page = Index.LookupKeyReverse(key1, Scheme.Groups[0], ESeek::Lower, &keyDefaults);
+                if (!key1Page || key1Page < last) {
                     return { true, true }; // first key is outside of bounds
                 }
-                if (first >= keyPage) {
-                    first = keyPage; // use the minimum
+                if (first >= key1Page) {
+                    first = key1Page; // use the minimum
                     startRow = Min(startRow, Index.GetLastRowId(first));
-                    firstPageLimits = false; // might find key1 after row1, don't count those rows
+                } else {
+                    key1Page = {};
                 }
             }
 
+            TIter key2Page;
             if (key2) {
                 // Last page to precharge (may contain key <= key2)
-                auto keyPage = Index.LookupKeyReverse(key2, Scheme.Groups[0], ESeek::Lower, &keyDefaults);
-                if (keyPage && keyPage >= last) {
-                    if (keyPage <= first) {
-                        last = keyPage; // precharge up to keyPage
+                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
                     }
@@ -411,7 +413,7 @@ namespace NTable {
                 }
             }
 
-            bool ready = DoPrechargeReverse(first, last, startRow, endRow, firstPageLimits, itemsLimit, bytesLimit);
+            bool ready = DoPrechargeReverse(key1, key2, key1Page, key2Page, first, last, startRow, endRow, keyDefaults, itemsLimit, bytesLimit);
 
             return { ready, overshot };
         }
@@ -420,11 +422,13 @@ namespace NTable {
         /**
          * Precharges data from first to last page inclusive
          *
-         * Assumes worst case first page will start iteration from startRowId
+         * Precharges data only in [ @param startRowId, @param endRowId ] range.
          *
-         * Assumes worst case last page will end iteration on endRowId
+         * If keys provided, precharges only foolproof needed pages between them.
          * 
-         * If items limit specified also touches [startRowId + itemsLimit] row.
+         * If items limit specified also touches [@param startRowId + itemsLimit] row.
+         *
+         * If @param key1Page specified, @param first should be the same.
          */
         bool DoPrecharge(const TCells key1, const TCells key2, 
                 const TIter key1Page, const TIter key2Page,
@@ -451,7 +455,9 @@ namespace NTable {
                     auto currentLastRowId = currentExt ? (currentExt->GetRowId() - 1) : Max<TRowId>();
 
                     auto page = Env->TryGetPage(Part, current->GetPageId());
-                    bytes += Part->GetPageSize(current->GetPageId());
+                    if (bytesLimit) {
+                        bytes += Part->GetPageSize(current->GetPageId());
+                    }
                     ready &= bool(page);
 
                     auto prechargeCurrentFirstRowId = Max(currentFirstRowId, startRowId);
@@ -472,7 +478,7 @@ namespace NTable {
                                 if (key2RowId) {
                                     prechargeCurrentLastRowId = Min(prechargeCurrentLastRowId, key2RowId - 1);
                                 } else {
-                                    break;
+                                    prechargeCurrentFirstRowId = Max<TRowId>(); // no precharge
                                 }
                             } else {
                                 prechargeCurrentFirstRowId = Max<TRowId>(); // no precharge
@@ -513,66 +519,96 @@ namespace NTable {
         /**
          * Precharges data from first to last page inclusive in reverse
          *
-         * Assumes worst case first page will start iteration from startRowId
+         * Precharges data only in [ @param endRowId, @param startRowId ] range.
+         *
+         * If keys provided, precharges only foolproof needed pages between them.
+         * 
+         * If items limit specified also touches [@param startRowId + itemsLimit] row.
+         *
+         * If @param key1Page specified, @param first should be the same.
          */
-        bool DoPrechargeReverse(TIter first, TIter last,
-                TRowId startRowId, TRowId endRowId, bool firstPageLimits,
-                ui64 itemsLimit, ui64 bytesLimit) const noexcept
+        bool DoPrechargeReverse(const TCells key1, const TCells key2, 
+                const TIter key1Page, const TIter key2Page,
+                TIter first, TIter last,
+                TRowId startRowId, TRowId endRowId,
+                const TKeyCellDefaults &keyDefaults, ui64 itemsLimit, ui64 bytesLimit) const noexcept
         {
             bool ready = true;
 
             if (first) {
                 Y_VERIFY_DEBUG(first >= last);
+                Y_VERIFY_DEBUG(!key1Page || key1Page == first);
 
                 ui64 items = 0;
                 ui64 bytes = 0;
 
-                TRowId prechargedFirstRowId = startRowId;
+                TRowId prechargedFirstRowId, prechargedLastRowId;
 
-                for (;;) {
-                    auto page = first->GetPageId();
-                    ready &= bool(Env->TryGetPage(Part, page));
+                for (auto current = first;
+                        current >= last && !LimitExceeded(items, itemsLimit) && !LimitExceeded(bytes, bytesLimit);
+                        current--) {
+                    auto currentExt = current + 1;
+                    auto currentFirstRowId = currentExt ? (currentExt->GetRowId() - 1) : Max<TRowId>();
+                    auto currentLastRowId = current->GetRowId();
 
-                    if (Groups) {
-                        TRowId lastRowId = Max(first->GetRowId(), endRowId);
-                        if (Y_LIKELY(lastRowId <= startRowId)) {
-                            if (itemsLimit && firstPageLimits) {
-                                ui64 left = itemsLimit - items + 1;
-                                if (startRowId - lastRowId > left - 1) {
-                                    lastRowId = startRowId - left - 1;
-                                }
-                            }
+                    auto page = Env->TryGetPage(Part, current->GetPageId());
+                    if (bytesLimit) {
+                        bytes += Part->GetPageSize(current->GetPageId());
+                    }
+                    ready &= bool(page);
 
-                            for (auto& g : Groups) {
-                                ready &= DoPrechargeGroupReverse(g, startRowId, lastRowId, bytes);
+                    auto prechargeCurrentFirstRowId = Min(currentFirstRowId, startRowId);
+                    auto prechargeCurrentLastRowId = Max(currentLastRowId, endRowId);
+
+                    if (key1Page && key1Page == current) {
+                        if (page) {
+                            auto key1RowId = LookupRowIdReverse(key1, page, ESeek::Lower, keyDefaults);
+                            prechargeCurrentFirstRowId = Min(prechargeCurrentFirstRowId, key1RowId);
+                        } else {
+                            prechargeCurrentLastRowId = Max<TRowId>(); // no precharge
+                        }
+                    }
+                    if (key2Page && key2Page >= current) {
+                        if (key2Page == current) {
+                            if (page) {
+                                auto key2RowId = LookupRowIdReverse(key2, page, ESeek::Upper, keyDefaults);
+                                if (key2RowId != Max<TRowId>()) { // Max<TRowId>() means that upper bound is before current page, so doesn't limit current page
+                                    prechargeCurrentLastRowId = Max(prechargeCurrentLastRowId, key2RowId + 1);
+                                }
+                            } else {
+                                prechargeCurrentLastRowId = Max<TRowId>(); // no precharge
                             }
+                        } else {
+                            prechargeCurrentLastRowId = Max<TRowId>(); // no precharge
                         }
                     }
 
-                    if (first == last || first.Off() == 0) {
-                        break; // this was the last page to precharge
+                    if (itemsLimit && prechargeCurrentFirstRowId >= prechargeCurrentLastRowId) {
+                        ui64 left = itemsLimit - items; // we count only foolprof taken rows, so here we may precharge some extra rows
+                        if (prechargeCurrentFirstRowId - prechargeCurrentLastRowId > left) {
+                            prechargeCurrentLastRowId = prechargeCurrentFirstRowId - left;
+                        }
                     }
 
-                    if (itemsLimit && firstPageLimits && first->GetRowId() <= startRowId) {
-                        items += startRowId + 1 - first->GetRowId();
-                        if (items > itemsLimit)
-                            break;
+                    if (prechargeCurrentFirstRowId >= prechargeCurrentLastRowId) {
+                        if (!items) {
+                            prechargedFirstRowId = prechargeCurrentFirstRowId;
+                        }
+                        prechargedLastRowId = prechargeCurrentLastRowId;
+                        if (Groups) {
+                            for (auto& g : Groups) {
+                                ready &= DoPrechargeGroupReverse(g, prechargeCurrentFirstRowId, prechargeCurrentLastRowId, bytes);
+                            }
+                        }
+                        items += prechargeCurrentFirstRowId - prechargeCurrentLastRowId + 1;
                     }
 
-                    if (bytesLimit) {
-                        bytes += Part->GetPageSize(page);
-                        if (bytes > bytesLimit)
-                            break;
+                    if (current.Off() == 0) {
+                        break;
                     }
-
-                    startRowId = first->GetRowId() - 1;
-                    firstPageLimits = true;
-                    --first;
                 }
 
-                if (HistoryIndex) {
-                    TRowId prechargedLastRowId = Max(first->GetRowId(), endRowId);
-
+                if (items && HistoryIndex) {
                     ready &= DoPrechargeHistory(prechargedFirstRowId, prechargedLastRowId);
                 }
             }
@@ -745,6 +781,17 @@ namespace NTable {
         }
 
     private:
+        TRowId LookupRowIdReverse(const TCells key, const TSharedData* page, ESeek seek, const TKeyCellDefaults &keyDefaults) const noexcept
+        {
+            auto data = TDataPage(page);
+            auto lookup = data.LookupKeyReverse(key, Scheme.Groups[0], seek, &keyDefaults);
+            auto rowId = lookup
+                ? data.BaseRow() + lookup.Off()
+                : Max<TRowId>();
+            return rowId;
+        }
+
+    private:
         bool LimitExceeded(ui64 value, ui64 limit) const noexcept {
             return limit && value > limit;
         }
diff --git a/ydb/core/tablet_flat/ut/ut_charge.cpp b/ydb/core/tablet_flat/ut/ut_charge.cpp
index 6fcaf83145..d93c27e884 100644
--- a/ydb/core/tablet_flat/ut/ut_charge.cpp
+++ b/ydb/core/tablet_flat/ut/ut_charge.cpp
@@ -127,18 +127,24 @@ namespace {
 
         void CheckByKeys(ui32 lower, ui32 upper, ui64 items, const TMap<TGroupId, TArr>& shouldPrecharge) const
         {
-            CheckPrechargeByKeys(lower, upper, items, false, shouldPrecharge);
-            CheckPrechargeByKeys(lower, upper, items, true, shouldPrecharge);
+            CheckPrechargeByKeys(lower, upper, items, false, shouldPrecharge, false);
+            CheckPrechargeByKeys(lower, upper, items, true, shouldPrecharge, false);
             CheckIterByKeys(lower, upper, items ? items : Max<ui32>(), shouldPrecharge);
         }
 
+        void CheckByKeysReverse(ui32 lower, ui32 upper, ui64 items, const TMap<TGroupId, TArr>& shouldPrecharge) const
+        {
+            CheckPrechargeByKeys(lower, upper, items, false, shouldPrecharge, true);
+            CheckPrechargeByKeys(lower, upper, items, true, shouldPrecharge, true);
+        }
+
         void CheckByRows(TPageId row1, TPageId row2, ui64 items, TMap<TGroupId, TArr> shouldPrecharge) const
         {
             CheckPrechargeByRows(row1, row2, items, false, shouldPrecharge);
             CheckPrechargeByRows(row1, row2, items, true, shouldPrecharge);
         }
 
-        void CheckPrechargeByKeys(ui32 lower, ui32 upper, ui64 items, bool fail, const TMap<TGroupId, TArr>& shouldPrecharge) const
+        void CheckPrechargeByKeys(ui32 lower, ui32 upper, ui64 items, bool fail, const TMap<TGroupId, TArr>& shouldPrecharge, bool reverse) const
         {
             Y_VERIFY(lower < Mass.Saved.Size() && upper < Mass.Saved.Size());
 
@@ -159,10 +165,11 @@ namespace {
                 tags.push_back(c.Tag);
             }
 
-            UNIT_ASSERT_VALUES_EQUAL_C(
-                !fail,
-                TCharge::Range(&env, from, to, run, keyDefaults, tags, items, Max<ui64>()),
-                AssertMesage(fail));
+            bool ready = !reverse
+                ? TCharge::Range(&env, from, to, run, keyDefaults, tags, items, Max<ui64>())
+                : TCharge::RangeReverse(&env, from, to, run, keyDefaults, tags, items, Max<ui64>());
+
+            UNIT_ASSERT_VALUES_EQUAL_C(!fail || env.Touched.empty(), ready, AssertMesage(fail));
 
             AssertEqual(env.Touched, shouldPrecharge, fail ? TPageIdFlags::IfFail : TPageIdFlags::IfNoFail);
         }
@@ -316,7 +323,7 @@ Y_UNIT_TEST_SUITE(Charge) {
         { /*_ 1xx: A set of some edge and basic cases */
 
             me.To(100).CheckByKeys(33, 34, 0, { 8 });     // keys on the last page
-            me.To(101).CheckByKeys(0, 0, 0, { 0, 1_f });    // keys before the part
+            me.To(101).CheckByKeys(0, 0, 0, { 0, 1_I });    // keys before the part
             me.To(102).CheckByKeys(3, 4, 0, { 0, 1 });
             me.To(103).CheckByKeys(3, 5, 0, { 0, 1, 2_I });
             me.To(104).CheckByKeys(4, 8, 0, { 0, 1, 2 });
@@ -389,7 +396,7 @@ Y_UNIT_TEST_SUITE(Charge) {
 
         /*_ 1xx: Play with first page */ {
             me.To(100).CheckByKeys(0, 0, 0, TMap<TGroupId, TArr>{
-                {TGroupId{0}, {0, 1_f}},
+                {TGroupId{0}, {0, 1_I}},
                 {TGroupId{1}, {}},
                 {TGroupId{2}, {}}
             });
@@ -705,6 +712,130 @@ Y_UNIT_TEST_SUITE(Charge) {
         }
     }
 
+    Y_UNIT_TEST(ByKeysReverse)
+    {
+        TModel me;
+
+        /*
+        
+        keys by pages:
+        group0 = |1 2 3|5 6 7|9 10 11|13 14 15|..
+        group1 = |1 2|3 5|6 7|9 10|11 13|14 15|..
+        group3 = |1|2|3|5|6|7|9|10|11|13|14|15|..
+        
+        */
+
+        me.To(100).CheckByKeysReverse(15, 15, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3}},
+            {TGroupId{2}, {11_g}}
+        });
+
+        me.To(101).CheckByKeysReverse(15, 5, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5_g, 4_g, 3_g}}
+        });
+
+        me.To(102).CheckByKeysReverse(15, 4, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1, 0}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5, 4, 3}}
+        });
+
+        me.To(103).CheckByKeysReverse(15, 3, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1, 0}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5, 4, 3, 2_g}}
+        });
+
+        me.To(104).CheckByKeysReverse(15, 1, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1, 0}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5, 4, 3, 2_g, 1_g, 0_g}}
+        });
+
+        me.To(105).CheckByKeysReverse(15, 0, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1, 0}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5, 4, 3, 2, 1, 0}}
+        });
+
+        me.To(106).CheckByKeysReverse(17, 17, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {4}},
+            {TGroupId{2}, {12_g}}
+        });
+        
+        me.To(107).CheckByKeysReverse(16, 16, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3}},
+            {TGroupId{2}, {}}
+        });
+
+        me.To(108).CheckByKeysReverse(35, 35, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {8}},
+            {TGroupId{2}, {26_g}}
+        });
+
+        me.To(109).CheckByKeysReverse(35, 33, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {8}},
+            {TGroupId{2}, {26_g, 25_g, 24_g}}
+        });
+
+        me.To(110).CheckByKeysReverse(35, 32, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {8, 7}},
+            {TGroupId{2}, {26_g, 25_g, 24_g}}
+        });
+
+        me.To(111).CheckByKeysReverse(4, 1, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {0}},
+            {TGroupId{2}, {2_g, 1_g, 0_g}}
+        });
+
+        me.To(112).CheckByKeysReverse(1, 1, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {0}},
+            {TGroupId{2}, {0_g}}
+        });
+
+        me.To(113).CheckByKeysReverse(1, 0, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {0}},
+            {TGroupId{2}, {0_g}}
+        });
+
+        me.To(114).CheckByKeysReverse(0, 0, 0, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {}},
+            {TGroupId{2}, {}}
+        });
+
+        me.To(200).CheckByKeysReverse(15, 3, 6, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1, 0_f}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5, 4_f, 3_f}} // here we touh extra pages, but it's fine
+        });
+
+        me.To(201).CheckByKeysReverse(15, 3, 5, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1_f}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6, 5_f, 4_f, 3_f}} // here we touh extra pages, but it's fine
+        });
+
+        me.To(202).CheckByKeysReverse(15, 5, 5, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1_f}},
+            {TGroupId{2}, {11_g, 10_g, 9_g, 8, 7, 6}}
+        });
+
+        me.To(203).CheckByKeysReverse(13, 3, 4, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1}},
+            {TGroupId{2}, {9_g, 8, 7, 6, 5, 4_f}} // here we touh extra pages, but it's fine
+        });
+
+        me.To(204).CheckByKeysReverse(13, 3, 3, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2, 1_f}},
+            {TGroupId{2}, {9_g, 8, 7, 6, 5_f}} // here we touh extra pages, but it's fine
+        });
+
+        me.To(205).CheckByKeysReverse(13, 3, 2, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2}},
+            {TGroupId{2}, {9_g, 8, 7, 6_f}} // here we touh extra pages, but it's fine
+        });
+
+        me.To(206).CheckByKeysReverse(13, 3, 1, TMap<TGroupId, TArr>{
+            {TGroupId{0}, {3, 2}},
+            {TGroupId{2}, {9_g, 8, 7_f}} // here we touh extra pages, but it's fine
+        });
+    }
+
     Y_UNIT_TEST(ByRows)
     {
         TModel me;