[KQP] Fix legacy Explain Analyze scan stats (#46498)

2 files changed, 103 insertions, 37 deletions

diff --git a/ydb/core/kqp/opt/kqp_query_plan.cpp b/ydb/core/kqp/opt/kqp_query_plan.cpp
index 6fa5645b155..16bc8d5ff1f 100644
--- a/ydb/core/kqp/opt/kqp_query_plan.cpp
+++ b/ydb/core/kqp/opt/kqp_query_plan.cpp
@@ -2481,7 +2481,7 @@ public:
     NJson::TJsonValue Reconstruct(
         const NJson::TJsonValue& plan
     ) {
-        auto reconstructed = ReconstructImpl(plan, 0, 0, false);
+        auto reconstructed = ReconstructImpl(plan, 0, 0, false, nullptr);
         return reconstructed;
     }
 
@@ -2490,16 +2490,21 @@ private:
         const NJson::TJsonValue& plan,
         int operatorIndex,
         int parentTaskCount,
-        bool fromBroadcast
+        bool fromBroadcast,
+        const NJson::TJsonValue* inheritedTableStats
     ) {
         int currentNodeId = NodeIDCounter++;
 
         int taskCount = parentTaskCount;
+        const NJson::TJsonValue* ownTableStats = nullptr;
         if (plan.GetMapSafe().contains("Stats")) {
             const auto& stats = plan.GetMapSafe().at("Stats").GetMapSafe();
             if (stats.contains("Tasks")) {
                 taskCount = stats.at("Tasks").GetIntegerSafe();
             }
+            if (stats.contains("Table")) {
+                ownTableStats = &stats.at("Table");
+            }
         }
 
         NJson::TJsonValue result;
@@ -2553,7 +2558,7 @@ private:
             lookupPlan["Operators"] = std::move(lookupOps);
 
 	        if (plan.GetMapSafe().contains("Plans")) {
-                newPlans.AppendValue(ReconstructImpl(plan.GetMapSafe().at("Plans").GetArraySafe()[0], 0, taskCount, false));
+                newPlans.AppendValue(ReconstructImpl(plan.GetMapSafe().at("Plans").GetArraySafe()[0], 0, taskCount, false, inheritedTableStats));
             }
 
             newPlans.AppendValue(std::move(lookupPlan));
@@ -2576,7 +2581,7 @@ private:
             if (plan.GetMapSafe().contains("CTE Name")) {
                 auto precompute = plan.GetMapSafe().at("CTE Name").GetStringSafe();
                 if (Precomputes.contains(precompute)) {
-                    planInputs.AppendValue(ReconstructImpl(Precomputes.at(precompute), 0, taskCount, false));
+                    planInputs.AppendValue(ReconstructImpl(Precomputes.at(precompute), 0, taskCount, false, nullptr));
                 }
             }
 
@@ -2614,10 +2619,10 @@ private:
                 if (!p.GetMapSafe().contains("Operators") && p.GetMapSafe().contains("CTE Name")) {
                     auto precompute = p.GetMapSafe().at("CTE Name").GetStringSafe();
                     if (Precomputes.contains(precompute)) {
-                        planInputs.AppendValue(ReconstructImpl(Precomputes.at(precompute), 0, taskCount, false));
+                        planInputs.AppendValue(ReconstructImpl(Precomputes.at(precompute), 0, taskCount, false, nullptr));
                     }
                 } else if (p.GetMapSafe().at("Node Type").GetStringSafe().find("Precompute") == TString::npos) {
-                    planInputs.AppendValue(ReconstructImpl(p, 0, taskCount, fromBroadcast));
+                    planInputs.AppendValue(ReconstructImpl(p, 0, taskCount, fromBroadcast, inheritedTableStats));
                 }
             }
             result["Plans"] = planInputs;
@@ -2631,7 +2636,7 @@ private:
                 return result;
             }
 
-            return ReconstructImpl(Precomputes.at(precompute), 0, taskCount, false);
+            return ReconstructImpl(Precomputes.at(precompute), 0, taskCount, false, nullptr);
         }
 
         auto ops = plan.GetMapSafe().at("Operators").GetArraySafe();
@@ -2654,7 +2659,7 @@ private:
                 processedExternalOperators.insert(inputPlanKey);
 
                 auto inputPlan = PlanIndex.at(inputPlanKey);
