Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix DB.CalculateReadSize #4728

Merged
merged 2 commits into from
May 22, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions ydb/core/engine/minikql/minikql_engine_host.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ bool TEngineHost::IsValidKey(TKeyDesc& key) const {
return NMiniKQL::IsValidKey(Scheme, localTableId, key);
}
ui64 TEngineHost::CalculateReadSize(const TVector<const TKeyDesc*>& keys) const {
NTable::TSizeEnv env;
auto env = Db.CreateSizeEnv();

for (const TKeyDesc* ki : keys) {
DoCalculateReadSize(*ki, env);
Expand Down Expand Up @@ -120,7 +120,7 @@ ui64 TEngineHost::CalculateResultSize(const TKeyDesc& key) const {
if (key.Range.Point) {
return Db.EstimateRowSize(localTid);
} else {
NTable::TSizeEnv env;
auto env = Db.CreateSizeEnv();
DoCalculateReadSize(key, env);
ui64 size = env.GetSize();

Expand Down
6 changes: 6 additions & 0 deletions ydb/core/tablet_flat/flat_database.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,12 @@ TSelectRowVersionResult TDatabase::SelectRowVersion(
return Require(table)->SelectRowVersion(key, Env, readFlags, visible, observer);
}

TSizeEnv TDatabase::CreateSizeEnv()
{
return TSizeEnv(Env);
}


void TDatabase::CalculateReadSize(TSizeEnv& env, ui32 table, TRawVals minKey, TRawVals maxKey,
TTagsRef tags, ui64 flg, ui64 items, ui64 bytes,
EDirection direction, TRowVersion snapshot)
Expand Down
1 change: 1 addition & 0 deletions ydb/core/tablet_flat/flat_database.h
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,7 @@ class TDatabase {
EDirection direction = EDirection::Forward,
TRowVersion snapshot = TRowVersion::Max());

TSizeEnv CreateSizeEnv();
void CalculateReadSize(TSizeEnv& env, ui32 table, TRawVals minKey, TRawVals maxKey,
TTagsRef tags, ui64 readFlags, ui64 itemsLimit, ui64 bytesLimit,
EDirection direction = EDirection::Forward,
Expand Down
30 changes: 24 additions & 6 deletions ydb/core/tablet_flat/flat_dbase_sz_env.h
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,11 @@ namespace NTable {
struct TSizeEnv : public IPages {
using TInfo = NTabletFlatExecutor::TPrivatePageCache::TInfo;

TSizeEnv(IPages* env)
: Env(env)
{
}

TResult Locate(const TMemTable*, ui64, ui32) noexcept override
{
Y_ABORT("IPages::Locate(TMemTable*, ...) shouldn't be used here");
Expand All @@ -20,32 +25,45 @@ namespace NTable {
{
auto *partStore = CheckedCast<const NTable::TPartStore*>(part);

return { true, Touch(partStore->Locate(lob, ref), ref) };
Touch(partStore->Locate(lob, ref), ref);

return { true, nullptr };
}

const TSharedData* TryGetPage(const TPart* part, TPageId page, TGroupId groupId) override
const TSharedData* TryGetPage(const TPart* part, TPageId pageId, TGroupId groupId) override
{
auto *partStore = CheckedCast<const NTable::TPartStore*>(part);

return Touch(partStore->PageCollections.at(groupId.Index).Get(), page);
auto info = partStore->PageCollections.at(groupId.Index).Get();
auto type = EPage(info->PageCollection->Page(pageId).Type);

if (type != EPage::FlatIndex) { // do not count flat index pages
kunga marked this conversation as resolved.
Show resolved Hide resolved
Touch(partStore->PageCollections.at(groupId.Index).Get(), pageId);
}

if (type == EPage::FlatIndex || type == EPage::BTreeIndex) {
// need page data to continue counting
return Env->TryGetPage(part, pageId, groupId);
}

return nullptr;
}

ui64 GetSize() const {
return Bytes;
}

private:
const TSharedData* Touch(TInfo *info, TPageId page) noexcept
void Touch(TInfo *info, TPageId page) noexcept
{
if (Touched[info].insert(page).second) {
Pages++;
Bytes += info->PageCollection->Page(page).Size;
}

return nullptr;
}

private:
IPages* Env;
THashMap<const void*, THashSet<TPageId>> Touched;
ui64 Pages = 0;
ui64 Bytes = 0;
Expand Down
106 changes: 106 additions & 0 deletions ydb/core/tablet_flat/flat_executor_ut.cpp
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
#include "flat_dbase_sz_env.h"
#include "flat_executor_ut_common.h"

namespace NKikimr {
Expand Down Expand Up @@ -5008,6 +5009,41 @@ Y_UNIT_TEST_SUITE(TFlatTableSnapshotWithCommits) {

Y_UNIT_TEST_SUITE(TFlatTableExecutorIndexLoading) {

struct TTxCalculateReadSize : public ITransaction {
ui32 Attempt = 0;
TVector<ui64>& ReadSizes;
ui64 MinKey, MaxKey;

TTxCalculateReadSize(TVector<ui64>& readSizes, ui64 minKey, ui64 maxKey)
: ReadSizes(readSizes)
, MinKey(minKey)
, MaxKey(maxKey)
{
ReadSizes.clear();
}


bool Execute(TTransactionContext &txc, const TActorContext &) override
{
UNIT_ASSERT_LE(++Attempt, 10);

const auto minKey = NScheme::TInt64::TInstance(MinKey);
const auto maxKey = NScheme::TInt64::TInstance(MaxKey);
const TVector<NTable::TTag> tags{ { TRowsModel::ColumnKeyId, TRowsModel::ColumnValueId } };

auto sizeEnv = txc.DB.CreateSizeEnv();
txc.DB.CalculateReadSize(sizeEnv, TRowsModel::TableId, { minKey }, { maxKey }, tags, 0, 0, 0);
ReadSizes.push_back(sizeEnv.GetSize());

return txc.DB.Precharge(TRowsModel::TableId, { minKey }, { maxKey }, tags, 0, 0, 0);
}

void Complete(const TActorContext &ctx) override
{
ctx.Send(ctx.SelfID, new NFake::TEvReturn);
}
};

struct TTxPrechargeAndSeek : public ITransaction {
ui32 Attempt = 0;
bool Pinned = false;
Expand Down Expand Up @@ -5056,6 +5092,76 @@ Y_UNIT_TEST_SUITE(TFlatTableExecutorIndexLoading) {
env->DispatchEvents(options);
}

Y_UNIT_TEST(CalculateReadSize_FlatIndex) {
TMyEnvBase env;
TRowsModel rows;
const ui32 rowsCount = 1024;

auto &appData = env->GetAppData();
appData.FeatureFlags.SetEnableLocalDBBtreeIndex(false);
appData.FeatureFlags.SetEnableLocalDBFlatIndex(true);

env.FireTablet(env.Edge, env.Tablet, [&env](const TActorId &tablet, TTabletStorageInfo *info) {
return new TTestFlatTablet(env.Edge, tablet, info);
});
env.WaitForWakeUp();
ZeroSharedCache(env);

env.SendSync(rows.MakeScheme(new TCompactionPolicy(), false));

env.SendSync(rows.MakeRows(rowsCount, 10*1024));

env.SendSync(new NFake::TEvCompact(TRowsModel::TableId));
env.WaitFor<NFake::TEvCompacted>();

TVector<ui64> sizes;

env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 0, 1) });
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{20566, 20566}));

env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 100, 200) });
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{1048866, 1048866}));

env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 300, 700) });
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{4133766, 4133766}));
}

