Skip to content

Commit 3e9854b

Browse files
authored
Merge d7b7f5b into 97d8627
2 parents 97d8627 + d7b7f5b commit 3e9854b

File tree

5 files changed

+140
-8
lines changed

5 files changed

+140
-8
lines changed

ydb/core/engine/minikql/minikql_engine_host.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ bool TEngineHost::IsValidKey(TKeyDesc& key) const {
7373
return NMiniKQL::IsValidKey(Scheme, localTableId, key);
7474
}
7575
ui64 TEngineHost::CalculateReadSize(const TVector<const TKeyDesc*>& keys) const {
76-
NTable::TSizeEnv env;
76+
auto env = Db.CreateSizeEnv();
7777

7878
for (const TKeyDesc* ki : keys) {
7979
DoCalculateReadSize(*ki, env);
@@ -120,7 +120,7 @@ ui64 TEngineHost::CalculateResultSize(const TKeyDesc& key) const {
120120
if (key.Range.Point) {
121121
return Db.EstimateRowSize(localTid);
122122
} else {
123-
NTable::TSizeEnv env;
123+
auto env = Db.CreateSizeEnv();
124124
DoCalculateReadSize(key, env);
125125
ui64 size = env.GetSize();
126126

ydb/core/tablet_flat/flat_database.cpp

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,12 @@ TSelectRowVersionResult TDatabase::SelectRowVersion(
270270
return Require(table)->SelectRowVersion(key, Env, readFlags, visible, observer);
271271
}
272272

273+
TSizeEnv TDatabase::CreateSizeEnv()
274+
{
275+
return TSizeEnv(Env);
276+
}
277+
278+
273279
void TDatabase::CalculateReadSize(TSizeEnv& env, ui32 table, TRawVals minKey, TRawVals maxKey,
274280
TTagsRef tags, ui64 flg, ui64 items, ui64 bytes,
275281
EDirection direction, TRowVersion snapshot)

ydb/core/tablet_flat/flat_database.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,7 @@ class TDatabase {
141141
EDirection direction = EDirection::Forward,
142142
TRowVersion snapshot = TRowVersion::Max());
143143

144+
TSizeEnv CreateSizeEnv();
144145
void CalculateReadSize(TSizeEnv& env, ui32 table, TRawVals minKey, TRawVals maxKey,
145146
TTagsRef tags, ui64 readFlags, ui64 itemsLimit, ui64 bytesLimit,
146147
EDirection direction = EDirection::Forward,

ydb/core/tablet_flat/flat_dbase_sz_env.h

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,11 @@ namespace NTable {
1111
struct TSizeEnv : public IPages {
1212
using TInfo = NTabletFlatExecutor::TPrivatePageCache::TInfo;
1313

14+
TSizeEnv(IPages* env)
15+
: Env(env)
16+
{
17+
}
18+
1419
TResult Locate(const TMemTable*, ui64, ui32) noexcept override
1520
{
1621
Y_ABORT("IPages::Locate(TMemTable*, ...) shouldn't be used here");
@@ -20,32 +25,46 @@ namespace NTable {
2025
{
2126
auto *partStore = CheckedCast<const NTable::TPartStore*>(part);
2227

23-
return { true, Touch(partStore->Locate(lob, ref), ref) };
28+
AddPageSize(partStore->Locate(lob, ref), ref);
29+
30+
return { true, nullptr };
2431
}
2532

26-
const TSharedData* TryGetPage(const TPart* part, TPageId page, TGroupId groupId) override
33+
const TSharedData* TryGetPage(const TPart* part, TPageId pageId, TGroupId groupId) override
2734
{
2835
auto *partStore = CheckedCast<const NTable::TPartStore*>(part);
2936

30-
return Touch(partStore->PageCollections.at(groupId.Index).Get(), page);
37+
auto info = partStore->PageCollections.at(groupId.Index).Get();
38+
auto type = EPage(info->PageCollection->Page(pageId).Type);
39+
40+
switch (type) {
41+
case EPage::FlatIndex:
42+
case EPage::BTreeIndex:
43+
// need index pages to continue counting
44+
// do not count index
45+
// if these pages are not in memory, data won't be counted in precharge
46+
return Env->TryGetPage(part, pageId, groupId);
47+
default:
48+
AddPageSize(partStore->PageCollections.at(groupId.Index).Get(), pageId);
49+
return nullptr;
50+
}
3151
}
3252

3353
ui64 GetSize() const {
3454
return Bytes;
3555
}
3656

3757
private:
38-
const TSharedData* Touch(TInfo *info, TPageId page) noexcept
58+
void AddPageSize(TInfo *info, TPageId page) noexcept
3959
{
4060
if (Touched[info].insert(page).second) {
4161
Pages++;
4262
Bytes += info->PageCollection->Page(page).Size;
4363
}
44-
45-
return nullptr;
4664
}
4765

4866
private:
67+
IPages* Env;
4968
THashMap<const void*, THashSet<TPageId>> Touched;
5069
ui64 Pages = 0;
5170
ui64 Bytes = 0;

ydb/core/tablet_flat/flat_executor_ut.cpp

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
#include "flat_dbase_sz_env.h"
12
#include "flat_executor_ut_common.h"
23

34
namespace NKikimr {
@@ -5008,6 +5009,41 @@ Y_UNIT_TEST_SUITE(TFlatTableSnapshotWithCommits) {
50085009

50095010
Y_UNIT_TEST_SUITE(TFlatTableExecutorIndexLoading) {
50105011

5012+
struct TTxCalculateReadSize : public ITransaction {
5013+
ui32 Attempt = 0;
5014+
TVector<ui64>& ReadSizes;
5015+
ui64 MinKey, MaxKey;
5016+
5017+
TTxCalculateReadSize(TVector<ui64>& readSizes, ui64 minKey, ui64 maxKey)
5018+
: ReadSizes(readSizes)
5019+
, MinKey(minKey)
5020+
, MaxKey(maxKey)
5021+
{
5022+
ReadSizes.clear();
5023+
}
5024+
5025+
5026+
bool Execute(TTransactionContext &txc, const TActorContext &) override
5027+
{
5028+
UNIT_ASSERT_LE(++Attempt, 10);
5029+
5030+
const auto minKey = NScheme::TInt64::TInstance(MinKey);
5031+
const auto maxKey = NScheme::TInt64::TInstance(MaxKey);
5032+
const TVector<NTable::TTag> tags{ { TRowsModel::ColumnKeyId, TRowsModel::ColumnValueId } };
5033+
5034+
auto sizeEnv = txc.DB.CreateSizeEnv();
5035+
txc.DB.CalculateReadSize(sizeEnv, TRowsModel::TableId, { minKey }, { maxKey }, tags, 0, 0, 0);
5036+
ReadSizes.push_back(sizeEnv.GetSize());
5037+
5038+
return txc.DB.Precharge(TRowsModel::TableId, { minKey }, { maxKey }, tags, 0, 0, 0);
5039+
}
5040+
5041+
void Complete(const TActorContext &ctx) override
5042+
{
5043+
ctx.Send(ctx.SelfID, new NFake::TEvReturn);
5044+
}
5045+
};
5046+
50115047
struct TTxPrechargeAndSeek : public ITransaction {
50125048
ui32 Attempt = 0;
50135049
bool Pinned = false;
@@ -5056,6 +5092,76 @@ Y_UNIT_TEST_SUITE(TFlatTableExecutorIndexLoading) {
50565092
env->DispatchEvents(options);
50575093
}
50585094

5095+
Y_UNIT_TEST(CalculateReadSize_FlatIndex) {
5096+
TMyEnvBase env;
5097+
TRowsModel rows;
5098+
const ui32 rowsCount = 1024;
5099+
5100+
auto &appData = env->GetAppData();
5101+
appData.FeatureFlags.SetEnableLocalDBBtreeIndex(false);
5102+
appData.FeatureFlags.SetEnableLocalDBFlatIndex(true);
5103+
5104+
env.FireTablet(env.Edge, env.Tablet, [&env](const TActorId &tablet, TTabletStorageInfo *info) {
5105+
return new TTestFlatTablet(env.Edge, tablet, info);
5106+
});
5107+
env.WaitForWakeUp();
5108+
ZeroSharedCache(env);
5109+
5110+
env.SendSync(rows.MakeScheme(new TCompactionPolicy(), false));
5111+
5112+
env.SendSync(rows.MakeRows(rowsCount, 10*1024));
5113+
5114+
env.SendSync(new NFake::TEvCompact(TRowsModel::TableId));
5115+
env.WaitFor<NFake::TEvCompacted>();
5116+
5117+
TVector<ui64> sizes;
5118+
5119+
env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 0, 1) });
5120+
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{20566, 20566}));
5121+
5122+
env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 100, 200) });
5123+
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{1048866, 1048866}));
5124+
5125+
env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 300, 700) });
5126+
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{4133766, 4133766}));
5127+
}
5128+
5129+
Y_UNIT_TEST(CalculateReadSize_BTreeIndex) {
5130+
TMyEnvBase env;
5131+
TRowsModel rows;
5132+
const ui32 rowsCount = 1024;
5133+
5134+
auto &appData = env->GetAppData();
5135+
appData.FeatureFlags.SetEnableLocalDBBtreeIndex(true);
5136+
appData.FeatureFlags.SetEnableLocalDBFlatIndex(false);
5137+
5138+
env.FireTablet(env.Edge, env.Tablet, [&env](const TActorId &tablet, TTabletStorageInfo *info) {
5139+
return new TTestFlatTablet(env.Edge, tablet, info);
5140+
});
5141+
env.WaitForWakeUp();
5142+
ZeroSharedCache(env);
5143+
5144+
auto policy = MakeIntrusive<TCompactionPolicy>();
5145+
policy->MinBTreeIndexNodeSize = 128;
5146+
env.SendSync(rows.MakeScheme(std::move(policy)));
5147+
5148+
env.SendSync(rows.MakeRows(rowsCount, 10*1024));
5149+
5150+
env.SendSync(new NFake::TEvCompact(TRowsModel::TableId));
5151+
env.WaitFor<NFake::TEvCompacted>();
5152+
5153+
TVector<ui64> sizes;
5154+
5155+
env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 0, 1) });
5156+
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{0, 0, 0, 0, 20566, 20566}));
5157+
5158+
env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 100, 200) });
5159+
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{0, 0, 0, 0, 1048866, 1048866}));
5160+
5161+
env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 300, 700) });
5162+
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{0, 0, 0, 0, 4133766, 4133766}));
5163+
}
5164+
50595165
Y_UNIT_TEST(PrechargeAndSeek_FlatIndex) {
50605166
TMyEnvBase env;
50615167
TRowsModel rows;

0 commit comments

Comments
 (0)