[NEW RBO] Fixed bug in statistics computation in Map (#30299)

Co-authored-by: Pavel Velikhov <pavelvelikhov@localhost.localdomain>

Author: pavelvelikhov

ydb/core/kqp/opt/rbo/kqp_rbo.cpp

@@ -56,6 +56,8 @@ void ComputeRequiredProps(TOpRoot &root, ui32 props, TRBOContext &ctx) {
 void TRuleBasedStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
     bool fired = true;
     int nMatches = 0;
+    bool needToLog = NYql::NLog::YqlLogger().NeedToLog(NYql::NLog::EComponent::CoreDq, NYql::NLog::ELevel::TRACE);
+
 
     while (fired && nMatches < 1000) {
         fired = false;
@@ -75,7 +77,7 @@ void TRuleBasedStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
                         root.SetInput(op);
                     }
 
-                    if (rule->LogRule) {
+                    if (needToLog && rule->LogRule) {
                         YQL_CLOG(TRACE, CoreDq) << "Plan after applying rule:\n" << root.PlanToString(ctx.ExprCtx);
                     }
 
@@ -96,23 +98,34 @@ void TRuleBasedStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
 }
 
 TExprNode::TPtr TRuleBasedOptimizer::Optimize(TOpRoot &root, TExprContext &ctx) {
-    YQL_CLOG(TRACE, CoreDq) << "Original plan:\n" << root.PlanToString(ctx);
+    bool needToLog = NYql::NLog::YqlLogger().NeedToLog(NYql::NLog::EComponent::CoreDq, NYql::NLog::ELevel::TRACE);
+
+    if (needToLog) {
+        YQL_CLOG(TRACE, CoreDq) << "Original plan:\n" << root.PlanToString(ctx);
+    }
 
     auto context = TRBOContext(KqpCtx, ctx, TypeCtx, RBOTypeAnnTransformer, FuncRegistry);
 
     for (size_t idx = 0; idx < Stages.size(); idx++) {
         auto stage = Stages[idx];
         YQL_CLOG(TRACE, CoreDq) << "Running stage: " << stage->StageName;
         ComputeRequiredProps(root, stage->Props, context);
+        if (needToLog) {
+            YQL_CLOG(TRACE, CoreDq) << "Before stage:\n" << root.PlanToString(ctx, EPrintPlanOptions::PrintFullMetadata | EPrintPlanOptions::PrintBasicStatistics);
+        }
         stage->RunStage(root, context);
-        YQL_CLOG(TRACE, CoreDq) << "After stage:\n" << root.PlanToString(ctx);
+        if (needToLog) {
+            YQL_CLOG(TRACE, CoreDq) << "After stage:\n" << root.PlanToString(ctx, EPrintPlanOptions::PrintFullMetadata | EPrintPlanOptions::PrintBasicStatistics);
+        }
     }
 
     YQL_CLOG(TRACE, CoreDq) << "New RBO finished, generating physical plan";
 
     auto convertProps = ERuleProperties::RequireParents | ERuleProperties::RequireTypes | ERuleProperties::RequireStatistics;
     ComputeRequiredProps(root, convertProps, context);
-    YQL_CLOG(TRACE, CoreDq) << "Final plan before generation:\n" << root.PlanToString(ctx, EPrintPlanOptions::PrintFullMetadata | EPrintPlanOptions::PrintBasicStatistics);
+    if (needToLog) {
+        YQL_CLOG(TRACE, CoreDq) << "Final plan before generation:\n" << root.PlanToString(ctx, EPrintPlanOptions::PrintFullMetadata | EPrintPlanOptions::PrintBasicStatistics);
+    }
 
     return ConvertToPhysical(root, context, TypeAnnTransformer, PeepholeTransformer);
 }

ydb/core/kqp/opt/rbo/kqp_rbo_compute_statistics.cpp

@@ -222,7 +222,7 @@ void TOpMap::ComputeMetadata(TRBOContext & ctx, TPlanProps & planProps) {
     if (Project) {
         Props.Metadata->ColumnsCount = MapElements.size();
     } else {
-        Props.Metadata->ColumnsCount += MapElements.size();
+        Props.Metadata->ColumnsCount += inputMetadata.ColumnsCount + MapElements.size();
     }
 
     auto renames = GetRenamesWithTransforms(planProps);
@@ -282,7 +282,14 @@ void TOpMap::ComputeStatistics(TRBOContext & ctx, TPlanProps & planProps) {
     int inputColumnsCount = GetInput()->Props.Metadata->ColumnsCount;
     if (Props.Metadata->ColumnsCount != inputColumnsCount) {
         double inputDataSize = Props.Statistics->DataSize;
-        Props.Statistics->DataSize = inputDataSize * Props.Metadata->ColumnsCount / (double)inputColumnsCount;
+        if (inputColumnsCount!=0) {
+            Props.Statistics->DataSize = inputDataSize * Props.Metadata->ColumnsCount / (double)inputColumnsCount;
+        }
+        // Input may have 0 columns (e.g. EmptySource), in such case the data size depends on the number of records
+        // and the number of columns in the output. We just assume each column contains 8 bytes
+        else {
+            Props.Statistics->DataSize = Props.Statistics->RecordsCount * Props.Metadata->ColumnsCount * 8;
+        }
     }
 }
 

ydb/core/kqp/ut/rbo/kqp_rbo_pg_ut.cpp

@@ -229,8 +229,6 @@ Y_UNIT_TEST_SUITE(KqpRboPg) {
         NKikimrConfig::TAppConfig appConfig;
         appConfig.MutableTableServiceConfig()->SetEnableNewRBO(true);
         appConfig.MutableTableServiceConfig()->SetEnableFallbackToYqlOptimizer(false);
-        // FIXME: Temporarily disabled cost-based optimizer, it break on this query
-        appConfig.MutableTableServiceConfig()->SetDefaultCostBasedOptimizationLevel(0);
 
         TKikimrRunner kikimr(NKqp::TKikimrSettings(appConfig).SetWithSampleTables(false));
         auto db = kikimr.GetTableClient();