-                planInputs.push_back( ReconstructImpl(inputPlan, 0, taskCount, inputPlan.GetMapSafe().at("Node Type").GetStringSafe() == "Broadcast") );
+                planInputs.push_back( ReconstructImpl(inputPlan, 0, taskCount, inputPlan.GetMapSafe().at("Node Type").GetStringSafe() == "Broadcast", ownTableStats) );
             } else if (opInput.GetMapSafe().contains("InternalOperatorId")) {
                 auto inputPlanId = opInput.GetMapSafe().at("InternalOperatorId").GetIntegerSafe();
 
@@ -2663,7 +2668,7 @@ private:
                 }
                 processedInternalOperators.insert(inputPlanId);
 
-                planInputs.push_back( ReconstructImpl(plan, inputPlanId, taskCount, false) );
+                planInputs.push_back( ReconstructImpl(plan, inputPlanId, taskCount, false, inheritedTableStats) );
             }
         }
 
@@ -2696,17 +2701,44 @@ private:
             }
 
             if (Precomputes.contains(maybePrecompute) && planInputs.empty()) {
-                planInputs.push_back(ReconstructImpl(Precomputes.at(maybePrecompute), 0, taskCount, false));
+                planInputs.push_back(ReconstructImpl(Precomputes.at(maybePrecompute), 0, taskCount, false, nullptr));
             }
         }
 
         result["Node Type"] = opName;
 
+        auto operatorSize = false;
+        auto operatorRows = false;
+        const auto applyTableStats = [&](const NJson::TJsonValue& tableStats) {
+            TString tablePath;
+            if (op.GetMapSafe().contains("Path")) {
+                tablePath = op.GetMapSafe().at("Path").GetStringSafe();
+            } else if (op.GetMapSafe().contains("Table")) {
+                tablePath = op.GetMapSafe().at("Table").GetStringSafe();
+            }
+            if (tablePath) {
+                for (auto& opStat : tableStats.GetArraySafe()) {
+                    if (opStat.IsMap()) {
+                        auto& opMap = opStat.GetMapSafe();
+                        if (opMap.contains("Path") && opMap.at("Path").GetStringSafe() == tablePath) {
+                            if (opMap.contains("ReadRows")) {
+                                op["A-Rows"] = opMap.at("ReadRows").GetMapSafe().at("Sum").GetDouble();
+                                operatorRows = true;
+                            }
+                            if (opMap.contains("ReadBytes")) {
+                                op["A-Size"] = opMap.at("ReadBytes").GetMapSafe().at("Sum").GetDouble();
+                                operatorSize = true;
+                            }
+                            break;
+                        }
+                    }
+                }
+            }
+        };
+
         if (plan.GetMapSafe().contains("Stats")) {
             const auto& stats = plan.GetMapSafe().at("Stats").GetMapSafe();
 
-            auto operatorSize = false;
-            auto operatorRows = false;
             TString opType;
             TString opId = "0";
 
@@ -2755,31 +2787,8 @@ private:
                 }
             }
 
-            if (opName == "TableFullScan" && stats.contains("Table")) {
-                TString tablePath;
-                if (op.GetMapSafe().contains("Path")) {
-                    tablePath = op.GetMapSafe().at("Path").GetStringSafe();
-                } else if (op.GetMapSafe().contains("Table")) {
-                    tablePath = op.GetMapSafe().at("Table").GetStringSafe();
-                }
-                if (tablePath) {
-                    for (auto& opStat : stats.at("Table").GetArraySafe()) {
-                        if (opStat.IsMap()) {
-                            auto& opMap = opStat.GetMapSafe();
-                            if (opMap.contains("Path") && opMap.at("Path").GetStringSafe() == tablePath) {
-                                if (opMap.contains("ReadRows")) {
-                                    op["A-Rows"] = opMap.at("ReadRows").GetMapSafe().at("Sum").GetDouble();
-                                    operatorRows = true;
-                                }
-                                if (opMap.contains("ReadBytes")) {
-                                    op["A-Size"] = opMap.at("ReadBytes").GetMapSafe().at("Sum").GetDouble();
-                                    operatorRows = true;
-                                }
-                                break;
-                            }
-                        }
-                    }
-                }
+            if (opName == "TableFullScan" && ownTableStats) {
+                applyTableStats(*ownTableStats);
             }
 
             if (opType && stats.contains("Operator")) {
@@ -2839,6 +2848,10 @@ private:
             }
         }
 