Y_UNIT_TEST(CalculateReadSize_BTreeIndex) {
TMyEnvBase env;
TRowsModel rows;
const ui32 rowsCount = 1024;

auto &appData = env->GetAppData();
appData.FeatureFlags.SetEnableLocalDBBtreeIndex(true);
appData.FeatureFlags.SetEnableLocalDBFlatIndex(false);

env.FireTablet(env.Edge, env.Tablet, [&env](const TActorId &tablet, TTabletStorageInfo *info) {
return new TTestFlatTablet(env.Edge, tablet, info);
});
env.WaitForWakeUp();
ZeroSharedCache(env);

auto policy = MakeIntrusive<TCompactionPolicy>();
policy->MinBTreeIndexNodeSize = 128;
env.SendSync(rows.MakeScheme(std::move(policy)));

env.SendSync(rows.MakeRows(rowsCount, 10*1024));

env.SendSync(new NFake::TEvCompact(TRowsModel::TableId));
env.WaitFor<NFake::TEvCompacted>();

TVector<ui64> sizes;

env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 0, 1) });
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{102, 1090, 1784, 2478, 23044, 23044}));

env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 100, 200) });
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{102, 1090, 2478, 8030, 1056896, 1056896}));

env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 300, 700) });
UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{102, 1090, 4560, 25033, 4158799, 4158799}));
}

Y_UNIT_TEST(PrechargeAndSeek_FlatIndex) {
TMyEnvBase env;
TRowsModel rows;
Expand Down
Loading