diff --git a/ydb/core/tablet_flat/flat_part_charge_btree_index.h b/ydb/core/tablet_flat/flat_part_charge_btree_index.h index b404c7450e20..baac5f3fafea 100644 --- a/ydb/core/tablet_flat/flat_part_charge_btree_index.h +++ b/ydb/core/tablet_flat/flat_part_charge_btree_index.h @@ -12,6 +12,7 @@ class TChargeBTreeIndex : public ICharge { using TRecIdx = NPage::TRecIdx; using TGroupId = NPage::TGroupId; using TChild = TBtreeIndexNode::TChild; + using TShortChild = TBtreeIndexNode::TShortChild; struct TChildState { TPageId PageId; @@ -69,11 +70,12 @@ class TChargeBTreeIndex : public ICharge { bool ready = true, overshot = true; bool chargeGroups = bool(Groups); // false value means that beginRowId, endRowId are invalid and shouldn't be used ui64 chargeGroupsItemsLimit = itemsLimit; // pessimistic items limit for groups + TRowId beginBytesLimitRowId = Max(); // first unloaded probably needed row const auto& meta = Part->IndexPages.BTreeGroups[0]; Y_ABORT_UNLESS(endRowId <= meta.RowCount); - TRowId sliceEndRowId = endRowId; + const TRowId sliceEndRowId = endRowId; if (Y_UNLIKELY(key1 && key2 && Compare(key1, key2, keyDefaults) > 0)) { key2 = key1; // will not go further than key1 chargeGroups = false; @@ -82,13 +84,14 @@ class TChargeBTreeIndex : public ICharge { TVector level, nextLevel(::Reserve(3)); TPageId key1PageId = key1 ? meta.PageId : Max(); TPageId key2PageId = key2 ? meta.PageId : Max(); + ui64 key1Items = 0, prevKey1Items = 0; const auto iterateLevel = [&](const auto& tryHandleChild) { // tryHandleChild may update them, copy for simplicity // always load beginRowId regardless of keys const TRowId levelBeginRowId = beginRowId, levelEndRowId = Max(endRowId, beginRowId + 1); + const TChild* levelFirstChild = nullptr; - const TChild* firstChild = nullptr; for (const auto &node : level) { if (node.EndRowId <= levelBeginRowId || node.BeginRowId >= levelEndRowId) { continue; @@ -104,21 +107,28 @@ class TChargeBTreeIndex : public ICharge { for (TRecIdx pos : xrange(from, to)) { auto child = node.GetChildRef(pos); auto prevChild = pos ? node.GetChildRef(pos - 1) : nullptr; - TRowId beginRowId = prevChild ? prevChild->RowCount : node.BeginRowId; - TRowId endRowId = child->RowCount; - ready &= tryHandleChild(TChildState(child->PageId, beginRowId, endRowId)); + TRowId childBeginRowId = prevChild ? prevChild->RowCount : node.BeginRowId; + TRowId childEndRowId = child->RowCount; + ready &= tryHandleChild(TChildState(child->PageId, childBeginRowId, childEndRowId)); if (itemsLimit || bytesLimit) { - if (!firstChild) { - // do not apply limit on the first child because beginRowId/key1 position is uncertain - firstChild = child; + if (!levelFirstChild) { + // do not apply limits on the first child because beginRowId/key1 position is uncertain + levelFirstChild = child; } else { if (itemsLimit) { - ui64 items = child->GetNonErasedRowCount() - firstChild->GetNonErasedRowCount(); + ui64 items = child->GetNonErasedRowCount() - levelFirstChild->GetNonErasedRowCount(); if (LimitExceeded(items, itemsLimit)) { overshot = false; return; } } + if (bytesLimit) { + ui64 bytes = child->DataSize - levelFirstChild->DataSize; + if (LimitExceeded(bytes, bytesLimit)) { + overshot = false; + return; + } + } } } } @@ -128,14 +138,16 @@ class TChargeBTreeIndex : public ICharge { const auto skipUnloadedRows = [&](const TChildState& child) { if (child.PageId == key1PageId) { if (chargeGroups && chargeGroupsItemsLimit) { - // TODO: use erased count - ui64 unloadedItems = child.EndRowId - child.BeginRowId; + ui64 unloadedItems = key1Items - prevKey1Items; if (unloadedItems < chargeGroupsItemsLimit) { chargeGroupsItemsLimit -= unloadedItems; } else { chargeGroups = false; } } + if (chargeGroups && bytesLimit) { + beginBytesLimitRowId = Max(beginRowId, child.BeginRowId); + } beginRowId = Max(beginRowId, child.EndRowId); } if (child.PageId == key2PageId) { @@ -149,14 +161,18 @@ class TChargeBTreeIndex : public ICharge { const auto& node = nextLevel.back(); if (child.PageId == key1PageId) { TRecIdx pos = node.Seek(ESeek::Lower, key1, Scheme.Groups[0].ColsKeyIdx, &keyDefaults); - key1PageId = node.GetShortChild(pos).PageId; + auto& key1Child = node.GetChild(pos); + key1PageId = key1Child.PageId; + key1Items = key1Child.GetNonErasedRowCount(); if (pos) { - beginRowId = Max(beginRowId, node.GetShortChild(pos - 1).RowCount); // move beginRowId to the first key >= key1 + auto& prevKey1Child = node.GetChild(pos - 1); + prevKey1Items = prevKey1Child.GetNonErasedRowCount(); + beginRowId = Max(beginRowId, prevKey1Child.RowCount); // move beginRowId to the first key >= key1 } } if (child.PageId == key2PageId) { TRecIdx pos = node.Seek(ESeek::Lower, key2, Scheme.Groups[0].ColsKeyIdx, &keyDefaults); - auto& key2Child = node.GetShortChild(pos); + auto& key2Child = node.GetChild(pos); key2PageId = key2Child.PageId; endRowId = Min(endRowId, key2Child.RowCount + 1); // move endRowId - 1 to the first key > key2 if (key2Child.RowCount <= beginRowId) { @@ -207,7 +223,7 @@ class TChargeBTreeIndex : public ICharge { } if (!ready) { // some index pages are missing, do not continue - ready &= DoPrechargeGroups(chargeGroups, beginRowId, endRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds + ready &= DoGroups(chargeGroups, beginRowId, endRowId, beginBytesLimitRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds return {ready, false}; } @@ -221,7 +237,7 @@ class TChargeBTreeIndex : public ICharge { iterateLevel(tryHandleDataPage); } - ready &= DoPrechargeGroups(chargeGroups, beginRowId, endRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds + ready &= DoGroups(chargeGroups, beginRowId, endRowId, beginBytesLimitRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds return {ready, overshot}; } @@ -234,11 +250,12 @@ class TChargeBTreeIndex : public ICharge { bool ready = true, overshot = true; bool chargeGroups = bool(Groups); // false value means that beginRowId, endRowId are invalid and shouldn't be used ui64 chargeGroupsItemsLimit = itemsLimit; // pessimistic items limit for groups + TRowId endBytesLimitRowId = Max(); // last unloaded probably needed row const auto& meta = Part->IndexPages.BTreeGroups[0]; Y_ABORT_UNLESS(endRowId <= meta.RowCount); - TRowId sliceBeginRowId = beginRowId; + const TRowId sliceBeginRowId = beginRowId; if (Y_UNLIKELY(key1 && key2 && Compare(key2, key1, keyDefaults) > 0)) { key2 = key1; // will not go further than key1 chargeGroups = false; @@ -248,14 +265,14 @@ class TChargeBTreeIndex : public ICharge { TVector level, nextLevel(::Reserve(3)); TPageId key1PageId = key1 ? meta.PageId : Max(); TPageId key2PageId = key2 ? meta.PageId : Max(); + ui64 prevKey1Items = 0, key1Items = 0; const auto iterateLevel = [&](const auto& tryHandleChild) { // tryHandleChild may update them, copy for simplicity // always load endRowId - 1 regardless of keys const TRowId levelBeginRowId = Min(beginRowId, endRowId - 1), levelEndRowId = endRowId; - - const TChild* lastChild = nullptr; - const TChild* prevLastChild = nullptr; + const TChild *levelLastChild = nullptr, *levelPrevLastChild = nullptr; + for (const auto &node : level) { if (node.EndRowId <= levelBeginRowId || node.BeginRowId >= levelEndRowId) { continue; @@ -271,26 +288,33 @@ class TChargeBTreeIndex : public ICharge { for (TRecIdx posExt = to; posExt > from; posExt--) { auto child = node.GetChildRef(posExt - 1); auto prevChild = posExt - 1 ? node.GetChildRef(posExt - 2) : nullptr; - TRowId beginRowId = prevChild ? prevChild->RowCount : node.BeginRowId; - TRowId endRowId = child->RowCount; + TRowId childBeginRowId = prevChild ? prevChild->RowCount : node.BeginRowId; + TRowId childEndRowId = child->RowCount; if (itemsLimit || bytesLimit) { - if (!lastChild) { - // do not apply limit on the last child because endRowId/key1 position is uncertain - lastChild = child; + if (!levelLastChild) { + // do not apply limits on the last child because endRowId/key1 position is uncertain + levelLastChild = child; } else { - if (!prevLastChild) { - prevLastChild = child; + if (!levelPrevLastChild) { + levelPrevLastChild = child; } if (itemsLimit) { - ui64 items = prevLastChild->GetNonErasedRowCount() - child->GetNonErasedRowCount(); + ui64 items = levelPrevLastChild->GetNonErasedRowCount() - child->GetNonErasedRowCount(); if (LimitExceeded(items, itemsLimit)) { overshot = false; return; } } + if (bytesLimit) { + ui64 bytes = levelPrevLastChild->DataSize - child->DataSize; + if (LimitExceeded(bytes, bytesLimit)) { + overshot = false; + return; + } + } } } - ready &= tryHandleChild(TChildState(child->PageId, beginRowId, endRowId)); + ready &= tryHandleChild(TChildState(child->PageId, childBeginRowId, childEndRowId)); } } }; @@ -298,14 +322,16 @@ class TChargeBTreeIndex : public ICharge { const auto skipUnloadedRows = [&](const TChildState& child) { if (child.PageId == key1PageId) { if (chargeGroups && chargeGroupsItemsLimit) { - // TODO: use erased count - ui64 unloadedItems = child.EndRowId - child.BeginRowId; + ui64 unloadedItems = key1Items - prevKey1Items; if (unloadedItems < chargeGroupsItemsLimit) { chargeGroupsItemsLimit -= unloadedItems; } else { chargeGroups = false; } } + if (chargeGroups && bytesLimit) { + endBytesLimitRowId = Min(endRowId, child.EndRowId); + } endRowId = Min(endRowId, child.BeginRowId); } if (child.PageId == key2PageId) { @@ -319,15 +345,20 @@ class TChargeBTreeIndex : public ICharge { const auto& node = nextLevel.back(); if (child.PageId == key1PageId) { TRecIdx pos = node.SeekReverse(ESeek::Lower, key1, Scheme.Groups[0].ColsKeyIdx, &keyDefaults); - auto& key1Child = node.GetShortChild(pos); + auto& key1Child = node.GetChild(pos); key1PageId = key1Child.PageId; + key1Items = key1Child.GetNonErasedRowCount(); + if (pos) { + auto& prevKey1Child = node.GetChild(pos - 1); + prevKey1Items = prevKey1Child.GetNonErasedRowCount(); + } endRowId = Min(endRowId, key1Child.RowCount); // move endRowId - 1 to the last key <= key1 } if (child.PageId == key2PageId) { TRecIdx pos = node.Seek(ESeek::Lower, key2, Scheme.Groups[0].ColsKeyIdx, &keyDefaults); - key2PageId = node.GetShortChild(pos).PageId; + key2PageId = node.GetChild(pos).PageId; if (pos) { - auto& prevKey2Child = node.GetShortChild(pos - 1); + auto& prevKey2Child = node.GetChild(pos - 1); beginRowId = Max(beginRowId, prevKey2Child.RowCount - 1); // move beginRowId to the last key < key2 if (prevKey2Child.RowCount >= endRowId) { chargeGroups = false; // key2 is after current slice @@ -388,7 +419,7 @@ class TChargeBTreeIndex : public ICharge { } if (!ready) { // some index pages are missing, do not continue - ready &= DoPrechargeGroupsReverse(chargeGroups, beginRowId, endRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds + ready &= DoGroupsReverse(chargeGroups, beginRowId, endRowId, endBytesLimitRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds return {ready, false}; } @@ -402,38 +433,44 @@ class TChargeBTreeIndex : public ICharge { iterateLevel(tryHandleDataPage); } - ready &= DoPrechargeGroupsReverse(chargeGroups, beginRowId, endRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds + ready &= DoGroupsReverse(chargeGroups, beginRowId, endRowId, endBytesLimitRowId, chargeGroupsItemsLimit, bytesLimit); // precharge groups using the latest row bounds return {ready, overshot}; } private: - bool DoPrechargeGroups(bool chargeGroups, TRowId beginRowId, TRowId endRowId, ui64 itemsLimit, ui64 bytesLimit) const noexcept { + bool DoGroups(bool chargeGroups, TRowId beginRowId, TRowId endRowId, TRowId beginBytesLimitRowId, ui64 itemsLimit, ui64 bytesLimit) const noexcept { bool ready = true; if (chargeGroups && beginRowId < endRowId) { if (itemsLimit && endRowId - beginRowId - 1 >= itemsLimit) { endRowId = beginRowId + itemsLimit + 1; } + if (beginBytesLimitRowId == Max()) { + beginBytesLimitRowId = beginRowId; + } for (auto groupId : Groups) { - ready &= DoPrechargeGroup(groupId, beginRowId, endRowId, bytesLimit); + ready &= DoGroup(groupId, beginRowId, endRowId, beginBytesLimitRowId, bytesLimit); } } return ready; } - bool DoPrechargeGroupsReverse(bool chargeGroups, TRowId beginRowId, TRowId endRowId, ui64 itemsLimit, ui64 bytesLimit) const noexcept { + bool DoGroupsReverse(bool chargeGroups, TRowId beginRowId, TRowId endRowId, TRowId endBytesLimitRowId, ui64 itemsLimit, ui64 bytesLimit) const noexcept { bool ready = true; if (chargeGroups && beginRowId < endRowId) { if (itemsLimit && endRowId - beginRowId - 1 >= itemsLimit) { beginRowId = endRowId - itemsLimit - 1; } + if (endBytesLimitRowId == Max()) { + endBytesLimitRowId = endRowId; + } for (auto groupId : Groups) { - ready &= DoPrechargeGroup(groupId, beginRowId, endRowId, bytesLimit); + ready &= DoGroupReverse(groupId, beginRowId, endRowId, endBytesLimitRowId, bytesLimit); } } @@ -441,30 +478,40 @@ class TChargeBTreeIndex : public ICharge { } private: - bool DoPrechargeGroup(TGroupId groupId, TRowId beginRowId, TRowId endRowId, ui64 bytesLimit) const noexcept { + bool DoGroup(TGroupId groupId, TRowId beginRowId, TRowId endRowId, TRowId beginBytesLimitRowId, ui64 bytesLimit) const noexcept { bool ready = true; - Y_UNUSED(bytesLimit); - const auto& meta = groupId.IsHistoric() ? Part->IndexPages.BTreeHistoric[groupId.Index] : Part->IndexPages.BTreeGroups[groupId.Index]; TVector level, nextLevel(::Reserve(3)); + ui64 prevBeginDataSize = 0; + ui64 prevBeginBytesLimitDataSize = bytesLimit ? GetPrevDataSize(meta, beginBytesLimitRowId) : 0; const auto iterateLevel = [&](const auto& tryHandleChild) { + ui64 prevChildDataSize = prevBeginDataSize; for (const auto &node : level) { TRecIdx from = 0, to = node.GetChildrenCount(); if (node.BeginRowId < beginRowId) { from = node.Seek(beginRowId); + if (from) { + prevChildDataSize = prevBeginDataSize = node.GetShortChild(from - 1).DataSize; + } } if (node.EndRowId > endRowId) { to = node.Seek(endRowId - 1) + 1; } for (TRecIdx pos : xrange(from, to)) { - auto child = node.GetShortChild(pos); + auto child = node.GetShortChildRef(pos); auto prevChild = pos ? node.GetShortChildRef(pos - 1) : nullptr; - TRowId beginRowId = prevChild ? prevChild->RowCount : node.BeginRowId; - TRowId endRowId = child.RowCount; - ready &= tryHandleChild(TChildState(child.PageId, beginRowId, endRowId)); + TRowId childBeginRowId = prevChild ? prevChild->RowCount : node.BeginRowId; + TRowId childEndRowId = child->RowCount; + if (bytesLimit) { + if (prevChildDataSize > prevBeginBytesLimitDataSize && LimitExceeded(prevChildDataSize - prevBeginBytesLimitDataSize, bytesLimit)) { + return; + } + } + ready &= tryHandleChild(TChildState(child->PageId, childBeginRowId, childEndRowId)); + prevChildDataSize = child->DataSize; } } }; @@ -500,6 +547,109 @@ class TChargeBTreeIndex : public ICharge { return ready; } + bool DoGroupReverse(TGroupId groupId, TRowId beginRowId, TRowId endRowId, TRowId endBytesLimitRowId, ui64 bytesLimit) const noexcept { + bool ready = true; + + const auto& meta = groupId.IsHistoric() ? Part->IndexPages.BTreeHistoric[groupId.Index] : Part->IndexPages.BTreeGroups[groupId.Index]; + + // level's nodes is in reverse order + TVector level, nextLevel(::Reserve(3)); + ui64 endBytesLimitDataSize = bytesLimit ? GetDataSize(meta, endBytesLimitRowId - 1) : 0; + + const auto iterateLevel = [&](const auto& tryHandleChild) { + for (const auto &node : level) { + TRecIdx from = 0, to = node.GetChildrenCount(); + if (node.BeginRowId < beginRowId) { + from = node.Seek(beginRowId); + } + if (node.EndRowId > endRowId) { + to = node.Seek(endRowId - 1) + 1; + } + for (TRecIdx posExt = to; posExt > from; posExt--) { + auto child = node.GetShortChildRef(posExt - 1); + auto prevChild = posExt - 1 ? node.GetShortChildRef(posExt - 2) : nullptr; + TRowId childBeginRowId = prevChild ? prevChild->RowCount : node.BeginRowId; + TRowId childEndRowId = child->RowCount; + if (bytesLimit) { + if (endBytesLimitDataSize > child->DataSize && LimitExceeded(endBytesLimitDataSize - child->DataSize, bytesLimit)) { + return; + } + } + ready &= tryHandleChild(TChildState(child->PageId, childBeginRowId, childEndRowId)); + } + } + }; + + const auto tryHandleNode = [&](TChildState child) -> bool { + return TryLoadNode(child, nextLevel); + }; + + const auto tryHandleDataPage = [&](TChildState child) -> bool { + return HasDataPage(child.PageId, groupId); + }; + + for (ui32 height = 0; height < meta.LevelCount && ready; height++) { + if (height == 0) { + ready &= tryHandleNode(TChildState(meta.PageId, 0, meta.RowCount)); + } else { + iterateLevel(tryHandleNode); + } + level.swap(nextLevel); + nextLevel.clear(); + } + + if (!ready) { // some index pages are missing, do not continue + return ready; + } + + if (meta.LevelCount == 0) { + ready &= tryHandleDataPage(TChildState(meta.PageId, 0, meta.RowCount)); + } else { + iterateLevel(tryHandleDataPage); + } + + return ready; + } + +private: + ui64 GetPrevDataSize(const TBtreeIndexMeta& meta, TRowId rowId) const { + TPageId pageId = meta.PageId; + ui64 result = 0; + + for (ui32 height = 0; height < meta.LevelCount; height++) { + auto page = Env->TryGetPage(Part, pageId); + if (!page) { + return result; + } + auto node = TBtreeIndexNode(*page); + auto pos = node.Seek(rowId); + pageId = node.GetShortChild(pos).PageId; + if (pos) { + result = node.GetShortChild(pos - 1).DataSize; + } + } + + return result; + } + + ui64 GetDataSize(TBtreeIndexMeta meta, TRowId rowId) const { + TPageId pageId = meta.PageId; + ui64 result = meta.DataSize; + + for (ui32 height = 0; height < meta.LevelCount; height++) { + auto page = Env->TryGetPage(Part, pageId); + if (!page) { + return result; + } + auto node = TBtreeIndexNode(*page); + auto pos = node.Seek(rowId); + pageId = node.GetShortChild(pos).PageId; + result = node.GetShortChild(pos).DataSize; + } + + return result; + } + private: const TSharedData* TryGetDataPage(TPageId pageId, TGroupId groupId) const noexcept { return Env->TryGetPage(Part, pageId, groupId); @@ -509,7 +659,7 @@ class TChargeBTreeIndex : public ICharge { return bool(Env->TryGetPage(Part, pageId, groupId)); } - bool TryLoadNode(TChildState& child, TVector& level) const noexcept { + bool TryLoadNode(const TChildState& child, TVector& level) const noexcept { auto page = Env->TryGetPage(Part, child.PageId); if (!page) { return false; diff --git a/ydb/core/tablet_flat/ut/ut_btree_index_iter_charge.cpp b/ydb/core/tablet_flat/ut/ut_btree_index_iter_charge.cpp index feb6d5606444..07c5c5cff587 100644 --- a/ydb/core/tablet_flat/ut/ut_btree_index_iter_charge.cpp +++ b/ydb/core/tablet_flat/ut/ut_btree_index_iter_charge.cpp @@ -91,30 +91,29 @@ namespace { NPage::TConf conf; switch (params.Levels) { case 0: - if (params.Groups) { - conf.Group(3).PageRows = 1; - } + conf.Group(0).PageRows = 999; break; case 1: + conf.Group(0).PageRows = 2; + break; case 3: conf.Group(0).PageRows = 2; - if (params.Groups) { - for (auto i : xrange(1, 4)) { - conf.Group(i).PageRows = 1; - } - } - if (params.Levels == 3) { - conf.Group(0).BTreeIndexNodeKeysMin = conf.Group(0).BTreeIndexNodeKeysMax = 2; - if (params.Groups) { - conf.Group(1).BTreeIndexNodeKeysMin = conf.Group(1).BTreeIndexNodeKeysMax = 2; - conf.Group(2).BTreeIndexNodeKeysMin = conf.Group(2).BTreeIndexNodeKeysMax = 2; - } - } + conf.Group(0).BTreeIndexNodeKeysMin = conf.Group(0).BTreeIndexNodeKeysMax = 2; break; default: Y_Fail("Unknown levels"); } + if (params.Groups) { + conf.Group(1).PageRows = params.Levels ? 1 : 999; + conf.Group(2).PageRows = 3; + conf.Group(3).PageRows = 1; + + conf.Group(1).BTreeIndexNodeKeysMin = conf.Group(1).BTreeIndexNodeKeysMax = conf.Group(0).BTreeIndexNodeKeysMax; + conf.Group(2).BTreeIndexNodeKeysMin = conf.Group(2).BTreeIndexNodeKeysMax = 2; + conf.Group(3).BTreeIndexNodeKeysMin = conf.Group(3).BTreeIndexNodeKeysMax = 999; + } + TLayoutCook lay; lay @@ -189,9 +188,8 @@ namespace { UNIT_ASSERT_VALUES_EQUAL(part.IndexPages.BTreeGroups[0].LevelCount, params.Levels); if (params.Groups) { - UNIT_ASSERT_VALUES_EQUAL(part.IndexPages.BTreeGroups[0].LevelCount, params.Levels); UNIT_ASSERT_VALUES_EQUAL(part.IndexPages.BTreeGroups[1].LevelCount, params.Levels); - UNIT_ASSERT_VALUES_EQUAL(part.IndexPages.BTreeGroups[2].LevelCount, params.Levels); + UNIT_ASSERT_VALUES_EQUAL(part.IndexPages.BTreeGroups[2].LevelCount, 2); UNIT_ASSERT_VALUES_EQUAL(part.IndexPages.BTreeGroups[3].LevelCount, 1); } @@ -426,7 +424,7 @@ Y_UNIT_TEST_SUITE(TChargeBTreeIndex) { void CheckChargeRowId(const TPartStore& part, TTagsRef tags, const TKeyCellDefaults *keyDefaults) { for (bool reverse : {false, true}) { - for (ui32 itemsLimit : TVector{0, 1, 2, 5, 13, 19, part.Stat.Rows - 2, part.Stat.Rows - 1}) { + for (ui64 itemsLimit : TVector{0, 1, 2, 5, 13, 19, part.Stat.Rows - 2, part.Stat.Rows - 1}) { for (TRowId rowId1 : xrange(0, part.Stat.Rows - 1)) { for (TRowId rowId2 : xrange(rowId1, part.Stat.Rows - 1)) { TTouchEnv bTreeEnv, flatEnv; @@ -446,7 +444,7 @@ Y_UNIT_TEST_SUITE(TChargeBTreeIndex) { void CheckChargeKeys(const TPartStore& part, TTagsRef tags, const TKeyCellDefaults *keyDefaults) { for (bool reverse : {false, true}) { - for (ui32 itemsLimit : TVector{0, 1, 2, 5, 13, 19, part.Stat.Rows - 2, part.Stat.Rows - 1}) { + for (ui64 itemsLimit : TVector{0, 1, 2, 5, 13, 19, part.Stat.Rows - 2, part.Stat.Rows - 1}) { for (ui32 firstCellKey1 : xrange(0, part.Stat.Rows / 7 + 1)) { for (ui32 secondCellKey1 : xrange(0, 14)) { for (ui32 firstCellKey2 : xrange(0, part.Stat.Rows / 7 + 1)) { @@ -487,6 +485,66 @@ Y_UNIT_TEST_SUITE(TChargeBTreeIndex) { } } + void CheckChargeBytesLimit(const TPartStore& part, TTagsRef tags, const TKeyCellDefaults *keyDefaults) { + for (bool reverse : {false, true}) { + for (ui64 bytesLimit : xrange(1, part.Stat.Bytes + 100, part.Stat.Bytes / 100)) { + for (ui32 firstCellKey1 : xrange(0, part.Stat.Rows / 7 + 1)) { + for (ui32 secondCellKey1 : xrange(0, 14)) { + TVector key1 = MakeKey(firstCellKey1, secondCellKey1); + + TTouchEnv limitedEnv, unlimitedEnv; + TChargeBTreeIndex limitedCharge(&limitedEnv, part, tags, true); + TChargeBTreeIndex unlimitedCharge(&unlimitedEnv, part, tags, true); + + TStringBuilder message = TStringBuilder() << (reverse ? "ChargeBytesLimitReverse " : "ChargeBytesLimit ") << "("; + for (auto c : key1) { + message << c.AsValue() << " "; + } + message << ") bytes " << bytesLimit; + + DoChargeKeys(part, limitedCharge, limitedEnv, key1, { }, 0, bytesLimit, reverse, *keyDefaults, message); + DoChargeKeys(part, unlimitedCharge, unlimitedEnv, key1, { }, 0, 0, reverse, *keyDefaults, message); + + TSet groupIds; + for (const auto &c : {limitedEnv.Loaded, unlimitedEnv.Loaded}) { + for (const auto &g : c) { + groupIds.insert(g.first); + } + } + for (auto groupId : groupIds) { + ui64 size = 0; + TSet expected, loaded; + TVector unlimitedLoaded(unlimitedEnv.Loaded[groupId].begin(), unlimitedEnv.Loaded[groupId].end()); + if (reverse) { + std::reverse(unlimitedLoaded.begin(), unlimitedLoaded.end()); + } + for (auto pageId : unlimitedLoaded) { + if (part.GetPageType(pageId, groupId) == EPage::DataPage) { + if (expected || !groupId.IsMain()) { + // do not count first main page + size += part.GetPageSize(pageId, groupId); + } + expected.insert(pageId); + if (size > bytesLimit) { + break; + } + } + } + for (auto pageId : limitedEnv.Loaded[groupId]) { + if (part.GetPageType(pageId, groupId) == EPage::DataPage) { + loaded.insert(pageId); + } + } + + UNIT_ASSERT_VALUES_EQUAL_C(expected, loaded, + TStringBuilder() << message << " Group {" << groupId.Index << "," << groupId.IsHistoric() << "}"); + } + } + } + } + } + } + void CheckPart(TMakePartParams params) { TPartEggs eggs = MakePart(params); const auto part = *eggs.Lone(); @@ -498,6 +556,7 @@ Y_UNIT_TEST_SUITE(TChargeBTreeIndex) { CheckChargeRowId(part, tags, eggs.Scheme->Keys.Get()); CheckChargeKeys(part, tags, eggs.Scheme->Keys.Get()); + CheckChargeBytesLimit(part, tags, eggs.Scheme->Keys.Get()); } Y_UNIT_TEST(NoNodes) {