Support progress stats with basic stats mode (#16773)

Author: shnikd

ydb/core/grpc_services/query/rpc_execute_query.cpp

@@ -386,7 +386,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
         if (NeedReportStats(*Request_->GetProtoRequest())) {
             if (record.HasQueryStats()) {
                 FillQueryStats(*response.mutable_exec_stats(), record.GetQueryStats());
-                response.mutable_exec_stats()->set_query_plan(NKqp::SerializeAnalyzePlan(record.GetQueryStats()));
+                response.mutable_exec_stats()->set_query_plan(record.GetQueryPlan());
             }
         }
 

ydb/core/kqp/compute_actor/kqp_compute_actor_factory.cpp

@@ -170,6 +170,7 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
         runtimeSettings.ExtraMemoryAllocationPool = args.MemoryPool;
         runtimeSettings.UseSpilling = args.WithSpilling;
         runtimeSettings.StatsMode = args.StatsMode;
+        runtimeSettings.WithProgressStats = args.WithProgressStats;
 
         if (runtimeSettings.UseSpilling) {
             args.Task->SetEnableSpilling(runtimeSettings.UseSpilling);

ydb/core/kqp/compute_actor/kqp_compute_actor_factory.h

@@ -122,6 +122,7 @@ struct IKqpNodeComputeActorFactory {
         const NKikimr::NKqp::NRm::EKqpMemoryPool MemoryPool;
         const bool WithSpilling;
         const NYql::NDqProto::EDqStatsMode StatsMode;
+        const bool WithProgressStats;
         const TInstant& Deadline;
         const bool ShareMailbox;
         const TMaybe<NYql::NDqProto::TRlPath>& RlPath;

ydb/core/kqp/executer_actor/kqp_executer_impl.h

@@ -1618,6 +1618,7 @@ class TKqpExecuterBase : public TActor<TDerived> {
             .UserToken = UserToken,
             .Deadline = Deadline.GetOrElse(TInstant::Zero()),
             .StatsMode = Request.StatsMode,
+            .WithProgressStats = Request.ProgressStatsPeriod != TDuration::Zero(),
             .RlPath = Request.RlPath,
             .ExecuterSpan =  ExecuterSpan,
             .ResourcesSnapshot = std::move(ResourcesSnapshot),

ydb/core/kqp/executer_actor/kqp_executer_stats.cpp

@@ -1326,91 +1326,116 @@ void TQueryExecutionStats::ExportAggAsyncBufferStats(TAsyncBufferStats& data, NY
 void TQueryExecutionStats::ExportExecStats(NYql::NDqProto::TDqExecutionStats& stats) {
 
     THashMap<ui32, NDqProto::TDqStageStats*> protoStages;
-    for (auto& [stageId, stagetype] : TasksGraph->GetStagesInfo()) {
-        protoStages.emplace(stageId.StageId, GetOrCreateStageStats(stageId, *TasksGraph, stats));
+
+    if (CollectFullStats(StatsMode)) {
+        for (auto& [stageId, stagetype] : TasksGraph->GetStagesInfo()) {
+            protoStages.emplace(stageId.StageId, GetOrCreateStageStats(stageId, *TasksGraph, stats));
+        }
     }
 
-    for (auto& [stageId, stageStat] : StageStats) {
-        auto& stageStats = *protoStages[stageStat.StageId.StageId];
-        stageStats.SetTotalTasksCount(stageStat.Task2Index.size());
-        stageStats.SetFinishedTasksCount(stageStat.FinishedCount);
-
-        stageStats.SetBaseTimeMs(BaseTimeMs);
-        stageStat.CpuTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableCpuTimeUs());
-        ExportAggStats(stageStat.SourceCpuTimeUs, *stageStats.MutableSourceCpuTimeUs());
-        stageStat.MaxMemoryUsage.ExportAggStats(BaseTimeMs, *stageStats.MutableMaxMemoryUsage());
-
-        ExportAggStats(stageStat.InputRows, *stageStats.MutableInputRows());
-        ExportAggStats(stageStat.InputBytes, *stageStats.MutableInputBytes());
-        ExportAggStats(stageStat.OutputRows, *stageStats.MutableOutputRows());
-        ExportAggStats(stageStat.OutputBytes, *stageStats.MutableOutputBytes());
-        ExportAggStats(stageStat.ResultRows, *stageStats.MutableResultRows());
-        ExportAggStats(stageStat.ResultBytes, *stageStats.MutableResultBytes());
-        ExportAggStats(stageStat.IngressRows, *stageStats.MutableIngressRows());
-        ExportAggStats(stageStat.IngressBytes, *stageStats.MutableIngressBytes());
-        ExportAggStats(stageStat.IngressDecompressedBytes, *stageStats.MutableIngressDecompressedBytes());
-        ExportAggStats(stageStat.EgressRows, *stageStats.MutableEgressRows());
-        ExportAggStats(stageStat.EgressBytes, *stageStats.MutableEgressBytes());
-
-        ExportOffsetAggStats(stageStat.StartTimeMs, *stageStats.MutableStartTimeMs(), BaseTimeMs);
-        ExportOffsetAggStats(stageStat.FinishTimeMs, *stageStats.MutableFinishTimeMs(), BaseTimeMs);
-        ExportAggStats(stageStat.DurationUs, *stageStats.MutableDurationUs());
-        stageStat.WaitInputTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableWaitInputTimeUs());
-        stageStat.WaitOutputTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableWaitOutputTimeUs());
-
-        stageStat.SpillingComputeBytes.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingComputeBytes());
-        stageStat.SpillingChannelBytes.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingChannelBytes());
-        stageStat.SpillingComputeTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingComputeTimeUs());
-        stageStat.SpillingChannelTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingChannelTimeUs());
-
-        FillStageDurationUs(stageStats);
+    std::unordered_map<TString, NYql::NDqProto::TDqTableStats*> currentTableStats;
 
+    for (auto& [stageId, stageStat] : StageStats) {
         for (auto& [path, t] : stageStat.Tables) {
-            auto& table = *stageStats.AddTables();
-            table.SetTablePath(path);
-            ExportAggStats(t.ReadRows, *table.MutableReadRows());
-            ExportAggStats(t.ReadBytes, *table.MutableReadBytes());
-            ExportAggStats(t.WriteRows, *table.MutableWriteRows());
-            ExportAggStats(t.WriteBytes, *table.MutableWriteBytes());
-            ExportAggStats(t.EraseRows, *table.MutableEraseRows());
-            ExportAggStats(t.EraseBytes, *table.MutableEraseBytes());
-            table.SetAffectedPartitions(ExportAggStats(t.AffectedPartitions));
-        }
-        for (auto& [id, i] : stageStat.Ingress) {
-            ExportAggAsyncBufferStats(i, (*stageStats.MutableIngress())[id]);
-        }
-        for (auto& [id, i] : stageStat.Input) {
-            ExportAggAsyncBufferStats(i, (*stageStats.MutableInput())[id]);
-        }
-        for (auto& [id, o] : stageStat.Output) {
-            ExportAggAsyncBufferStats(o, (*stageStats.MutableOutput())[id]);
-        }
-        for (auto& [id, e] : stageStat.Egress) {
-            ExportAggAsyncBufferStats(e, (*stageStats.MutableEgress())[id]);
-        }
-        for (auto& [id, j] : stageStat.Joins) {
-            auto& joinStat = (*stageStats.MutableOperatorJoin())[id];
-            joinStat.SetOperatorId(id);
-            ExportAggStats(j.Bytes, *joinStat.MutableBytes());
-            ExportAggStats(j.Rows, *joinStat.MutableRows());
-        }
-        for (auto& [id, f] : stageStat.Filters) {
-            auto& filterStat = (*stageStats.MutableOperatorFilter())[id];
-            filterStat.SetOperatorId(id);
-            ExportAggStats(f.Bytes, *filterStat.MutableBytes());
-            ExportAggStats(f.Rows, *filterStat.MutableRows());
+            NYql::NDqProto::TDqTableStats* tableAggr = nullptr;
+            if (auto it = currentTableStats.find(path); it != currentTableStats.end()) {
+                tableAggr = it->second;
+            } else {
+                tableAggr = stats.AddTables();
+                tableAggr->SetTablePath(path);
+                currentTableStats.emplace(path, tableAggr);
+            }
+
+            tableAggr->SetReadRows(tableAggr->GetReadRows() + ExportAggStats(t.ReadRows));
+            tableAggr->SetReadBytes(tableAggr->GetReadBytes() + ExportAggStats(t.ReadBytes));
+            tableAggr->SetWriteRows(tableAggr->GetWriteRows() + ExportAggStats(t.WriteRows));
+            tableAggr->SetWriteBytes(tableAggr->GetWriteBytes() + ExportAggStats(t.WriteBytes));
+            tableAggr->SetEraseRows(tableAggr->GetEraseRows() + ExportAggStats(t.EraseRows));
+            tableAggr->SetAffectedPartitions(tableAggr->GetAffectedPartitions() + ExportAggStats(t.AffectedPartitions));
+
         }
-        for (auto& [id, a] : stageStat.Aggregations) {
-            auto& aggrStat = (*stageStats.MutableOperatorAggregation())[id];
-            aggrStat.SetOperatorId(id);
-            ExportAggStats(a.Bytes, *aggrStat.MutableBytes());
-            ExportAggStats(a.Rows, *aggrStat.MutableRows());
+
+
+        if (CollectFullStats(StatsMode)) {
+            auto& stageStats = *protoStages[stageStat.StageId.StageId];
+            stageStats.SetTotalTasksCount(stageStat.Task2Index.size());
+            stageStats.SetFinishedTasksCount(stageStat.FinishedCount);
+
+            stageStats.SetBaseTimeMs(BaseTimeMs);
+            stageStat.CpuTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableCpuTimeUs());
+            ExportAggStats(stageStat.SourceCpuTimeUs, *stageStats.MutableSourceCpuTimeUs());
+            stageStat.MaxMemoryUsage.ExportAggStats(BaseTimeMs, *stageStats.MutableMaxMemoryUsage());
+
+            ExportAggStats(stageStat.InputRows, *stageStats.MutableInputRows());
+            ExportAggStats(stageStat.InputBytes, *stageStats.MutableInputBytes());
+            ExportAggStats(stageStat.OutputRows, *stageStats.MutableOutputRows());
+            ExportAggStats(stageStat.OutputBytes, *stageStats.MutableOutputBytes());
+            ExportAggStats(stageStat.ResultRows, *stageStats.MutableResultRows());
+            ExportAggStats(stageStat.ResultBytes, *stageStats.MutableResultBytes());
+            ExportAggStats(stageStat.IngressRows, *stageStats.MutableIngressRows());
+            ExportAggStats(stageStat.IngressBytes, *stageStats.MutableIngressBytes());
+            ExportAggStats(stageStat.IngressDecompressedBytes, *stageStats.MutableIngressDecompressedBytes());
+            ExportAggStats(stageStat.EgressRows, *stageStats.MutableEgressRows());
+            ExportAggStats(stageStat.EgressBytes, *stageStats.MutableEgressBytes());
+
+            ExportOffsetAggStats(stageStat.StartTimeMs, *stageStats.MutableStartTimeMs(), BaseTimeMs);
+            ExportOffsetAggStats(stageStat.FinishTimeMs, *stageStats.MutableFinishTimeMs(), BaseTimeMs);
+            ExportAggStats(stageStat.DurationUs, *stageStats.MutableDurationUs());
+            stageStat.WaitInputTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableWaitInputTimeUs());
+            stageStat.WaitOutputTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableWaitOutputTimeUs());
+
+            stageStat.SpillingComputeBytes.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingComputeBytes());
+            stageStat.SpillingChannelBytes.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingChannelBytes());
+            stageStat.SpillingComputeTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingComputeTimeUs());
+            stageStat.SpillingChannelTimeUs.ExportAggStats(BaseTimeMs, *stageStats.MutableSpillingChannelTimeUs());
+
+            FillStageDurationUs(stageStats);
+
+            for (auto& [path, t] : stageStat.Tables) {
+                auto& table = *stageStats.AddTables();
+                table.SetTablePath(path);
+                ExportAggStats(t.ReadRows, *table.MutableReadRows());
+                ExportAggStats(t.ReadBytes, *table.MutableReadBytes());
+                ExportAggStats(t.WriteRows, *table.MutableWriteRows());
+                ExportAggStats(t.WriteBytes, *table.MutableWriteBytes());
+                ExportAggStats(t.EraseRows, *table.MutableEraseRows());
+                ExportAggStats(t.EraseBytes, *table.MutableEraseBytes());
+                table.SetAffectedPartitions(ExportAggStats(t.AffectedPartitions));
+            }
+            for (auto& [id, i] : stageStat.Ingress) {
+                ExportAggAsyncBufferStats(i, (*stageStats.MutableIngress())[id]);
+            }
+            for (auto& [id, i] : stageStat.Input) {
+                ExportAggAsyncBufferStats(i, (*stageStats.MutableInput())[id]);
+            }
+            for (auto& [id, o] : stageStat.Output) {
+                ExportAggAsyncBufferStats(o, (*stageStats.MutableOutput())[id]);
+            }
+            for (auto& [id, e] : stageStat.Egress) {
+                ExportAggAsyncBufferStats(e, (*stageStats.MutableEgress())[id]);
+            }
+            for (auto& [id, j] : stageStat.Joins) {
+                auto& joinStat = (*stageStats.MutableOperatorJoin())[id];
+                joinStat.SetOperatorId(id);
+                ExportAggStats(j.Bytes, *joinStat.MutableBytes());
+                ExportAggStats(j.Rows, *joinStat.MutableRows());
+            }
+            for (auto& [id, f] : stageStat.Filters) {
+                auto& filterStat = (*stageStats.MutableOperatorFilter())[id];
+                filterStat.SetOperatorId(id);
+                ExportAggStats(f.Bytes, *filterStat.MutableBytes());
+                ExportAggStats(f.Rows, *filterStat.MutableRows());
+            }
+            for (auto& [id, a] : stageStat.Aggregations) {
+                auto& aggrStat = (*stageStats.MutableOperatorAggregation())[id];
+                aggrStat.SetOperatorId(id);
+                ExportAggStats(a.Bytes, *aggrStat.MutableBytes());
+                ExportAggStats(a.Rows, *aggrStat.MutableRows());
+            }
         }
     }
 
