[Refactor] Refactor Optimizer&OptimizerContext interface (#55117)

Signed-off-by: Seaven <[email protected]>

fe/fe-core/src/main/java/com/starrocks/authorization/ColumnPrivilege.java

@@ -37,7 +37,8 @@
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.OptExpressionVisitor;
 import com.starrocks.sql.optimizer.Optimizer;
-import com.starrocks.sql.optimizer.OptimizerConfig;
+import com.starrocks.sql.optimizer.OptimizerFactory;
+import com.starrocks.sql.optimizer.OptimizerOptions;
 import com.starrocks.sql.optimizer.base.ColumnRefFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.base.PhysicalPropertySet;
@@ -110,13 +111,14 @@ public static void check(ConnectContext context, QueryStatement stmt, List<Table
             TransformerContext transformerContext = new TransformerContext(columnRefFactory, context, mvTransformerContext);
             logicalPlan = new RelationTransformer(transformerContext).transformWithSelectLimit(stmt.getQueryRelation());
 
-            OptimizerConfig optimizerConfig = new OptimizerConfig(OptimizerConfig.OptimizerAlgorithm.RULE_BASED);
-            optimizerConfig.disableRule(RuleType.GP_SINGLE_TABLE_MV_REWRITE);
-            optimizerConfig.disableRule(RuleType.GP_MULTI_TABLE_MV_REWRITE);
-            optimizerConfig.disableRule(RuleType.GP_PRUNE_EMPTY_OPERATOR);
-            Optimizer optimizer = new Optimizer(optimizerConfig);
-            optimizedPlan = optimizer.optimize(context, logicalPlan.getRoot(),
-                    new PhysicalPropertySet(), new ColumnRefSet(logicalPlan.getOutputColumn()), columnRefFactory);
+            OptimizerOptions optimizerOptions = new OptimizerOptions(OptimizerOptions.OptimizerStrategy.RULE_BASED);
+            optimizerOptions.disableRule(RuleType.GP_SINGLE_TABLE_MV_REWRITE);
+            optimizerOptions.disableRule(RuleType.GP_MULTI_TABLE_MV_REWRITE);
+            optimizerOptions.disableRule(RuleType.GP_PRUNE_EMPTY_OPERATOR);
+            Optimizer optimizer =
+                    OptimizerFactory.create(OptimizerFactory.initContext(context, columnRefFactory, optimizerOptions));
+            optimizedPlan = optimizer.optimize(logicalPlan.getRoot(),
+                    new PhysicalPropertySet(), new ColumnRefSet(logicalPlan.getOutputColumn()));
 
             optimizedPlan.getOp().accept(new ScanColumnCollector(tableObjectToTableName, scanColumns), optimizedPlan, null);
         }

fe/fe-core/src/main/java/com/starrocks/sql/DeletePlanner.java

@@ -37,6 +37,7 @@
 import com.starrocks.sql.ast.QueryRelation;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.Optimizer;
+import com.starrocks.sql.optimizer.OptimizerFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.base.PhysicalPropertySet;
@@ -72,13 +73,11 @@ public ExecPlan plan(DeleteStmt deleteStatement, ConnectContext session) {
             // Non-query must use the strategy assign scan ranges per driver sequence, which local shuffle agg cannot use.
             session.getSessionVariable().setEnableLocalShuffleAgg(false);
 
-            Optimizer optimizer = new Optimizer();
+            Optimizer optimizer = OptimizerFactory.create(OptimizerFactory.initContext(session, columnRefFactory));
             OptExpression optimizedPlan = optimizer.optimize(
-                    session,
                     logicalPlan.getRoot(),
                     new PhysicalPropertySet(),
-                    new ColumnRefSet(logicalPlan.getOutputColumn()),
-                    columnRefFactory);
+                    new ColumnRefSet(logicalPlan.getOutputColumn()));
             ExecPlan execPlan = PlanFragmentBuilder.createPhysicalPlan(optimizedPlan, session,
                     logicalPlan.getOutputColumn(), columnRefFactory,
                     colNames, TResultSinkType.MYSQL_PROTOCAL, false);

fe/fe-core/src/main/java/com/starrocks/sql/InsertPlanner.java

@@ -78,6 +78,7 @@
 import com.starrocks.sql.common.TypeManager;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.Optimizer;
+import com.starrocks.sql.optimizer.OptimizerFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.base.DistributionProperty;
@@ -511,18 +512,16 @@ private ExecPlan buildExecPlanWithRetry(InsertStmt insertStmt, ConnectContext se
     private ExecPlan buildExecPlan(InsertStmt insertStmt, ConnectContext session, List<ColumnRefOperator> outputColumns,
                                    LogicalPlan logicalPlan, ColumnRefFactory columnRefFactory,
                                    QueryRelation queryRelation, Table targetTable) {
-        Optimizer optimizer = new Optimizer();
         PhysicalPropertySet requiredPropertySet = createPhysicalPropertySet(insertStmt, outputColumns,
                 session.getSessionVariable());
         OptExpression optimizedPlan;
 
         try (Timer ignore2 = Tracers.watchScope("Optimizer")) {
+            Optimizer optimizer = OptimizerFactory.create(OptimizerFactory.initContext(session, columnRefFactory));
             optimizedPlan = optimizer.optimize(
-                    session,
                     logicalPlan.getRoot(),
                     requiredPropertySet,
-                    new ColumnRefSet(logicalPlan.getOutputColumn()),
-                    columnRefFactory);
+                    new ColumnRefSet(logicalPlan.getOutputColumn()));
         }
 
         //8. Build fragment exec plan

fe/fe-core/src/main/java/com/starrocks/sql/ShortCircuitPlanner.java

@@ -24,7 +24,6 @@
 import com.starrocks.sql.optimizer.OptExpressionVisitor;
 import com.starrocks.sql.optimizer.operator.ColumnFilterConverter;
 import com.starrocks.sql.optimizer.operator.Operator;
-import com.starrocks.sql.optimizer.operator.OperatorType;
 import com.starrocks.sql.optimizer.operator.logical.LogicalFilterOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalLimitOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalOlapScanOperator;
@@ -73,19 +72,12 @@ public static BaseLogicalPlanChecker createLogicalPlanChecker(OptExpression root
         }
     }
 
-    public static OptExpression checkSupportShortCircuitRead(OptExpression root, ConnectContext connectContext) {
+    public static boolean checkSupportShortCircuitRead(OptExpression root, ConnectContext connectContext) {
         if (!connectContext.getSessionVariable().isEnableShortCircuit()) {
-            root.setShortCircuit(false);
-            return root;
-        }
-        boolean supportShortCircuit = root.getOp().accept(new LogicalPlanChecker(), root, null);
-        if (supportShortCircuit && OperatorType.LOGICAL_LIMIT.equals(root.getOp().getOpType())) {
-            root = root.getInputs().get(0);
+            return false;
         }
-        root.setShortCircuit(supportShortCircuit);
-        return root;
+        return root.getOp().accept(new LogicalPlanChecker(), root, null);
     }
-
     protected static boolean isRedundant(Map<ColumnRefOperator, ScalarOperator> projections) {
         for (Map.Entry<ColumnRefOperator, ScalarOperator> entry : projections.entrySet()) {
             if (!entry.getKey().equals(entry.getValue())) {

fe/fe-core/src/main/java/com/starrocks/sql/StatementPlanner.java

@@ -60,6 +60,9 @@
 import com.starrocks.sql.common.UnsupportedException;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.Optimizer;
+import com.starrocks.sql.optimizer.OptimizerContext;
+import com.starrocks.sql.optimizer.OptimizerFactory;
+import com.starrocks.sql.optimizer.OptimizerOptions;
 import com.starrocks.sql.optimizer.OptimizerTraceUtil;
 import com.starrocks.sql.optimizer.base.ColumnRefFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
@@ -234,7 +237,7 @@ private static ExecPlan createQueryPlan(StatementBase stmt,
                                             ConnectContext session,
                                             TResultSinkType resultSinkType) {
         QueryStatement queryStmt = (QueryStatement) stmt;
-        QueryRelation query = (QueryRelation) queryStmt.getQueryRelation();
+        QueryRelation query = queryStmt.getQueryRelation();
         List<String> colNames = query.getColumnOutputNames();
         // 1. Build Logical plan
         ColumnRefFactory columnRefFactory = new ColumnRefFactory();
@@ -247,20 +250,21 @@ private static ExecPlan createQueryPlan(StatementBase stmt,
             logicalPlan = new RelationTransformer(transformerContext).transformWithSelectLimit(query);
         }
 
-        OptExpression root = ShortCircuitPlanner.checkSupportShortCircuitRead(logicalPlan.getRoot(), session);
-
+        boolean isShortCircuit = ShortCircuitPlanner.checkSupportShortCircuitRead(logicalPlan.getRoot(), session);
         OptExpression optimizedPlan;
         try (Timer ignored = Tracers.watchScope("Optimizer")) {
             // 2. Optimize logical plan and build physical plan
-            Optimizer optimizer = new Optimizer();
-            optimizedPlan = optimizer.optimize(
-                    session,
-                    root,
-                    mvTransformerContext,
-                    stmt,
+            OptimizerContext optimizerContext = OptimizerFactory.initContext(session, columnRefFactory);
+            optimizerContext.setMvTransformerContext(mvTransformerContext);
+            optimizerContext.setStatement(stmt);
+            if (isShortCircuit) {
+                optimizerContext.setOptimizerOptions(OptimizerOptions.newShortCircuitOpt());
+            }
+
+            Optimizer optimizer = OptimizerFactory.create(optimizerContext);
+            optimizedPlan = optimizer.optimize(logicalPlan.getRoot(),
                     new PhysicalPropertySet(),
-                    new ColumnRefSet(logicalPlan.getOutputColumn()),
-                    columnRefFactory);
+                    new ColumnRefSet(logicalPlan.getOutputColumn()));
         }
 
         try (Timer ignored = Tracers.watchScope("ExecPlanBuild")) {
@@ -273,7 +277,7 @@ private static ExecPlan createQueryPlan(StatementBase stmt,
             ExecPlan execPlan = PlanFragmentBuilder.createPhysicalPlan(
                     optimizedPlan, session, logicalPlan.getOutputColumn(), columnRefFactory, colNames,
                     resultSinkType,
-                    !session.getSessionVariable().isSingleNodeExecPlan());
+                    !session.getSessionVariable().isSingleNodeExecPlan(), isShortCircuit);
             execPlan.setLogicalPlan(logicalPlan);
             execPlan.setColumnRefFactory(columnRefFactory);
             return execPlan;
@@ -312,25 +316,26 @@ public static ExecPlan createQueryPlanWithReTry(QueryStatement queryStmt,
                 logicalPlan = new RelationTransformer(transformerContext).transformWithSelectLimit(query);
             }
 
-            OptExpression root = ShortCircuitPlanner.checkSupportShortCircuitRead(logicalPlan.getRoot(), session);
-
+            boolean isShortCircuit = ShortCircuitPlanner.checkSupportShortCircuitRead(logicalPlan.getRoot(), session);
             OptExpression optimizedPlan;
             try (Timer ignored = Tracers.watchScope("Optimizer")) {
+                OptimizerContext optimizerContext = OptimizerFactory.initContext(session, columnRefFactory);
                 // 2. Optimize logical plan and build physical plan
-                Optimizer optimizer = new Optimizer();
                 // FIXME: refactor this into Optimizer.optimize() method.
                 // set query tables into OptimizeContext so can be added for mv rewrite
                 if (Config.skip_whole_phase_lock_mv_limit >= 0) {
-                    optimizer.setQueryTables(olapTables);
+                    optimizerContext.setQueryTables(olapTables);
+                }
+
+                if (isShortCircuit) {
+                    optimizerContext.setOptimizerOptions(OptimizerOptions.newShortCircuitOpt());
                 }
-                optimizedPlan = optimizer.optimize(
-                        session,
-                        root,
-                        mvTransformerContext,
-                        queryStmt,
-                        new PhysicalPropertySet(),
-                        new ColumnRefSet(logicalPlan.getOutputColumn()),
-                        columnRefFactory);
+                optimizerContext.setMvTransformerContext(mvTransformerContext);
+                optimizerContext.setStatement(queryStmt);
+
+                Optimizer optimizer = OptimizerFactory.create(optimizerContext);
+                optimizedPlan = optimizer.optimize(logicalPlan.getRoot(), new PhysicalPropertySet(),
+                        new ColumnRefSet(logicalPlan.getOutputColumn()));
             }
 
             try (Timer ignored = Tracers.watchScope("ExecPlanBuild")) {
@@ -341,7 +346,7 @@ public static ExecPlan createQueryPlanWithReTry(QueryStatement queryStmt,
                 ExecPlan plan = PlanFragmentBuilder.createPhysicalPlan(
                         optimizedPlan, session, logicalPlan.getOutputColumn(), columnRefFactory, colNames,
                         resultSinkType,
-                        !session.getSessionVariable().isSingleNodeExecPlan());
+                        !session.getSessionVariable().isSingleNodeExecPlan(), isShortCircuit);
                 final long finalPlanStartTime = planStartTime;
                 isSchemaValid = olapTables.stream().allMatch(t -> OptimisticVersion.validateTableUpdate(t,
                         finalPlanStartTime));

fe/fe-core/src/main/java/com/starrocks/sql/UpdatePlanner.java

@@ -40,6 +40,8 @@
 import com.starrocks.sql.ast.UpdateStmt;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.Optimizer;
+import com.starrocks.sql.optimizer.OptimizerContext;
+import com.starrocks.sql.optimizer.OptimizerFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.base.PhysicalPropertySet;
@@ -95,15 +97,14 @@ public ExecPlan plan(UpdateStmt updateStmt, ConnectContext session) {
             session.getSessionVariable().setEnableLocalShuffleAgg(false);
 
             long tableId = targetTable.getId();
-            Optimizer optimizer = new Optimizer();
-            optimizer.setUpdateTableId(tableId);
+            OptimizerContext optimizerContext = OptimizerFactory.initContext(session, columnRefFactory);
+            optimizerContext.setUpdateTableId(tableId);
 
+            Optimizer optimizer = OptimizerFactory.create(optimizerContext);
             OptExpression optimizedPlan = optimizer.optimize(
-                    session,
                     optExprBuilder.getRoot(),
                     new PhysicalPropertySet(),
-                    new ColumnRefSet(outputColumns),
-                    columnRefFactory);
+                    new ColumnRefSet(outputColumns));
             ExecPlan execPlan = PlanFragmentBuilder.createPhysicalPlan(optimizedPlan, session,
                     outputColumns, columnRefFactory, colNames, TResultSinkType.MYSQL_PROTOCAL, false);
             DescriptorTable descriptorTable = execPlan.getDescTbl();

fe/fe-core/src/main/java/com/starrocks/sql/analyzer/AnalyzeStmtAnalyzer.java

@@ -42,9 +42,7 @@
 import com.starrocks.sql.ast.DropStatsStmt;
 import com.starrocks.sql.ast.StatementBase;
 import com.starrocks.sql.common.MetaUtils;
-import com.starrocks.sql.optimizer.Memo;
-import com.starrocks.sql.optimizer.OptimizerConfig;
-import com.starrocks.sql.optimizer.OptimizerContext;
+import com.starrocks.sql.optimizer.OptimizerFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefFactory;
 import com.starrocks.sql.optimizer.statistics.Statistics;
 import com.starrocks.statistic.StatisticUtils;
@@ -381,8 +379,7 @@ private void analyzeAnalyzeTypeDesc(ConnectContext session, StatementBase statem
                     }
 
                     Statistics tableStats = session.getGlobalStateMgr().getMetadataMgr().
-                            getTableStatistics(new OptimizerContext(new Memo(), new ColumnRefFactory(), session,
-                                            OptimizerConfig.defaultConfig()),
+                            getTableStatistics(OptimizerFactory.initContext(session, new ColumnRefFactory()),
                                     tableName.getCatalog(), analyzeTable, Maps.newHashMap(), keys, null);
                     totalRows = tableStats.getOutputRowCount();
                 }

fe/fe-core/src/main/java/com/starrocks/sql/analyzer/MaterializedViewAnalyzer.java

@@ -96,6 +96,8 @@
 import com.starrocks.sql.common.PListCell;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.Optimizer;
+import com.starrocks.sql.optimizer.OptimizerContext;
+import com.starrocks.sql.optimizer.OptimizerFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefFactory;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.base.PhysicalPropertySet;
@@ -457,14 +459,13 @@ private void planMVQuery(CreateMaterializedViewStatement createStmt, QueryStatem
                 // Build logical plan for view query
                 OptExprBuilder optExprBuilder = logicalPlan.getRootBuilder();
                 logicalPlan = new LogicalPlan(optExprBuilder, outputColumns, logicalPlan.getCorrelation());
-                Optimizer optimizer = new Optimizer();
+                OptimizerContext optimizerContext = OptimizerFactory.initContext(ctx, columnRefFactory);
+                Optimizer optimizer = OptimizerFactory.create(optimizerContext);
                 PhysicalPropertySet requiredPropertySet = PhysicalPropertySet.EMPTY;
                 OptExpression optimizedPlan = optimizer.optimize(
-                        ctx,
                         logicalPlan.getRoot(),
                         requiredPropertySet,
-                        new ColumnRefSet(logicalPlan.getOutputColumn()),
-                        columnRefFactory);
+                        new ColumnRefSet(logicalPlan.getOutputColumn()));
                 optimizedPlan.deriveMVProperty();
 
                 // TODO: refine rules for mv plan