Add query perf test for KV read. ()

3 files changed, 85 insertions, 0 deletions

diff --git a/ydb/core/kqp/ut/common/kqp_ut_common.cpp b/ydb/core/kqp/ut/common/kqp_ut_common.cpp
index ea6ae04fae..ae3826794d 100644
--- a/ydb/core/kqp/ut/common/kqp_ut_common.cpp
+++ b/ydb/core/kqp/ut/common/kqp_ut_common.cpp
@@ -818,12 +818,40 @@ void FindPlanNodesImpl(const NJson::TJsonValue& node, const TString& key, std::v
     }
 }
 
+void FindPlanStagesImpl(const NJson::TJsonValue& node, std::vector<NJson::TJsonValue>& stages) {
+    if (node.IsArray()) {
+        for (const auto& item: node.GetArray()) {
+            FindPlanStagesImpl(item, stages);
+        }
+    }
+
+    if (!node.IsMap()) {
+        return;
+    }
+
+    auto map = node.GetMap();
+    // TODO: Use explicit PlanNodeType for stages
+    if (map.contains("Node Type") && !map.contains("PlanNodeType")) {
+        stages.push_back(node);
+    }
+
+    for (const auto& [_, value]: map) {
+        FindPlanStagesImpl(value, stages);
+    }
+}
+
 std::vector<NJson::TJsonValue> FindPlanNodes(const NJson::TJsonValue& plan, const TString& key) {
     std::vector<NJson::TJsonValue> results;
     FindPlanNodesImpl(plan, key, results);
     return results;
 }
 
+std::vector<NJson::TJsonValue> FindPlanStages(const NJson::TJsonValue& plan) {
+    std::vector<NJson::TJsonValue> stages;
+    FindPlanStagesImpl(plan, stages);
+    return stages;
+}
+
 void CreateSampleTablesWithIndex(TSession& session) {
     auto res = session.ExecuteSchemeQuery(R"(
         --!syntax_v1
diff --git a/ydb/core/kqp/ut/common/kqp_ut_common.h b/ydb/core/kqp/ut/common/kqp_ut_common.h
index cd15942a43..d9a7ef2b94 100644
--- a/ydb/core/kqp/ut/common/kqp_ut_common.h
+++ b/ydb/core/kqp/ut/common/kqp_ut_common.h
@@ -240,6 +240,7 @@ TString StreamResultToYson(NYdb::NTable::TTablePartIterator& it);
 ui32 CountPlanNodesByKv(const NJson::TJsonValue& plan, const TString& key, const TString& value);
 NJson::TJsonValue FindPlanNodeByKv(const NJson::TJsonValue& plan, const TString& key, const TString& value);
 std::vector<NJson::TJsonValue> FindPlanNodes(const NJson::TJsonValue& plan, const TString& key);
+std::vector<NJson::TJsonValue> FindPlanStages(const NJson::TJsonValue& plan);
 
 TString ReadTableToYson(NYdb::NTable::TSession session, const TString& table);
 TString ReadTablePartToYson(NYdb::NTable::TSession session, const TString& table);
diff --git a/ydb/core/kqp/ut/perf/kqp_query_perf_ut.cpp b/ydb/core/kqp/ut/perf/kqp_query_perf_ut.cpp
index 7f8d5c3bf9..bc411fa861 100644
--- a/ydb/core/kqp/ut/perf/kqp_query_perf_ut.cpp
+++ b/ydb/core/kqp/ut/perf/kqp_query_perf_ut.cpp
@@ -120,6 +120,62 @@ TParams BuildInsertIndexParams(TTableClient& client) {
 } // namespace
 
 Y_UNIT_TEST_SUITE(KqpQueryPerf) {
+    Y_UNIT_TEST_TWIN(KvRead, EnableStreamLookup) {
+        NKikimrConfig::TAppConfig appConfig;
+        appConfig.MutableTableServiceConfig()->SetEnableKqpDataQueryStreamLookup(EnableStreamLookup);
+        auto settings = TKikimrSettings()
+            .SetAppConfig(appConfig);
+        TKikimrRunner kikimr{settings};
+
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto params = db.GetParamsBuilder()
+            .AddParam("$key").Uint64(102).Build()
+            .Build();
+
+        NYdb::NTable::TExecDataQuerySettings execSettings;
+        execSettings.CollectQueryStats(ECollectQueryStatsMode::Full);
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            DECLARE $key AS Uint64;
+
+            SELECT * FROM EightShard
+            WHERE Key = $key;
+        )"), TTxControl::BeginTx().CommitTx(), params, execSettings).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        // Cerr << stats.query_plan() << Endl;
+
+        // TODO: Fix stream lookup case
+        if (!EnableStreamLookup) {
+            AssertTableStats(result, "/Root/EightShard", {
+                .ExpectedReads = 1,
+            });
+        }
+
+        auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+
+        UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), 1);
+        for (const auto& phase : stats.query_phases()) {
+            UNIT_ASSERT(phase.affected_shards() <= 1);
+        }
+
+        NJson::TJsonValue plan;
+        NJson::ReadJsonTree(stats.query_plan(), &plan, true);
+
+        auto stages = FindPlanStages(plan);
+        // TODO: Fix stream lookup case
+        UNIT_ASSERT_VALUES_EQUAL(stages.size(), EnableStreamLookup ? 3 : 2);
+
+        i64 totalTasks = 0;
+        for (const auto& stage : stages) {
+            totalTasks += stage.GetMapSafe().at("Stats").GetMapSafe().at("TotalTasks").GetIntegerSafe();
+        }
+        // TODO: Fix stream lookup case
+        UNIT_ASSERT_VALUES_EQUAL(totalTasks, EnableStreamLookup ? 3 : 2);
+    }
+
     Y_UNIT_TEST(Upsert) {
         auto kikimr = DefaultKikimrRunner();
         auto db = kikimr.GetTableClient();