-    for (const auto& [_, tableStats] : TableStats) {
-        stats.AddTables()->CopyFrom(*tableStats);
-    }
+    stats.SetDurationUs(TInstant::Now().MicroSeconds() - StartTs.MicroSeconds());
 }
 
 void TQueryExecutionStats::AdjustExternalAggr(NYql::NDqProto::TDqExternalAggrStats& stats) {

ydb/core/kqp/executer_actor/kqp_planner.cpp

@@ -89,6 +89,7 @@ TKqpPlanner::TKqpPlanner(TKqpPlanner::TArgs&& args)
     , UserToken(args.UserToken)
     , Deadline(args.Deadline)
     , StatsMode(args.StatsMode)
+    , WithProgressStats(args.WithProgressStats)
     , RlPath(args.RlPath)
     , ResourcesSnapshot(std::move(args.ResourcesSnapshot))
     , ExecuterSpan(args.ExecuterSpan)
@@ -224,6 +225,7 @@ std::unique_ptr<TEvKqpNode::TEvStartKqpTasksRequest> TKqpPlanner::SerializeReque
     }
 
     request.MutableRuntimeSettings()->SetStatsMode(GetDqStatsMode(StatsMode));
+    request.MutableRuntimeSettings()->SetWithProgressStats(WithProgressStats);
     request.SetStartAllOrFail(true);
     request.MutableRuntimeSettings()->SetExecType(NYql::NDqProto::TComputeRuntimeSettings::DATA);
     request.MutableRuntimeSettings()->SetUseSpilling(TasksGraph.GetMeta().AllowWithSpilling);
