KIKIMR-18845 Cut keys test Slices

3 files changed, 135 insertions, 5 deletions

diff --git a/ydb/core/tablet_flat/flat_page_index.h b/ydb/core/tablet_flat/flat_page_index.h
index 362d63a298..1151120cd8 100644
--- a/ydb/core/tablet_flat/flat_page_index.h
+++ b/ydb/core/tablet_flat/flat_page_index.h
@@ -270,6 +270,15 @@ namespace NPage {
             return LastKey;
         }
 
+        const TRecord* At(TRecIdx index) const noexcept
+        {
+            Y_VERIFY(index <= Page.Count);
+            auto it = Page.Begin() + index;
+            return it
+                ? it.GetRecord()
+                : GetLastKeyRecord();
+        }
+
         TRowId GetEndRowId() const noexcept
         {
             return EndRowId;
diff --git a/ydb/core/tablet_flat/flat_part_dump.cpp b/ydb/core/tablet_flat/flat_part_dump.cpp
index 7c923e2954..4565221620 100644
--- a/ydb/core/tablet_flat/flat_part_dump.cpp
+++ b/ydb/core/tablet_flat/flat_part_dump.cpp
@@ -121,15 +121,15 @@ namespace {
                 break;
             }
 
-            auto iter = i == items - 1 ? part.Index.GetLastKeyRecord() : (part.Index->Begin() + i).GetRecord();
+            auto record = part.Index.At(i);
             for (const auto &info: part.Scheme->Groups[0].ColsKeyIdx)
-                Key.push_back(iter->Cell(info));
+                Key.push_back(record->Cell(info));
 
             Out
-                << " | " << (Printf(Out, " %4u", iter->GetPageId()), " ")
-                << (Printf(Out, " %6lu", iter->GetRowId()), " ");
+                << " | " << (Printf(Out, " %4u", record->GetPageId()), " ")
+                << (Printf(Out, " %6lu", record->GetRowId()), " ");
 
-            if (auto *page = Env->TryGetPage(&part, iter->GetPageId())) {
+            if (auto *page = Env->TryGetPage(&part, record->GetPageId())) {
                 Printf(Out, " %6zub  ", page->size());
             } else {
                 Out << "~none~  ";
diff --git a/ydb/core/tablet_flat/ut/ut_part.cpp b/ydb/core/tablet_flat/ut/ut_part.cpp
index 3bbf4705a6..c9ad7e7167 100644
--- a/ydb/core/tablet_flat/ut/ut_part.cpp
+++ b/ydb/core/tablet_flat/ut/ut_part.cpp
@@ -848,6 +848,127 @@ Y_UNIT_TEST_SUITE(TPart) {
         }
     }
 
+    Y_UNIT_TEST(CutKeys_SeekSlices)
+    {
+        TLayoutCook lay;
+
+        lay
+            .Col(0, 0,  NScheme::NTypeIds::Uint32)
+            .Col(0, 1,  NScheme::NTypeIds::String)
+            .Key({0, 1});
+
+        TVector<std::pair<ui32, TString>> fullRows = {
+            {1, "aaa"}, // -> (1, "aaa")
+            {1, "aba"}, // -> (1, "ab")
+            {1, "aca"}, // -> (1, "ac")
+            {1, "baa"}, // -> (1, "b")
+            {1, "bba"}, // -> (1, "bb")
+            {2, "aaa"}, // -> (2, null)
+            {2, "aba"}, // -> (2, "ab")
+            {2, "aca"}, // -> (2, "ac")
+            {2, "baa"}, // -> (2, "b")
+            {2, "bba"}, // -> (2, "bba")
+        };
+
+        NPage::TConf cutConf{ true, 8192 }, fullConf{ true, 8192 };
+        cutConf.CutIndexKeys = true;
+        fullConf.CutIndexKeys = false;
+        cutConf.Group(0).PageRows = fullConf.Group(0).PageRows = 1;
+
+        TPartCook cutCookTmp(lay, cutConf), cutCook(lay, cutConf), cutCookR(lay, cutConf), fullCook(lay, fullConf), fullCookR(lay, fullConf);
+        for (auto r : fullRows) {
+            cutCookTmp.Add(*TSchemedCookRow(*lay).Col(r.first, r.second));
+            cutCook.Add(*TSchemedCookRow(*lay).Col(r.first, r.second));
+            cutCookR.Add(*TSchemedCookRow(*lay).Col(r.first, r.second));
+            fullCook.Add(*TSchemedCookRow(*lay).Col(r.first, r.second));
+            fullCookR.Add(*TSchemedCookRow(*lay).Col(r.first, r.second));
+        }
+
+        TCheckIt cutWrapTmp(cutCookTmp.Finish(), { });
+        auto cutPartTmp = (*cutWrapTmp).Eggs.Lone();
+
+        auto getKey = [cutPartTmp] (const NPage::TIndex::TRecord* record) {
+            TSmallVec<TCell> key;
+            for (const auto& info : cutPartTmp->Scheme->Groups[0].ColsKeyIdx) {
+                key.push_back(record->Cell(info));
+            }
+            return TSerializedCellVec(key);
+        };
+
+        auto slices = MakeIntrusive<TSlices>();
+        for (size_t rowId = 0; rowId < fullRows.size();) {
+            TSlice slice;
+            slice.FirstInclusive = true;
+            slice.FirstRowId = rowId;
+            if (rowId > 0)
+                slice.FirstKey = getKey(cutPartTmp->Index.At(slice.FirstRowId));
+            slice.LastInclusive = false;
+            slice.LastRowId = rowId + RandomNumber<ui32>(2) + 1;
+            if (slice.LastRowId < fullRows.size())
+                slice.LastKey = getKey(cutPartTmp->Index.At(slice.LastRowId));
+            slices->push_back(slice);
+            rowId = slice.LastRowId;
+        }
+
+        Cerr << "======= SLICES =======" << Endl;
+        slices->Describe(Cerr);
+        Cerr << Endl;
+
+        TCheckIt cutWrap(cutCook.Finish(), { new TTouchEnv() }, slices), fullWrap(fullCook.Finish(), { new TTouchEnv() });
+        TCheckReverseIt cutWrapR(cutCookR.Finish(), { new TTouchEnv() }, slices), fullWrapR(fullCookR.Finish(), { new TTouchEnv() });
+
+        auto cutPart = (*cutWrap).Eggs.Lone();
+        auto fullPart = (*fullWrap).Eggs.Lone();
+
+        Cerr << "======= CUT =======" << Endl;
+        Cerr << DumpPart(*cutPart, 2) << Endl;
+        
+        Cerr << "======= FULL =======" << Endl;
+        Cerr << DumpPart(*fullPart, 2) << Endl;
+
+        UNIT_ASSERT_GT(fullPart->IndexesRawSize, cutPart->IndexesRawSize);
+        
+        for (auto r : fullRows) {
+            cutWrap.Has(*TSchemedCookRow(*lay).Col(r.first, r.second));
+            fullWrap.Has(*TSchemedCookRow(*lay).Col(r.first, r.second));
+        }
+
+        for (size_t rowId = 0; rowId < fullRows.size(); rowId++)
+        for (auto seekMode : {ESeek::Exact, ESeek::Lower, ESeek::Upper})
+        for (auto transformMode : {ESeek::Exact, ESeek::Lower, ESeek::Upper}) {
+            auto str = fullRows[rowId].second;
+
+            switch (transformMode) {
+            case ESeek::Exact:
+                break;
+            case ESeek::Lower:
+                str[str.size() - 1] = '#';
+                UNIT_ASSERT_LT(str, fullRows[rowId].second);
+                break;
+            case ESeek::Upper:
+                str += '#';
+                UNIT_ASSERT_GT(str, fullRows[rowId].second);
+                break;
+            }
+
+            auto seekRow = *TSchemedCookRow(*lay).Col(fullRows[rowId].first, str);
+
+            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());
+        }
+    }
+
     Y_UNIT_TEST(CutKeys_CutString)
     {
         TLayoutCook lay;