+        if (opName == "TableFullScan" && !ownTableStats && inheritedTableStats) {
+            applyTableStats(*inheritedTableStats);
+        }
+
         // erase some redundant info from the table scan
         if (op.IsMap()) {
             auto& map = op.GetMapSafe();
diff --git a/ydb/core/kqp/ut/query/kqp_stats_ut.cpp b/ydb/core/kqp/ut/query/kqp_stats_ut.cpp
index 389e3f9db21..7487922a3b5 100644
--- a/ydb/core/kqp/ut/query/kqp_stats_ut.cpp
+++ b/ydb/core/kqp/ut/query/kqp_stats_ut.cpp
@@ -381,6 +381,59 @@ Y_UNIT_TEST(RequestUnitForExecute) {
     }
 }
 
+Y_UNIT_TEST(LegacySimplifiedPlanTableFullScanActualStats) {
+    NKikimrConfig::TAppConfig app;
+    app.MutableTableServiceConfig()->SetEnableNewRBO(false);
+
+    TKikimrRunner kikimr{TKikimrSettings(app)};
+    auto db = kikimr.GetTableClient();
+    auto session = db.CreateSession().GetValueSync().GetSession();
+
+    TExecDataQuerySettings settings;
+    settings.CollectQueryStats(ECollectQueryStatsMode::Full);
+
+    const auto execute = [&](const TString& query) {
+        auto result = session.ExecuteDataQuery(
+            query,
+            TTxControl::BeginTx().CommitTx(),
+            settings
+        ).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        NJson::TJsonValue plan;
+        UNIT_ASSERT_C(NJson::ReadJsonTree(result.GetQueryPlan(), &plan, true), result.GetQueryPlan());
+        return plan.GetMapSafe().at("SimplifiedPlan");
+    };
+
+    const auto fullScanPlan = execute(R"(
+        SELECT Key, Value1 FROM `/Root/TwoShard`;
+    )");
+    const auto fullScan = FindPlanNodeByKv(fullScanPlan, "Name", "TableFullScan");
+    UNIT_ASSERT_C(fullScan.IsDefined(), fullScanPlan);
+    UNIT_ASSERT_C(fullScan.GetMapSafe().contains("A-Rows"), fullScanPlan);
+    UNIT_ASSERT_VALUES_EQUAL_C(fullScan.GetMapSafe().at("A-Rows").GetDoubleSafe(), 6, fullScanPlan);
+    UNIT_ASSERT_C(fullScan.GetMapSafe().contains("A-Size"), fullScanPlan);
+    UNIT_ASSERT_C(fullScan.GetMapSafe().at("A-Size").GetDoubleSafe() > 0, fullScanPlan);
+
+    const auto limitedPlan = execute(R"(
+        SELECT Key, Value1 FROM `/Root/TwoShard` LIMIT 3;
+    )");
+    const auto limitNode = FindPlanNodeByKv(limitedPlan, "Node Type", "Limit");
+    UNIT_ASSERT_C(limitNode.IsDefined(), limitedPlan);
+
+    const auto limit = FindPlanNodeByKv(limitNode, "Name", "Limit");
+    UNIT_ASSERT_C(limit.IsDefined(), limitedPlan);
+    UNIT_ASSERT_C(limit.GetMapSafe().contains("A-Rows"), limitedPlan);
+    UNIT_ASSERT_VALUES_EQUAL_C(limit.GetMapSafe().at("A-Rows").GetDoubleSafe(), 3, limitedPlan);
+
+    const auto limitedScan = FindPlanNodeByKv(limitNode, "Name", "TableFullScan");
+    UNIT_ASSERT_C(limitedScan.IsDefined(), limitedPlan);
+    UNIT_ASSERT_C(limitedScan.GetMapSafe().contains("A-Rows"), limitedPlan);
+    UNIT_ASSERT_C(limitedScan.GetMapSafe().at("A-Rows").GetDoubleSafe() > 0, limitedPlan);
+    UNIT_ASSERT_C(limitedScan.GetMapSafe().contains("A-Size"), limitedPlan);
+    UNIT_ASSERT_C(limitedScan.GetMapSafe().at("A-Size").GetDoubleSafe() > 0, limitedPlan);
+}
+
 Y_UNIT_TEST(StatsProfile) {
     auto kikimr = DefaultKikimrRunner();
     auto db = kikimr.GetTableClient();