Fix DB.CalculateReadSize (#4728)

5 files changed, 140 insertions, 8 deletions

diff --git a/ydb/core/engine/minikql/minikql_engine_host.cpp b/ydb/core/engine/minikql/minikql_engine_host.cpp
index 00eee7d23e..d0f09d1c03 100644
--- a/ydb/core/engine/minikql/minikql_engine_host.cpp
+++ b/ydb/core/engine/minikql/minikql_engine_host.cpp
@@ -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);
@@ -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();
 
diff --git a/ydb/core/tablet_flat/flat_database.cpp b/ydb/core/tablet_flat/flat_database.cpp
index fca12cc0c8..e6290f1a03 100644
--- a/ydb/core/tablet_flat/flat_database.cpp
+++ b/ydb/core/tablet_flat/flat_database.cpp
@@ -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)
diff --git a/ydb/core/tablet_flat/flat_database.h b/ydb/core/tablet_flat/flat_database.h
index f305e01242..0e75fedec0 100644
--- a/ydb/core/tablet_flat/flat_database.h
+++ b/ydb/core/tablet_flat/flat_database.h
@@ -141,6 +141,7 @@ public:
                         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,
diff --git a/ydb/core/tablet_flat/flat_dbase_sz_env.h b/ydb/core/tablet_flat/flat_dbase_sz_env.h
index 9a5d36538a..02f8f172c6 100644
--- a/ydb/core/tablet_flat/flat_dbase_sz_env.h
+++ b/ydb/core/tablet_flat/flat_dbase_sz_env.h
@@ -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");
@@ -20,14 +25,29 @@ namespace NTable {
         {
             auto *partStore = CheckedCast<const NTable::TPartStore*>(part);
 
-            return { true, Touch(partStore->Locate(lob, ref), ref) };
+            AddPageSize(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);
+            
+            switch (type) {
+                case EPage::FlatIndex:
+                case EPage::BTreeIndex:
+                    // need index pages to continue counting
+                    // do not count index
+                    // if these pages are not in memory, data won't be counted in precharge
+                    return Env->TryGetPage(part, pageId, groupId);
+                default:
+                    AddPageSize(partStore->PageCollections.at(groupId.Index).Get(), pageId);
+                    return nullptr;
+            }
         }
 
         ui64 GetSize() const {
@@ -35,17 +55,16 @@ namespace NTable {
         }
 
     private:
-        const TSharedData* Touch(TInfo *info, TPageId page) noexcept
+        void AddPageSize(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;
diff --git a/ydb/core/tablet_flat/flat_executor_ut.cpp b/ydb/core/tablet_flat/flat_executor_ut.cpp
index 839597319c..ad7a7e308a 100644
--- a/ydb/core/tablet_flat/flat_executor_ut.cpp
+++ b/ydb/core/tablet_flat/flat_executor_ut.cpp
@@ -1,3 +1,4 @@
+#include "flat_dbase_sz_env.h"
 #include "flat_executor_ut_common.h"
 
 namespace NKikimr {
@@ -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;
@@ -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>{0, 0, 0, 0, 20566, 20566}));
+
+        env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 100, 200) });
+        UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{0, 0, 0, 0, 1048866, 1048866}));
+
+        env.SendSync(new NFake::TEvExecute{ new TTxCalculateReadSize(sizes, 300, 700) });
+        UNIT_ASSERT_VALUES_EQUAL(sizes, (TVector<ui64>{0, 0, 0, 0, 4133766, 4133766}));
+    }
+
     Y_UNIT_TEST(PrechargeAndSeek_FlatIndex) {
         TMyEnvBase env;
         TRowsModel rows;