Skip to content

Commit

Permalink
xxx
Browse files Browse the repository at this point in the history
Signed-off-by: Seaven <[email protected]>
  • Loading branch information
Seaven committed Jan 15, 2025
1 parent beeb28e commit e618b99
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 66 deletions.
27 changes: 10 additions & 17 deletions fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,6 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
Expand All @@ -143,18 +142,16 @@ public class Optimizer {
private final OptimizerContext context;
private OptimizerConfig optimizerConfig;
private MvRewriteStrategy mvRewriteStrategy = new MvRewriteStrategy();
private TaskScheduler scheduler = new TaskScheduler();
private final TaskScheduler scheduler = new TaskScheduler();
private Memo memo;

// collect all LogicalOlapScanOperators in the query before any optimization
private final List<LogicalOlapScanOperator> allLogicalOlapScanOperators = Lists.newArrayList();

Optimizer(OptimizerContext context) {
this.context = context;
}

@VisibleForTesting
public OptimizerConfig getOptimizerConfig() {
return optimizerConfig;
}

public OptimizerContext getContext() {
return context;
}
Expand All @@ -167,14 +164,15 @@ public MvRewriteStrategy getMvRewriteStrategy() {
private void prepare(OptExpression logicOperatorTree) {
optimizerConfig = context.getOptimizerConfig();

if (!context.getOptimizerConfig().isRuleBased()) {
if (!optimizerConfig.isRuleBased()) {
memo = new Memo();
context.setMemo(memo);
}
context.setTaskScheduler(scheduler);

// collect all olap scan operator
collectAllLogicalOlapScanOperators(logicOperatorTree, context);
Utils.extractOperator(logicOperatorTree, allLogicalOlapScanOperators,
op -> op instanceof LogicalOlapScanOperator);
}

public OptExpression optimize(OptExpression logicOperatorTree, ColumnRefSet requiredColumns) {
Expand Down Expand Up @@ -307,9 +305,10 @@ private OptExpression optimizeByCost(ConnectContext connectContext,
// valid the final plan
PlanValidator.getInstance().validatePlan(finalPlan, rootTaskContext);
// validate mv and log tracer if needed
MVRewriteValidator.getInstance().validateMV(connectContext, finalPlan, rootTaskContext);
MVRewriteValidator mvRewriteValidator = new MVRewriteValidator(allLogicalOlapScanOperators);
mvRewriteValidator.validateMV(connectContext, finalPlan, rootTaskContext);
// audit mv
MVRewriteValidator.getInstance().auditMv(connectContext, finalPlan, rootTaskContext);
mvRewriteValidator.auditMv(connectContext, finalPlan, rootTaskContext);
return finalPlan;
}
}
Expand Down Expand Up @@ -995,12 +994,6 @@ private OptExpression extractBestPlan(PhysicalPropertySet requiredProperty,
return expression;
}

private void collectAllLogicalOlapScanOperators(OptExpression tree, OptimizerContext optimizerContext) {
List<LogicalOlapScanOperator> list = Lists.newArrayList();
Utils.extractOperator(tree, list, op -> op instanceof LogicalOlapScanOperator);
optimizerContext.setAllLogicalOlapScanOperators(Collections.unmodifiableList(list));
}

private List<PhysicalOlapScanOperator> collectAllPhysicalOlapScanOperators(OptExpression tree) {
List<PhysicalOlapScanOperator> list = Lists.newArrayList();
Utils.extractOperator(tree, list, op -> op instanceof PhysicalOlapScanOperator);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
import com.starrocks.sql.common.StarRocksPlannerException;
import com.starrocks.sql.optimizer.base.ColumnRefFactory;
import com.starrocks.sql.optimizer.dump.DumpInfo;
import com.starrocks.sql.optimizer.operator.logical.LogicalOlapScanOperator;
import com.starrocks.sql.optimizer.operator.scalar.IsNullPredicateOperator;
import com.starrocks.sql.optimizer.rewrite.JoinPredicatePushdown;
import com.starrocks.sql.optimizer.rule.RuleSet;
Expand All @@ -49,14 +48,15 @@ public class OptimizerContext {
private Set<OlapTable> queryTables;
private long updateTableId = -1;

private OptimizerConfig optimizerConfig;
private VectorSearchOptions vectorSearchOptions;

// ============================ Optimizer ============================
private Memo memo;
private final RuleSet ruleSet = new RuleSet();
private final RuleSet ruleSet;
private TaskScheduler taskScheduler;
private OptimizerConfig optimizerConfig;
private VectorSearchOptions vectorSearchOptions = new VectorSearchOptions();

private final CTEContext cteContext = new CTEContext();
private final CTEContext cteContext;
private TaskContext currentTaskContext;
private final QueryMaterializationContext queryMaterializationContext = new QueryMaterializationContext();

Expand All @@ -82,15 +82,17 @@ public class OptimizerContext {
// which should be kept to be used to convert outer join into inner join.
private final List<IsNullPredicateOperator> pushdownNotNullPredicates = Lists.newArrayList();

// collect all LogicalOlapScanOperators in the query before any optimization
private List<LogicalOlapScanOperator> allLogicalOlapScanOperators;

OptimizerContext(ConnectContext context) {
this.connectContext = context;
this.ruleSet = new RuleSet();
this.cteContext = new CTEContext();
cteContext.reset();
this.cteContext.setEnableCTE(getSessionVariable().isCboCteReuse());
this.cteContext.setInlineCTERatio(getSessionVariable().getCboCTERuseRatio());
this.cteContext.setMaxCTELimit(getSessionVariable().getCboCTEMaxLimit());

this.vectorSearchOptions = new VectorSearchOptions();
this.optimizerConfig = OptimizerConfig.defaultConfig();
}

// @VisibleForTesting
Expand Down Expand Up @@ -178,6 +180,22 @@ public long getUpdateTableId() {
return updateTableId;
}

public OptimizerConfig getOptimizerConfig() {
return optimizerConfig;
}

public void setOptimizerConfig(OptimizerConfig optimizerConfig) {
this.optimizerConfig = optimizerConfig;
}

public void setVectorSearchOptions(VectorSearchOptions vectorSearchOptions) {
this.vectorSearchOptions = vectorSearchOptions;
}

public VectorSearchOptions getVectorSearchOptions() {
return vectorSearchOptions;
}

// ============================ Optimizer ============================
public Memo getMemo() {
return memo;
Expand All @@ -199,22 +217,6 @@ public void setTaskScheduler(TaskScheduler taskScheduler) {
this.taskScheduler = taskScheduler;
}

public OptimizerConfig getOptimizerConfig() {
return optimizerConfig;
}

public void setOptimizerConfig(OptimizerConfig optimizerConfig) {
this.optimizerConfig = optimizerConfig;
}

public void setVectorSearchOptions(VectorSearchOptions vectorSearchOptions) {
this.vectorSearchOptions = vectorSearchOptions;
}

public VectorSearchOptions getVectorSearchOptions() {
return vectorSearchOptions;
}

public CTEContext getCteContext() {
return cteContext;
}
Expand Down Expand Up @@ -268,14 +270,6 @@ public boolean isInMemoPhase() {
return this.inMemoPhase;
}

public void setAllLogicalOlapScanOperators(List<LogicalOlapScanOperator> allScanOperators) {
this.allLogicalOlapScanOperators = allScanOperators;
}

public List<LogicalOlapScanOperator> getAllLogicalOlapScanOperators() {
return allLogicalOlapScanOperators;
}

/**
* Get all valid candidate materialized views for the query:
* - The materialized view is valid to rewrite by rule(SPJG)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import com.starrocks.sql.optimizer.MaterializationContext;
import com.starrocks.sql.optimizer.OptExpression;
import com.starrocks.sql.optimizer.OptimizerContext;
import com.starrocks.sql.optimizer.operator.logical.LogicalOlapScanOperator;
import com.starrocks.sql.optimizer.rule.transformation.materialization.MaterializedViewRewriter;
import com.starrocks.sql.optimizer.task.TaskContext;
import org.apache.commons.collections.CollectionUtils;
Expand All @@ -36,10 +37,15 @@
import static com.starrocks.sql.optimizer.rule.transformation.materialization.MvUtils.collectMaterializedViews;

public class MVRewriteValidator {
private static final MVRewriteValidator INSTANCE = new MVRewriteValidator();

public static MVRewriteValidator getInstance() {
return INSTANCE;
// private static final MVRewriteValidator INSTANCE = new MVRewriteValidator();
//
// public static MVRewriteValidator getInstance() {
// return INSTANCE;
// }
private List<LogicalOlapScanOperator> allLogicalOlapScanOperators;

public MVRewriteValidator(List<LogicalOlapScanOperator> allLogicalOlapScanOperators) {
this.allLogicalOlapScanOperators = allLogicalOlapScanOperators;
}

private static boolean isUpdateMaterializedViewMetrics(ConnectContext connectContext) {
Expand Down Expand Up @@ -79,7 +85,7 @@ public void auditMv(ConnectContext connectContext, OptExpression physicalPlan, T
List<MaterializedView> mvs = collectMaterializedViews(physicalPlan);
// To avoid queries that query the materialized view directly, only consider materialized views
// that are not used in rewriting before.
Set<Long> beforeTableIds = optimizerContext.getAllLogicalOlapScanOperators().stream()
Set<Long> beforeTableIds = allLogicalOlapScanOperators.stream()
.map(op -> op.getTable().getId())
.collect(Collectors.toSet());
if (CollectionUtils.isNotEmpty(mvs)) {
Expand Down Expand Up @@ -130,7 +136,7 @@ public void validateMV(ConnectContext connectContext, OptExpression physicalPlan
}

List<MaterializedView> mvs = collectMaterializedViews(physicalPlan);
Set<Long> beforeTableIds = taskContext.getOptimizerContext().getAllLogicalOlapScanOperators().stream()
Set<Long> beforeTableIds = allLogicalOlapScanOperators.stream()
.map(op -> op.getTable().getId())
.collect(Collectors.toSet());

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -104,26 +104,29 @@ public void testOptimizer() throws Exception {
LogicalPlan logicalPlan = new RelationTransformer(columnRefFactory, connectContext)
.transformWithSelectLimit(query.getQueryRelation());

Optimizer optimizer = OptimizerFactory.create(OptimizerFactory.mockContext(connectContext, columnRefFactory));
Assert.assertFalse(optimizer.getOptimizerConfig().isRuleBased());
Assert.assertFalse(optimizer.getOptimizerConfig().isRuleDisable(RuleType.TF_MERGE_TWO_PROJECT));
Assert.assertFalse(optimizer.getOptimizerConfig().isRuleDisable(RuleType.GP_AGGREGATE_REWRITE));
OptimizerContext optimizerContext = OptimizerFactory.mockContext(connectContext, columnRefFactory);
Optimizer optimizer = OptimizerFactory.create(optimizerContext);
Assert.assertFalse(optimizerContext.getOptimizerConfig().isRuleBased());
Assert.assertFalse(optimizerContext.getOptimizerConfig().isRuleDisable(RuleType.TF_MERGE_TWO_PROJECT));
Assert.assertFalse(optimizerContext.getOptimizerConfig().isRuleDisable(RuleType.GP_AGGREGATE_REWRITE));

OptExpression expr = optimizer.optimize(logicalPlan.getRoot(), new PhysicalPropertySet(),
new ColumnRefSet(logicalPlan.getOutputColumn()));
Assert.assertTrue(expr.getInputs().get(0).getOp() instanceof PhysicalOlapScanOperator);
Assert.assertNotNull(expr.getInputs().get(0).getOp().getPredicate());


OptimizerConfig optimizerConfig = new OptimizerConfig(OptimizerConfig.OptimizerAlgorithm.RULE_BASED);
optimizerConfig.disableRule(RuleType.TF_MERGE_TWO_PROJECT);
optimizerConfig.disableRule(RuleType.GP_PUSH_DOWN_PREDICATE);
Optimizer optimizer1 = OptimizerFactory.create(OptimizerFactory.mockContext(connectContext, columnRefFactory,
optimizerConfig));
Assert.assertTrue(optimizer1.getOptimizerConfig().isRuleBased());
Assert.assertFalse(optimizer1.getOptimizerConfig().isRuleDisable(RuleType.TF_MERGE_TWO_AGG_RULE));
Assert.assertTrue(optimizer1.getOptimizerConfig().isRuleDisable(RuleType.TF_MERGE_TWO_PROJECT));
Assert.assertFalse(optimizer1.getOptimizerConfig().isRuleDisable(RuleType.GP_COLLECT_CTE));
Assert.assertTrue(optimizer1.getOptimizerConfig().isRuleDisable(RuleType.GP_PUSH_DOWN_PREDICATE));
OptimizerContext optimizerContext1 = OptimizerFactory.mockContext(connectContext, columnRefFactory,
optimizerConfig);
Optimizer optimizer1 = OptimizerFactory.create(optimizerContext1);
Assert.assertTrue(optimizerContext1.getOptimizerConfig().isRuleBased());
Assert.assertFalse(optimizerContext1.getOptimizerConfig().isRuleDisable(RuleType.TF_MERGE_TWO_AGG_RULE));
Assert.assertTrue(optimizerContext1.getOptimizerConfig().isRuleDisable(RuleType.TF_MERGE_TWO_PROJECT));
Assert.assertFalse(optimizerContext1.getOptimizerConfig().isRuleDisable(RuleType.GP_COLLECT_CTE));
Assert.assertTrue(optimizerContext1.getOptimizerConfig().isRuleDisable(RuleType.GP_PUSH_DOWN_PREDICATE));

OptExpression expr1 = optimizer1.optimize(logicalPlan.getRoot(), new PhysicalPropertySet(),
new ColumnRefSet(logicalPlan.getOutputColumn()));
Expand Down

0 comments on commit e618b99

Please sign in to comment.