[NEW RBO] Recompute properties only when needed (#28818)

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

Author: pavelvelikhov

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

@@ -74,6 +74,10 @@ namespace {
 namespace NKikimr {
 namespace NKqp {
 
+TConstantFoldingStage::TConstantFoldingStage() {
+    Props = ERuleProperties::RequireParents |  ERuleProperties::RequireTypes;
+}
+
 void TConstantFoldingStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
     TVector<TExprNode::TPtr> lambdasWithConstExpr;
     bool foldUdfs = ctx.KqpCtx.Config->EnableFoldUdfs;

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

@@ -6,7 +6,7 @@
 namespace NKikimr {
 namespace NKqp {
 
-bool TSimplifiedRule::TestAndApply(std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) {
+bool ISimplifiedRule::TestAndApply(std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) {
 
     auto output = SimpleTestAndApply(input, ctx, props);
     if (input != output) {
@@ -16,6 +16,24 @@ bool TSimplifiedRule::TestAndApply(std::shared_ptr<IOperator> &input, TRBOContex
         return false;
     }
 }
+
+TRuleBasedStage::TRuleBasedStage(TVector<std::shared_ptr<IRule>> rules) : Rules(rules) {
+    for (auto & r : Rules) {
+        Props |= r->Props;
+    }
+}
+
+void ComputeRequiredProps(TOpRoot &root, ui32 props, TRBOContext &ctx) {
+    if (props & ERuleProperties::RequireParents) {
+        root.ComputeParents();
+    }
+    if (props & ERuleProperties::RequireTypes) {
+        if (root.ComputeTypes(ctx) != IGraphTransformer::TStatus::Ok) {
+            Y_ENSURE(false, "RBO type annotation failed");
+        }
+    }
+}
+
 /**
  * Run a rule-based stage
  *
@@ -31,10 +49,6 @@ void TRuleBasedStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
     bool fired = true;
     int nMatches = 0;
 
-    if (root.ComputeTypes(ctx) != IGraphTransformer::TStatus::Ok) {
-        Y_ENSURE(false, "RBO type annotation failed");
-    }
-
     while (fired && nMatches < 1000) {
         fired = false;
 
@@ -53,15 +67,7 @@ void TRuleBasedStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
                         root.Children[0] = op;
                     }
 
-                    if (rule->RequiresParentRecompute) {
-                        root.ComputeParents();
-                    }
-
-                    YQL_CLOG(TRACE, CoreDq) << "After rule:\n" << root.PlanToString(ctx.ExprCtx);
-
-                    if (root.ComputeTypes(ctx) != IGraphTransformer::TStatus::Ok) {
-                        Y_ENSURE(false, "RBO type annotation failed");
-                    }
+                    ComputeRequiredProps(root, Props, ctx);
 
                     nMatches++;
                     break;
@@ -82,22 +88,18 @@ TExprNode::TPtr TRuleBasedOptimizer::Optimize(TOpRoot &root, TExprContext &ctx)
 
     auto context = TRBOContext(KqpCtx,ctx,TypeCtx, RBOTypeAnnTransformer, FuncRegistry);
 
-    if (root.ComputeTypes(context) != IGraphTransformer::TStatus::Ok) {
-        Y_ENSURE(false, "RBO type annotation failed");
-    }
-    
     for (size_t idx = 0; idx < Stages.size(); idx++) {
         YQL_CLOG(TRACE, CoreDq) << "Running stage: " << idx;
         auto stage = Stages[idx];
+        ComputeRequiredProps(root, stage->Props, context);
         stage->RunStage(root, context);
         YQL_CLOG(TRACE, CoreDq) << "After stage:\n" << root.PlanToString(ctx);
-        if (root.ComputeTypes(context) != IGraphTransformer::TStatus::Ok) {
-            Y_ENSURE(false, "RBO type annotation failed");
-        }
     }
 
     YQL_CLOG(TRACE, CoreDq) << "New RBO finished, generating physical plan";
 
+    ComputeRequiredProps(root, ERuleProperties::RequireParents | ERuleProperties::RequireTypes, context);
+
     return ConvertToPhysical(root, context, TypeAnnTransformer, PeepholeTransformer);
 }
 } // namespace NKqp

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

@@ -9,6 +9,13 @@ namespace NKqp {
 
 using namespace NOpt;
 
+enum ERuleProperties: ui32 {
+    RequireParents = 0x01,
+    RequireTypes = 0x02,
+    RequireTableMeta = 0x04,
+    RequireCosts = 0x08,
+};
+
 /**
  * Interface for transformation rule:
  *
@@ -18,31 +25,31 @@ using namespace NOpt;
  */
 class IRule {
   public:
-    IRule(TString name, bool parentRecompute = true) : RuleName(name), RequiresParentRecompute(parentRecompute) {}
+    IRule(TString name) : RuleName(name) {}
+    IRule(TString name, ui32 props) : RuleName(name), Props(props) {}
 
     virtual bool TestAndApply(std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) = 0;
 
     virtual ~IRule() = default;
 
     TString RuleName;
-    bool RequiresParentRecompute = true;
+    ui32 Props;
 };
 
 /**
  * A Simplified rule does not alter the original subplan that it matched, but instead returns a new
  * subplan that replaces the old one
  */
-class TSimplifiedRule : public IRule {
+class ISimplifiedRule : public IRule {
   public:
-    TSimplifiedRule(TString name, bool parentRecompute = true) : IRule(name, parentRecompute) {}
+    ISimplifiedRule(TString name) : IRule(name) {}
+    ISimplifiedRule(TString name, ui32 props) : IRule(name, props) {}
 
     virtual std::shared_ptr<IOperator> SimpleTestAndApply(const std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) = 0;
 
     virtual bool TestAndApply(std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) override;
 };
 
-class TRuleBasedOptimizer;
-
 /**
  * Stage Interface
  *
@@ -53,27 +60,20 @@ class IRBOStage {
   public:
     virtual void RunStage(TOpRoot &root, TRBOContext &ctx) = 0;
     virtual ~IRBOStage() = default;
+    ui32 Props;
 };
 
 /**
  * Rule based stage is just a collection of rules
  */
 class TRuleBasedStage : public IRBOStage {
   public:
-    TRuleBasedStage(TVector<std::shared_ptr<IRule>> rules) : Rules(rules) {}
+    TRuleBasedStage(TVector<std::shared_ptr<IRule>> rules);
     virtual void RunStage(TOpRoot &root, TRBOContext &ctx) override;
 
     TVector<std::shared_ptr<IRule>> Rules;
 };
 
-/**
- * A Global stage uses its own logic to transform the entire plan
- */
-class ISinglePassStage : public IRBOStage {
-  public:
-    virtual void RunStage(TOpRoot &root, TRBOContext &ctx) override = 0;
-};
-
 /**
  * A rule based optimizer is a collection of rule-based and global stages
  */

ydb/core/kqp/opt/rbo/kqp_rbo_rules.h

@@ -27,7 +27,7 @@ class TExtractJoinExpressionsRule : public IRule {
  */
 class TInlineScalarSubplanRule : public IRule {
   public:
-    TInlineScalarSubplanRule() : IRule("Inline scalar subplan") {}
+    TInlineScalarSubplanRule() : IRule("Inline scalar subplan", ERuleProperties::RequireParents | ERuleProperties::RequireTypes) {}
 
     virtual bool TestAndApply(std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) override;
 };
@@ -36,9 +36,9 @@ class TInlineScalarSubplanRule : public IRule {
  * Push down a non-projecting map operator
  * Currently only pushes below joins that are immediately below
  */
-class TPushMapRule : public TSimplifiedRule {
+class TPushMapRule : public ISimplifiedRule {
   public:
-    TPushMapRule() : TSimplifiedRule("Push map operator") {}
+    TPushMapRule() : ISimplifiedRule("Push map operator", ERuleProperties::RequireParents) {}
 
     virtual std::shared_ptr<IOperator> SimpleTestAndApply(const std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) override;
 };
@@ -47,9 +47,9 @@ class TPushMapRule : public TSimplifiedRule {
  * Push down filter through joins, adding join conditions to the join operator and potentially
  * converting left join into inner join
  */
-class TPushFilterRule : public TSimplifiedRule {
+class TPushFilterRule : public ISimplifiedRule {
   public:
-    TPushFilterRule() : TSimplifiedRule("Push filter") {}
+    TPushFilterRule() : ISimplifiedRule("Push filter", ERuleProperties::RequireParents) {}
 
     virtual std::shared_ptr<IOperator> SimpleTestAndApply(const std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) override;
 };
@@ -59,7 +59,7 @@ class TPushFilterRule : public TSimplifiedRule {
  */
 class TAssignStagesRule : public IRule {
   public:
-    TAssignStagesRule() : IRule("Assign stages") {}
+    TAssignStagesRule() : IRule("Assign stages", ERuleProperties::RequireParents) {}
 
     virtual bool TestAndApply(std::shared_ptr<IOperator> &input, TRBOContext &ctx, TPlanProps &props) override;
 };
@@ -70,16 +70,18 @@ extern TRuleBasedStage RuleStage2;
 /**
  * Separate global stage to remove extra renames and project out unneeded columns
  */
-class TRenameStage : public ISinglePassStage {
+class TRenameStage : public IRBOStage {
   public:
+    TRenameStage();
     virtual void RunStage(TOpRoot &root, TRBOContext &ctx) override;
 };
 
 /**
  * Separate global constant folding stage
  */
-class TConstantFoldingStage : public ISinglePassStage {
+class TConstantFoldingStage : public IRBOStage {
   public:
+    TConstantFoldingStage();
     virtual void RunStage(TOpRoot &root, TRBOContext &ctx) override;
 };
 

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

@@ -123,6 +123,11 @@ struct TIntTUnitPairHash {
 /**
  * Global stage that removed unnecessary renames and unused columns
  */
+
+TRenameStage::TRenameStage() {
+    Props = ERuleProperties::RequireParents;
+}
+
 void TRenameStage::RunStage(TOpRoot &root, TRBOContext &ctx) {
 
     YQL_CLOG(TRACE, CoreDq) << "Before compute parents";