[Enhancement] Add enable_optimizer_rule_debug session variable for rule failure diagnosis (#63693)

stephen-shelby · mergify[bot] · commit 76ff0cfe71a6 · 2025-10-23T11:02:31.000Z
Signed-off-by: stephen <stephen5217@163.com> (cherry picked from commit 6a318c2)
diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/SessionVariable.java b/fe/fe-core/src/main/java/com/starrocks/qe/SessionVariable.java
@@ -755,6 +755,8 @@ public static MaterializedViewRewriteMode parse(String str) {
 
     public static final String ENABLE_PLAN_VALIDATION = "enable_plan_validation";
 
+    public static final String ENABLE_OPTIMIZER_RULE_DEBUG = "enable_optimizer_rule_debug";
+
     public static final String ENABLE_STRICT_TYPE = "enable_strict_type";
 
     public static final String PARTIAL_UPDATE_MODE = "partial_update_mode";
@@ -2332,6 +2334,9 @@ public long getConnectorSinkTargetMaxFileSize() {
     @VarAttr(name = ENABLE_PLAN_VALIDATION, flag = VariableMgr.INVISIBLE)
     private boolean enablePlanValidation = true;
 
+    @VarAttr(name = ENABLE_OPTIMIZER_RULE_DEBUG)
+    private boolean enableOptimizerRuleDebug = false;
+
     @VarAttr(name = SCAN_OR_TO_UNION_LIMIT, flag = VariableMgr.INVISIBLE)
     private int scanOrToUnionLimit = 4;
 
@@ -4570,6 +4575,14 @@ public void setEnablePlanValidation(boolean val) {
         this.enablePlanValidation = val;
     }
 
+    public boolean enableOptimizerRuleDebug() {
+        return this.enableOptimizerRuleDebug;
+    }
+
+    public void setEnableOptimizerRuleDebug(boolean enableOptimizerRuleDebug) {
+        this.enableOptimizerRuleDebug = enableOptimizerRuleDebug;
+    }
+
     public boolean isCboPruneSubfield() {
         return cboPruneSubfield;
     }
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/QueryOptimizer.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/QueryOptimizer.java
@@ -124,7 +124,6 @@
 import com.starrocks.sql.optimizer.task.TaskScheduler;
 import com.starrocks.sql.optimizer.validate.MVRewriteValidator;
 import com.starrocks.sql.optimizer.validate.OptExpressionValidator;
-import com.starrocks.sql.optimizer.validate.PlanValidator;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -303,8 +302,9 @@ private OptExpression optimizeByCost(ConnectContext connectContext,
         mvScan.stream().map(scan -> ((MaterializedView) scan.getTable()).getDbId()).forEach(currentSqlDbIds::add);
 
         try (Timer ignored = Tracers.watchScope("PlanValidate")) {
+            rootTaskContext.getPlanValidator().enableAllCheckers();
             // valid the final plan
-            PlanValidator.getInstance().validatePlan(finalPlan, rootTaskContext);
+            rootTaskContext.getPlanValidator().validatePlan(finalPlan, rootTaskContext);
             // validate mv and log tracer if needed
             MVRewriteValidator mvRewriteValidator = new MVRewriteValidator(allLogicalOlapScanOperators);
             mvRewriteValidator.validateMV(connectContext, finalPlan, rootTaskContext);
@@ -550,6 +550,9 @@ private OptExpression logicalRuleRewrite(
         rootTaskContext.setRequiredColumns(requiredColumns.clone());
         scheduler.rewriteOnce(tree, rootTaskContext, RuleSet.PRUNE_COLUMNS_RULES);
 
+        // Enable input dependencies checker after column pruning for rule transformation validation
+        rootTaskContext.getPlanValidator().enableInputDependenciesChecker();
+
         pruneTables(tree, rootTaskContext, requiredColumns);
 
         scheduler.rewriteIterative(tree, rootTaskContext, new PruneEmptyWindowRule());
@@ -625,7 +628,6 @@ private OptExpression logicalRuleRewrite(
 
         scheduler.rewriteDownTop(tree, rootTaskContext, OnPredicateMoveAroundRule.INSTANCE);
         scheduler.rewriteIterative(tree, rootTaskContext, RuleSet.PUSH_DOWN_PREDICATE_RULES);
-
         scheduler.rewriteIterative(tree, rootTaskContext, new PartitionColumnMinMaxRewriteRule());
         scheduler.rewriteOnce(tree, rootTaskContext, RuleSet.PARTITION_PRUNE_RULES);
         scheduler.rewriteIterative(tree, rootTaskContext, new RewriteMultiDistinctRule());
@@ -656,10 +658,8 @@ private OptExpression logicalRuleRewrite(
         scheduler.rewriteOnce(tree, rootTaskContext, new PartitionColumnValueOnlyOnScanRule());
         // before MergeProjectWithChildRule, after INLINE_CTE and MergeApplyWithTableFunction
         scheduler.rewriteIterative(tree, rootTaskContext, RewriteUnnestBitmapRule.getInstance());
-
         // After this rule, we shouldn't generate logical project operator
         scheduler.rewriteIterative(tree, rootTaskContext, new MergeProjectWithChildRule());
-
         scheduler.rewriteOnce(tree, rootTaskContext, new EliminateSortColumnWithEqualityPredicateRule());
         scheduler.rewriteOnce(tree, rootTaskContext, new PushDownTopNBelowOuterJoinRule());
         // intersect rewrite depend on statistics
@@ -675,7 +675,6 @@ private OptExpression logicalRuleRewrite(
 
         // this rewrite rule should be after mv.
         scheduler.rewriteOnce(tree, rootTaskContext, RewriteSimpleAggToHDFSScanRule.SCAN_NO_PROJECT);
-
         // NOTE: This rule should be after MV Rewrite because MV Rewrite cannot handle
         // select count(distinct c) from t group by a, b
         // if this rule has applied before MV.
@@ -905,7 +904,6 @@ void memoOptimize(ConnectContext connectContext, Memo memo, TaskContext rootTask
     private OptExpression physicalRuleRewrite(ConnectContext connectContext, TaskContext rootTaskContext,
                                               OptExpression result) {
         Preconditions.checkState(result.getOp().isPhysical());
-
         int planCount = result.getPlanCount();
 
         // Since there may be many different plans in the logic phase, it's possible
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/RowOutputInfo.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/RowOutputInfo.java
@@ -137,7 +137,7 @@ public List<ColumnOutputInfo> getCommonColInfo() {
     }
 
     public List<ColumnRefOperator> getOutputColRefs() {
-        return chooseOutputMap().values().stream().map(e -> e.getColumnRef()).collect(Collectors.toList());
+        return chooseOutputMap().values().stream().map(ColumnOutputInfo::getColumnRef).collect(Collectors.toList());
     }
 
     public Map<ColumnRefOperator, ScalarOperator> getColumnRefMap() {
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/base/ColumnRefSet.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/base/ColumnRefSet.java
@@ -47,7 +47,7 @@ public ColumnRefSet(Collection<ColumnRefOperator> refs) {
 
     public static ColumnRefSet createByIds(Collection<Integer> colIds) {
         ColumnRefSet columnRefSet = new ColumnRefSet();
-        colIds.stream().forEach(columnRefSet::union);
+        colIds.forEach(columnRefSet::union);
         return columnRefSet;
     }
 
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/PushDownAggregateGroupingSetsRule.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/PushDownAggregateGroupingSetsRule.java
@@ -247,7 +247,7 @@ private OptExpression buildSelectConsume(ColumnRefFactory factory,
     /*
      * select a, b, c, d, null, sum(x) x from t group by rollup(a, b, c, d)
      */
-    private OptExpression buildSubRepeatConsume(ColumnRefFactory factory,
+    public OptExpression buildSubRepeatConsume(ColumnRefFactory factory,
                                                 Map<ColumnRefOperator, ColumnRefOperator> outputs,
                                                 LogicalAggregationOperator aggregate, LogicalRepeatOperator repeat,
                                                 int cteId) {
@@ -357,4 +357,4 @@ private OptExpression buildSubRepeatConsume(ColumnRefFactory factory,
         return OptExpression.create(projectOperator,
                 OptExpression.create(newAggregate, OptExpression.create(newRepeat, OptExpression.create(consume))));
     }
-}
+}
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/task/ApplyRuleTask.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/task/ApplyRuleTask.java
@@ -89,6 +89,7 @@ public void execute() {
             }
             List<OptExpression> targetExpressions;
             OptimizerTraceUtil.logApplyRuleBefore(context.getOptimizerContext(), rule, extractExpr);
+            
             try (Timer ignore = Tracers.watchScope(Tracers.Module.OPTIMIZER, rule.getClass().getSimpleName())) {
                 targetExpressions = rule.transform(extractExpr, context.getOptimizerContext());
             } catch (StarRocksPlannerException e) {
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/task/RewriteTreeTask.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/task/RewriteTreeTask.java
@@ -23,6 +23,7 @@
 import com.starrocks.sql.optimizer.operator.OperatorType;
 import com.starrocks.sql.optimizer.operator.pattern.Pattern;
 import com.starrocks.sql.optimizer.rule.Rule;
+import com.starrocks.sql.optimizer.validate.PlanValidator;
 
 import java.util.List;
 
@@ -62,8 +63,22 @@ public void execute() {
                 .getOptimizerOptions().isRuleDisable(rule.type()))) {
             return;
         }
+
+        // Save whole tree state before rule group execution for validation
+        OptExpression wholePlanBefore = null;
+        if (context.getOptimizerContext().getSessionVariable().enableOptimizerRuleDebug()) {
+            wholePlanBefore = planTree.getInputs().get(0);
+        }
+
         // first node must be RewriteAnchorNode
         rewrite(planTree, 0, planTree.getInputs().get(0));
+
+        // Validate after the entire rule group execution
+        if (wholePlanBefore != null && change > 0) {
+            OptExpression wholePlanAfter = planTree.getInputs().get(0);
+            PlanValidator.validateAfterRule(wholePlanBefore, wholePlanAfter, rules.get(0), context);
+        }
+        
         // pushdownNotNullPredicates should task-bind, reset it before another RewriteTreeTask
         // TODO: refactor TaskContext to make it local to support this requirement better?
         context.getOptimizerContext().clearNotNullPredicates();
@@ -97,6 +112,7 @@ protected OptExpression applyRules(OptExpression parent, int childIndex, OptExpr
             }
 
             OptimizerTraceUtil.logApplyRuleBefore(context.getOptimizerContext(), rule, root);
+            
             List<OptExpression> result;
             try (Timer ignore = Tracers.watchScope(Tracers.Module.OPTIMIZER, rule.toString())) {
                 result = rule.transform(root, context.getOptimizerContext());
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/task/TaskContext.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/task/TaskContext.java
@@ -18,13 +18,15 @@
 import com.starrocks.sql.optimizer.OptimizerContext;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.base.PhysicalPropertySet;
+import com.starrocks.sql.optimizer.validate.PlanValidator;
 
 // The context for optimizer task
 public class TaskContext {
     private final OptimizerContext optimizerContext;
     private final PhysicalPropertySet requiredProperty;
     private ColumnRefSet requiredColumns;
     private double upperBoundCost;
+    private final PlanValidator planValidator;
 
     public TaskContext(OptimizerContext context,
                        PhysicalPropertySet physicalPropertySet,
@@ -34,6 +36,11 @@ public TaskContext(OptimizerContext context,
         this.requiredProperty = physicalPropertySet;
         this.requiredColumns = requiredColumns;
         this.upperBoundCost = cost;
+        this.planValidator = new PlanValidator();
+    }
+
+    public PlanValidator getPlanValidator() {
+        return planValidator;
     }
 
     public OptimizerContext getOptimizerContext() {
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/validate/CTEUniqueChecker.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/validate/CTEUniqueChecker.java
@@ -43,7 +43,6 @@ public static CTEUniqueChecker getInstance() {
         return INSTANCE;
     }
 
-
     @Override
     public void validate(OptExpression physicalPlan, TaskContext taskContext) {
         CTEUniqueChecker.Visitor visitor = new CTEUniqueChecker.Visitor();
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/validate/InputDependenciesChecker.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/validate/InputDependenciesChecker.java
@@ -202,15 +202,15 @@ private void checkInputCols(ColumnRefSet inputCols, ColumnRefSet usedCols, OptEx
             ColumnRefSet missedCols = usedCols.clone();
             missedCols.except(inputCols);
             if (!missedCols.isEmpty()) {
-                String message = String.format("Invalid plan:%s%s%s The required cols %s cannot obtain from input cols %s.",
+                String message = String.format("Invalid plan:%s%s\n%s \nThe required cols %s cannot obtain from input cols %s.",
                         System.lineSeparator(), optExpr.debugString(), PREFIX, missedCols, inputCols);
                 throw new StarRocksPlannerException(message, ErrorType.INTERNAL_ERROR);
             }
         }
 
         private void checkInputType(ColumnRefOperator inputCol, ColumnRefOperator outputCol, OptExpression optExpression) {
             if (!outputCol.getType().isFullyCompatible(inputCol.getType())) {
-                String message = String.format("Invalid plan:%s%s%s Type of output col %s is not fully compatible with " +
+                String message = String.format("Invalid plan:%s%s\n%s \nType of output col %s is not fully compatible with " +
                                 "type of input col %s.",
                         System.lineSeparator(), optExpression.debugString(), PREFIX, outputCol, inputCol);
                 throw new StarRocksPlannerException(message, ErrorType.INTERNAL_ERROR);
@@ -219,7 +219,7 @@ private void checkInputType(ColumnRefOperator inputCol, ColumnRefOperator output
 
         private void checkChildNumberOfSet(int inputSize, int requiredSize, OptExpression optExpression) {
             if (inputSize != requiredSize) {
-                String message = String.format("Invalid plan:%s%s%s. The required number of children is %d but found %d.",
+                String message = String.format("Invalid plan:%s%s\n%s. \nThe required number of children is %d but found %d.",
                         System.lineSeparator(), optExpression.debugString(), PREFIX, requiredSize, inputSize);
                 throw new StarRocksPlannerException(message, ErrorType.INTERNAL_ERROR);
             }
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/validate/PlanValidator.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/validate/PlanValidator.java
@@ -21,6 +21,7 @@
 import com.starrocks.sql.common.ErrorType;
 import com.starrocks.sql.common.StarRocksPlannerException;
 import com.starrocks.sql.optimizer.OptExpression;
+import com.starrocks.sql.optimizer.rule.Rule;
 import com.starrocks.sql.optimizer.task.TaskContext;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -31,22 +32,35 @@ public final class PlanValidator {
 
     private static final Logger LOGGER = LogManager.getLogger(PlanValidator.class);
 
-    private static final PlanValidator INSTANCE = new PlanValidator();
+    private List<Checker> checkerList;
+    private boolean enableInputDependenciesChecker = false;
 
-    private final List<Checker> checkerList;
+    public PlanValidator() {
+        checkerList = ImmutableList.of(CTEUniqueChecker.getInstance());
+    }
+
+    /**
+     * Enable inputDependencies checker for rule transformation phase.
+     * This method should be called after column pruning to ensure InputDependenciesChecker
+     * can properly validate column dependencies in the optimized plan.
+     */
+    public void enableInputDependenciesChecker() {
+        if (!enableInputDependenciesChecker) {
+            checkerList = ImmutableList.of(
+                    CTEUniqueChecker.getInstance(),
+                    InputDependenciesChecker.getInstance());
+            enableInputDependenciesChecker = true;
+        }
+    }
 
-    private PlanValidator() {
+    public void enableAllCheckers() {
         checkerList = ImmutableList.of(
-                InputDependenciesChecker.getInstance(),
                 TypeChecker.getInstance(),
                 CTEUniqueChecker.getInstance(),
+                InputDependenciesChecker.getInstance(),
                 ColumnReuseChecker.getInstance());
     }
 
-    public static PlanValidator getInstance() {
-        return INSTANCE;
-    }
-
     public void validatePlan(OptExpression optExpression, TaskContext taskContext) {
         boolean enablePlanValidation = ConnectContext.get().getSessionVariable().getEnablePlanValidation();
         try {
@@ -58,7 +72,7 @@ public void validatePlan(OptExpression optExpression, TaskContext taskContext) {
         } catch (IllegalArgumentException e) {
             String message = e.getMessage();
             if (!message.contains("Invalid plan")) {
-                message = "Invalid plan:\n" + optExpression.debugString() + message;
+                message = "Invalid plan:\n" + optExpression.debugString() + "\n" + message;
             }
             LOGGER.debug("Failed to validate plan.", e);
             if (enablePlanValidation) {
@@ -77,6 +91,44 @@ public void validatePlan(OptExpression optExpression, TaskContext taskContext) {
         }
     }
 
+    /**
+     * Validates the query plan after applying an optimizer rule for debugging purposes.
+     * Provides detailed error information including rule name and before/after expressions.
+     * 
+     * @param beforeExpression The query plan expression before rule transformation
+     * @param afterExpression The query plan expression after rule transformation
+     * @param rule The optimizer rule that was applied
+     * @param taskContext The task context containing session variables
+     * @throws IllegalStateException if validation fails, with detailed rule information
+     */
+    public static void validateAfterRule(OptExpression beforeExpression, OptExpression afterExpression, 
+                                       Rule rule, TaskContext taskContext) {
+        if (!taskContext.getOptimizerContext().getSessionVariable().enableOptimizerRuleDebug()) {
+            return;
+        }
+        
+        try {
+            OptExpression clonedOpt = OptExpression.create(afterExpression.getOp(), afterExpression.getInputs());
+            clonedOpt.clearStatsAndInitOutputInfo();
+            taskContext.getPlanValidator().validatePlan(clonedOpt, taskContext);
+        } catch (Exception e) {
+            String errorMsg = String.format("Optimizer rule debug: Plan validation failed after applying rule [%s].\n" +
+                            "Validation error: %s\n" +
+                            "Hint: This error was caught by enable_optimizer_rule_debug=true\n" +
+                            "\n=== Query Plan Before Rule Application ===\n%s\n" +
+                            "\n=== Query Plan After Rule Application ===\n%s\n",
+                    rule.type(),
+                    e.getMessage(),
+                    beforeExpression.debugString(),
+                    afterExpression.debugString()
+            );
+            LOGGER.error(errorMsg, e);
+            
+            throw new IllegalStateException(errorMsg, e);
+        }
+    }
+
+
     public interface Checker {
         void validate(OptExpression physicalPlan, TaskContext taskContext);
     }
diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/plan/GroupingSetsTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/plan/GroupingSetsTest.java

Original file line number	Diff line number	Diff line change
`@@ -137,7 +137,7 @@ public List<ColumnOutputInfo> getCommonColInfo() {`
`137`	`137`	`}`
`138`	`138`
`139`	`139`	`public List<ColumnRefOperator> getOutputColRefs() {`
`140`		`- return chooseOutputMap().values().stream().map(e -> e.getColumnRef()).collect(Collectors.toList());`
	`140`	`+ return chooseOutputMap().values().stream().map(ColumnOutputInfo::getColumnRef).collect(Collectors.toList());`
`141`	`141`	`}`
`142`	`142`
`143`	`143`	`public Map<ColumnRefOperator, ScalarOperator> getColumnRefMap() {`
Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ public ColumnRefSet(Collection<ColumnRefOperator> refs) {`
`47`	`47`
`48`	`48`	`public static ColumnRefSet createByIds(Collection<Integer> colIds) {`
`49`	`49`	`ColumnRefSet columnRefSet = new ColumnRefSet();`
`50`		`- colIds.stream().forEach(columnRefSet::union);`
	`50`	`+ colIds.forEach(columnRefSet::union);`
`51`	`51`	`return columnRefSet;`
`52`	`52`	`}`
`53`	`53`
Original file line number	Diff line number	Diff line change
`@@ -89,6 +89,7 @@ public void execute() {`
`89`	`89`	`}`
`90`	`90`	`List<OptExpression> targetExpressions;`
`91`	`91`	`OptimizerTraceUtil.logApplyRuleBefore(context.getOptimizerContext(), rule, extractExpr);`
	`92`	`+`
`92`	`93`	`try (Timer ignore = Tracers.watchScope(Tracers.Module.OPTIMIZER, rule.getClass().getSimpleName())) {`
`93`	`94`	`targetExpressions = rule.transform(extractExpr, context.getOptimizerContext());`
`94`	`95`	`} catch (StarRocksPlannerException e) {`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,6 @@ public static CTEUniqueChecker getInstance() {`
`43`	`43`	`return INSTANCE;`
`44`	`44`	`}`
`45`	`45`
`46`		`-`
`47`	`46`	`@Override`
`48`	`47`	`public void validate(OptExpression physicalPlan, TaskContext taskContext) {`
`49`	`48`	`CTEUniqueChecker.Visitor visitor = new CTEUniqueChecker.Visitor();`