[CBO] UnionAll stats propagation fix (#20746)

Author: pashandor789

ydb/library/yql/dq/opt/dq_opt_stat.cpp

@@ -650,6 +650,8 @@ void InferStatisticsForExtendBase(const TExprNode::TPtr& input, TTypeAnnotationC
     for (const auto& input: input->Children()) {
         if (auto inputStats = typeCtx->GetStats(input.Get())) {
             stats->Nrows += inputStats->Nrows;
+            stats->ByteSize += inputStats->ByteSize;
+            stats->Cost += inputStats->Cost;
         }
     }
 
@@ -849,6 +851,22 @@ void InferStatisticsForDqMerge(const TExprNode::TPtr& input, TTypeAnnotationCont
     typeCtx->SetStats(merge.Raw(), newStats);
 }
 
+void InferStatisticsForUnionAll(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx) {
+    auto inputNode = TExprBase(input);
+    auto unionAll = inputNode.Cast<TCoUnionAll>();
+
+    auto stats = std::make_shared<TOptimizerStatistics>();
+    for (const auto& input: input->Children()) {
+        if (auto inputStats = typeCtx->GetStats(input.Get())) {
+            stats->Nrows += inputStats->Nrows;
+            stats->ByteSize += inputStats->ByteSize;
+            stats->Cost += inputStats->Cost;
+        }
+    }
+
+    typeCtx->SetStats(inputNode.Raw(), std::move(stats));
+}
+
 /**
  * Just update the sorted order with alias
  */
@@ -865,8 +883,6 @@ void InferStatisticsForDqPhyCrossJoin(const TExprNode::TPtr& input, TTypeAnnotat
     typeCtx->SetStats(cross.Raw(), outputStats);
 }
 
-
-
 std::shared_ptr<TOptimizerStatistics> RemoveSorting(const std::shared_ptr<TOptimizerStatistics>& stats) {
     if (stats->SortingOrderings.HasState()) {
         auto newStats = *stats;

ydb/library/yql/dq/opt/dq_opt_stat.h

@@ -29,6 +29,7 @@ void InferStatisticsForTopBase(const TExprNode::TPtr& input, TTypeAnnotationCont
 void InferStatisticsForSortBase(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx);
 bool InferStatisticsForListParam(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx);
 void InferStatisticsForEquiJoin(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx);
+void InferStatisticsForUnionAll(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx);
 std::shared_ptr<TOptimizerStatistics> RemoveSorting(const std::shared_ptr<TOptimizerStatistics>& stats);
 std::shared_ptr<TOptimizerStatistics> RemoveSorting(const std::shared_ptr<TOptimizerStatistics>& stats, const TExprNode::TPtr& input);
 

ydb/library/yql/dq/opt/dq_opt_stat_transformer_base.cpp

@@ -139,6 +139,9 @@ bool TDqStatisticsTransformerBase::BeforeLambdas(const TExprNode::TPtr& input, T
     else if (auto sortBase = TMaybeNode<TCoSortBase>(input)) {
         InferStatisticsForSortBase(input, TypeCtx);
     }
+    else if (TCoUnionAll::Match(input.Get())) {
+        InferStatisticsForUnionAll(input, TypeCtx);
+    }
     else {
         matched = false;
     }