@@ -499,6 +501,7 @@ TString TKqpPlanner::ExecuteDataComputeTask(ui64 taskId, ui32 computeTasksSize)
         .MemoryPool = NRm::EKqpMemoryPool::DataQuery,
         .WithSpilling = TasksGraph.GetMeta().AllowWithSpilling,
         .StatsMode = GetDqStatsMode(StatsMode),
+        .WithProgressStats = WithProgressStats,
         .Deadline = Deadline,
         .ShareMailbox = (computeTasksSize <= 1),
         .RlPath = Nothing(),

ydb/core/kqp/executer_actor/kqp_planner.h

@@ -48,6 +48,7 @@ class TKqpPlanner {
         const TIntrusiveConstPtr<NACLib::TUserToken>& UserToken;
         const TInstant Deadline;
         const Ydb::Table::QueryStatsCollection::Mode& StatsMode;
+        const bool WithProgressStats;
         const TMaybe<NKikimrKqp::TRlPath>& RlPath;
         NWilson::TSpan& ExecuterSpan;
         TVector<NKikimrKqp::TKqpNodeResources> ResourcesSnapshot;
@@ -110,6 +111,7 @@ class TKqpPlanner {
     const TIntrusiveConstPtr<NACLib::TUserToken> UserToken;
     const TInstant Deadline;
     const Ydb::Table::QueryStatsCollection::Mode StatsMode;
+    const bool WithProgressStats;
     const TMaybe<NKikimrKqp::TRlPath> RlPath;
     THashSet<ui32> TrackingNodes;
     TVector<NKikimrKqp::TKqpNodeResources> ResourcesSnapshot;

ydb/core/kqp/node_service/kqp_node_service.cpp

@@ -274,6 +274,7 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
                 .MemoryPool = memoryPool,
                 .WithSpilling = msgRtSettings.GetUseSpilling(),
                 .StatsMode = msgRtSettings.GetStatsMode(),
+                .WithProgressStats = msgRtSettings.GetWithProgressStats(),
                 .Deadline = TInstant(),
                 .ShareMailbox = false,
                 .RlPath = rlPath,

ydb/core/kqp/session_actor/kqp_session_actor.cpp

@@ -1567,14 +1567,13 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             }
 
             if (QueryState->ReportStats()) {
-                if (QueryState->GetStatsMode() >= Ydb::Table::QueryStatsCollection::STATS_COLLECTION_FULL) {
-                    NKqpProto::TKqpStatsQuery& stats = *ev->Get()->Record.MutableQueryStats();
-                    NKqpProto::TKqpStatsQuery executionStats;
-                    executionStats.Swap(&stats);
-                    stats = QueryState->QueryStats.ToProto();
-                    stats.MutableExecutions()->MergeFrom(executionStats.GetExecutions());
-                    ev->Get()->Record.SetQueryPlan(SerializeAnalyzePlan(stats, QueryState->UserRequestContext->PoolId));
-                }
+                NKqpProto::TKqpStatsQuery& stats = *ev->Get()->Record.MutableQueryStats();
+                NKqpProto::TKqpStatsQuery executionStats;
+                executionStats.Swap(&stats);
+                stats = QueryState->QueryStats.ToProto();
+                stats.MutableExecutions()->MergeFrom(executionStats.GetExecutions());
+                ev->Get()->Record.SetQueryPlan(SerializeAnalyzePlan(stats, QueryState->UserRequestContext->PoolId));
+                stats.SetDurationUs((TInstant::Now() - QueryState->StartTime).MicroSeconds());
             }
 
             LOG_D("Forwarded TEvExecuterProgress to " << QueryState->RequestActorId);

ydb/library/yql/dq/actors/compute/dq_compute_actor.h

@@ -256,6 +256,8 @@ struct TComputeRuntimeSettings {
 
     i64 AsyncInputPushLimit = std::numeric_limits<i64>::max();
 
+    bool WithProgressStats = false;
+
     inline bool CollectNone() const {
         return StatsMode <= NDqProto::DQ_STATS_MODE_NONE;
     }