From 041d36595a4c9bcddccf42823ae67a8f70cd8a41 Mon Sep 17 00:00:00 2001 From: Pavel Ivanov Date: Mon, 7 Jul 2025 15:42:48 +0000 Subject: [PATCH 1/2] [CBO] UnionAll stats propagation fix --- ydb/library/yql/dq/opt/dq_opt_stat.cpp | 18 ++++++++++++++++-- ydb/library/yql/dq/opt/dq_opt_stat.h | 1 + .../dq/opt/dq_opt_stat_transformer_base.cpp | 3 +++ 3 files changed, 20 insertions(+), 2 deletions(-) diff --git a/ydb/library/yql/dq/opt/dq_opt_stat.cpp b/ydb/library/yql/dq/opt/dq_opt_stat.cpp index ea8fdfba5d9b..624f47104a19 100644 --- a/ydb/library/yql/dq/opt/dq_opt_stat.cpp +++ b/ydb/library/yql/dq/opt/dq_opt_stat.cpp @@ -650,6 +650,7 @@ 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; } } @@ -849,6 +850,21 @@ 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(); + + auto stats = std::make_shared(); + for (const auto& input: input->Children()) { + if (auto inputStats = typeCtx->GetStats(input.Get())) { + stats->Nrows += inputStats->Nrows; + stats->ByteSize += inputStats->ByteSize; + } + } + + typeCtx->SetStats(inputNode.Raw(), std::move(stats)); +} + /** * Just update the sorted order with alias */ @@ -865,8 +881,6 @@ void InferStatisticsForDqPhyCrossJoin(const TExprNode::TPtr& input, TTypeAnnotat typeCtx->SetStats(cross.Raw(), outputStats); } - - std::shared_ptr RemoveSorting(const std::shared_ptr& stats) { if (stats->SortingOrderings.HasState()) { auto newStats = *stats; diff --git a/ydb/library/yql/dq/opt/dq_opt_stat.h b/ydb/library/yql/dq/opt/dq_opt_stat.h index e6d0966f742a..fc9a147c7c95 100644 --- a/ydb/library/yql/dq/opt/dq_opt_stat.h +++ b/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 RemoveSorting(const std::shared_ptr& stats); std::shared_ptr RemoveSorting(const std::shared_ptr& stats, const TExprNode::TPtr& input); diff --git a/ydb/library/yql/dq/opt/dq_opt_stat_transformer_base.cpp b/ydb/library/yql/dq/opt/dq_opt_stat_transformer_base.cpp index 54c00c11e989..eb876bc9e102 100644 --- a/ydb/library/yql/dq/opt/dq_opt_stat_transformer_base.cpp +++ b/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(input)) { InferStatisticsForSortBase(input, TypeCtx); } + else if (TCoUnionAll::Match(input.Get())) { + InferStatisticsForUnionAll(input, TypeCtx); + } else { matched = false; } From 4810093fc28fbf56d33c8495273930f9dd9707d3 Mon Sep 17 00:00:00 2001 From: Pavel Ivanov Date: Mon, 7 Jul 2025 18:52:42 +0000 Subject: [PATCH 2/2] [CBO] UnionAll stats propagation fix --- ydb/library/yql/dq/opt/dq_opt_stat.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ydb/library/yql/dq/opt/dq_opt_stat.cpp b/ydb/library/yql/dq/opt/dq_opt_stat.cpp index 624f47104a19..2a1c3f8b26d1 100644 --- a/ydb/library/yql/dq/opt/dq_opt_stat.cpp +++ b/ydb/library/yql/dq/opt/dq_opt_stat.cpp @@ -651,6 +651,7 @@ void InferStatisticsForExtendBase(const TExprNode::TPtr& input, TTypeAnnotationC if (auto inputStats = typeCtx->GetStats(input.Get())) { stats->Nrows += inputStats->Nrows; stats->ByteSize += inputStats->ByteSize; + stats->Cost += inputStats->Cost; } } @@ -859,6 +860,7 @@ void InferStatisticsForUnionAll(const TExprNode::TPtr& input, TTypeAnnotationCon if (auto inputStats = typeCtx->GetStats(input.Get())) { stats->Nrows += inputStats->Nrows; stats->ByteSize += inputStats->ByteSize; + stats->Cost += inputStats->Cost; } }