Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,11 @@ public enum TableFrom {
private final Map<CTEId, List<Pair<Multimap<Slot, Slot>, Group>>> cteIdToConsumerGroup = new HashMap<>();
private final Map<CTEId, LogicalPlan> rewrittenCteProducer = new HashMap<>();
private final Map<CTEId, LogicalPlan> rewrittenCteConsumer = new HashMap<>();
// The exprIds which a rule renamed while it rewrote the body of a CTE producer, keyed by the previous exprId.
// The consumers of that CTE map their slots to the outputs of the producer by exprId and are rewritten before
// the producer subtree, so every rule which renames an output of a producer body has to publish it here, and
// RewriteCteChildren refreshes the consumer maps with it. ConstantPropagation is currently the only publisher.
private final Map<CTEId, Map<ExprId, ExprId>> producerExprIdRemap = new HashMap<>();
private final Set<String> viewDdlSqlSet = Sets.newHashSet();
private final SqlCacheContext sqlCacheContext;

Expand Down Expand Up @@ -772,6 +777,19 @@ public Map<CTEId, LogicalPlan> getRewrittenCteConsumer() {
return rewrittenCteConsumer;
}

/**
* Record that a rule renamed exprIds while it rewrote the body of a CTE producer. Rules may rename more than
* once, so the mappings are accumulated instead of replaced, and a reader has to follow the chain to its end.
*/
public void addProducerExprIdRemap(CTEId cteId, Map<ExprId, ExprId> renamedExprIds) {
producerExprIdRemap.computeIfAbsent(cteId, k -> new HashMap<>()).putAll(renamedExprIds);
}

/** The exprIds renamed while rewriting the body of a CTE producer, keyed by their previous exprId. */
public Map<ExprId, ExprId> getProducerExprIdRemap(CTEId cteId) {
return producerExprIdRemap.getOrDefault(cteId, Collections.emptyMap());
}

/** Clear CTE-related rewrite and memo state before rebuilding it from a new plan tree. */
public void clearCteEnvironment() {
cteIdToConsumers.clear();
Expand All @@ -782,6 +800,7 @@ public void clearCteEnvironment() {
cteIdToConsumerGroup.clear();
rewrittenCteProducer.clear();
rewrittenCteConsumer.clear();
producerExprIdRemap.clear();
}

/**
Expand All @@ -796,7 +815,8 @@ public CteEnvironmentSnapshot cacheCteEnvironment() {
new HashMap<>(consumerIdToLimitRows),
copyMapOfLists(cteIdToConsumerGroup),
new HashMap<>(rewrittenCteProducer),
new HashMap<>(rewrittenCteConsumer));
new HashMap<>(rewrittenCteConsumer),
copyMapOfMaps(producerExprIdRemap));
}

/** Restore CTE-related environment from snapshot. */
Expand Down Expand Up @@ -824,6 +844,17 @@ public void restoreCteEnvironment(CteEnvironmentSnapshot snapshot) {

rewrittenCteConsumer.clear();
rewrittenCteConsumer.putAll(snapshot.rewrittenCteConsumer);

producerExprIdRemap.clear();
producerExprIdRemap.putAll(snapshot.producerExprIdRemap);
}

private static <K, V, W> Map<K, Map<V, W>> copyMapOfMaps(Map<K, Map<V, W>> source) {
Map<K, Map<V, W>> copied = new HashMap<>();
for (Map.Entry<K, Map<V, W>> entry : source.entrySet()) {
copied.put(entry.getKey(), new HashMap<>(entry.getValue()));
}
return copied;
}

private static <K, V> Map<K, Set<V>> copyMapOfSets(Map<K, Set<V>> source) {
Expand Down Expand Up @@ -852,6 +883,7 @@ public static class CteEnvironmentSnapshot {
private final Map<CTEId, List<Pair<Multimap<Slot, Slot>, Group>>> cteIdToConsumerGroup;
private final Map<CTEId, LogicalPlan> rewrittenCteProducer;
private final Map<CTEId, LogicalPlan> rewrittenCteConsumer;
private final Map<CTEId, Map<ExprId, ExprId>> producerExprIdRemap;

/**
* cte related structures in StatementContext
Expand All @@ -864,7 +896,8 @@ public CteEnvironmentSnapshot(
Map<RelationId, Long> consumerIdToLimitRows,
Map<CTEId, List<Pair<Multimap<Slot, Slot>, Group>>> cteIdToConsumerGroup,
Map<CTEId, LogicalPlan> rewrittenCteProducer,
Map<CTEId, LogicalPlan> rewrittenCteConsumer) {
Map<CTEId, LogicalPlan> rewrittenCteConsumer,
Map<CTEId, Map<ExprId, ExprId>> producerExprIdRemap) {
this.cteIdToConsumers = cteIdToConsumers;
this.cteIdToOutputIds = cteIdToOutputIds;
this.cteIdToProducer = cteIdToProducer;
Expand All @@ -873,6 +906,7 @@ public CteEnvironmentSnapshot(
this.cteIdToConsumerGroup = cteIdToConsumerGroup;
this.rewrittenCteProducer = rewrittenCteProducer;
this.rewrittenCteConsumer = rewrittenCteConsumer;
this.producerExprIdRemap = producerExprIdRemap;
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,17 @@
import org.apache.doris.nereids.rules.expression.rules.FoldConstantRule;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.And;
import org.apache.doris.nereids.trees.expressions.CTEId;
import org.apache.doris.nereids.trees.expressions.EqualPredicate;
import org.apache.doris.nereids.trees.expressions.EqualTo;
import org.apache.doris.nereids.trees.expressions.ExprId;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.Match;
import org.apache.doris.nereids.trees.expressions.NamedExpression;
import org.apache.doris.nereids.trees.expressions.Or;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral;
import org.apache.doris.nereids.trees.expressions.literal.ComparableLiteral;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
Expand Down Expand Up @@ -71,6 +74,7 @@
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
* constant propagation, like: a = 10 and a + b > 30 => a = 10 and 10 + b > 30,
Expand All @@ -79,19 +83,51 @@
* finally use the combine uniforms to replace this plan's expression's slot with literals.
*/
public class ConstantPropagation extends DefaultPlanRewriter<CascadesContext> implements CustomRewriter {
/**
* When this rule replaces a slot with a constant, the new slot describes another value than the replaced
* slot: the nullable and the data type may differ. Reusing the exprId of the replaced slot would make two
* different definitions share one exprId, and every rule which only compares exprIds (e.g.
* EliminateUnnecessaryProject) would treat them as the same column.
*
* <p>So the constant slot gets a fresh exprId, and this map records {@code replaced exprId -> fresh exprId}
* to refresh the references in the plan above.
*/
private final Map<ExprId, ExprId> replacedExprIds = Maps.newHashMap();
private ExprIdRewriter exprIdRewriter;

@Override
public Plan rewriteRoot(Plan plan, JobContext jobContext) {
// logical apply uniform maybe not correct.
if (plan.containsType(LogicalApply.class)) {
return plan;
}
return plan.accept(this, jobContext.getCascadesContext());
replacedExprIds.clear();
CascadesContext context = jobContext.getCascadesContext();
// a present current tree means this plan is the body of a CTE producer, see RewriteCteChildren
Optional<CTEId> cteId = context.getCurrentTree();
exprIdRewriter = new ExprIdRewriter(new ExprIdRewriter.ReplaceRule(replacedExprIds, true), jobContext);
Plan rewrittenPlan = plan.accept(this, context);
if (cteId.isPresent() && !replacedExprIds.isEmpty()) {
// The consumers of this CTE map their slots to the outputs of the producer by exprId, and they are
// rewritten before the producer subtree, so publish the renamed exprIds for RewriteCteChildren to
// refresh their maps. A rule may rename more than once, so the mappings are accumulated.
context.getStatementContext().addProducerExprIdRemap(cteId.get(), replacedExprIds);
}
return rewrittenPlan;
}

@Override
public Plan visit(Plan plan, CascadesContext context) {
plan = visitChildren(this, plan, context);
// the children may have replaced some slots by constants with fresh exprIds, refresh the references of
// this node before rewriting it
return exprIdRewriter.rewriteExpr(plan, replacedExprIds);
}

@Override
public Plan visitLogicalFilter(LogicalFilter<? extends Plan> filter, CascadesContext context) {
filter = visitChildren(this, filter, context);
filter = (LogicalFilter<? extends Plan>) exprIdRewriter.rewriteExpr(filter, replacedExprIds);
Expression oldPredicate = filter.getPredicate();
Expression newPredicate = replaceConstantsAndRewriteExpr(filter, oldPredicate, context);
if (isExprEqualIgnoreOrder(oldPredicate, newPredicate)) {
Expand All @@ -105,6 +141,7 @@ public Plan visitLogicalFilter(LogicalFilter<? extends Plan> filter, CascadesCon
@Override
public Plan visitLogicalHaving(LogicalHaving<? extends Plan> having, CascadesContext context) {
having = visitChildren(this, having, context);
having = (LogicalHaving<? extends Plan>) exprIdRewriter.rewriteExpr(having, replacedExprIds);
Expression oldPredicate = having.getPredicate();
Expression newPredicate = replaceConstantsAndRewriteExpr(having, oldPredicate, context);
if (isExprEqualIgnoreOrder(oldPredicate, newPredicate)) {
Expand All @@ -118,14 +155,17 @@ public Plan visitLogicalHaving(LogicalHaving<? extends Plan> having, CascadesCon
@Override
public Plan visitLogicalProject(LogicalProject<? extends Plan> project, CascadesContext context) {
project = visitChildren(this, project, context);
project = (LogicalProject<? extends Plan>) exprIdRewriter.rewriteExpr(project, replacedExprIds);
ExpressionRewriteContext rewriteContext = new ExpressionRewriteContext(project, context);
Pair<ImmutableEqualSet<Slot>, Map<Slot, Literal>> childEqualTrait =
getChildEqualSetAndConstants(project, rewriteContext);
ImmutableList.Builder<NamedExpression> newProjectsBuilder
= ImmutableList.builderWithExpectedSize(project.getProjects().size());
Set<ExprId> childOutputExprIds = project.child().getOutput().stream()
.map(Slot::getExprId).collect(Collectors.toSet());
for (NamedExpression expr : project.getProjects()) {
newProjectsBuilder.add(replaceNameExpressionConstants(
expr, rewriteContext, childEqualTrait.first, childEqualTrait.second));
expr, rewriteContext, childEqualTrait.first, childEqualTrait.second, childOutputExprIds));
}

List<NamedExpression> newProjects = newProjectsBuilder.build();
Expand All @@ -135,6 +175,7 @@ public Plan visitLogicalProject(LogicalProject<? extends Plan> project, Cascades
@Override
public Plan visitLogicalSort(LogicalSort<? extends Plan> sort, CascadesContext context) {
sort = visitChildren(this, sort, context);
sort = (LogicalSort<? extends Plan>) exprIdRewriter.rewriteExpr(sort, replacedExprIds);
ExpressionRewriteContext rewriteContext = new ExpressionRewriteContext(sort, context);
Pair<ImmutableEqualSet<Slot>, Map<Slot, Literal>> childEqualTrait
= getChildEqualSetAndConstants(sort, rewriteContext);
Expand Down Expand Up @@ -163,6 +204,7 @@ public Plan visitLogicalSort(LogicalSort<? extends Plan> sort, CascadesContext c
@Override
public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> aggregate, CascadesContext context) {
aggregate = visitChildren(this, aggregate, context);
aggregate = (LogicalAggregate<? extends Plan>) exprIdRewriter.rewriteExpr(aggregate, replacedExprIds);
ExpressionRewriteContext rewriteContext = new ExpressionRewriteContext(aggregate, context);
Pair<ImmutableEqualSet<Slot>, Map<Slot, Literal>> childEqualTrait =
getChildEqualSetAndConstants(aggregate, rewriteContext);
Expand Down Expand Up @@ -192,7 +234,6 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> aggregate, Ca
List<NamedExpression> newOutputExprs = Lists.newArrayListWithExpectedSize(oldOutputExprs.size());
ImmutableList.Builder<NamedExpression> projectBuilder
= ImmutableList.builderWithExpectedSize(oldOutputExprs.size());

boolean containsConstantOutput = false;

// after normal agg, group by expressions and output expressions are slots,
Expand All @@ -208,7 +249,13 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> aggregate, Ca
newOutputExprs.add((NamedExpression) newOutputExpr);
}

if (replacedExpr.isConstant()) {
if (replacedExpr.isConstant() && !replacedExpr.equals(expr)) {
// The constant is materialized above the aggregate, but it keeps the exprId of the aggregate
// output instead of getting a fresh one and refreshing the plan above: when the aggregate still
// outputs that column, the value is the same constant, and when it removed the column, this
// project item is the only definition of the exprId. Renaming the exprId would break the
// consumers which rely on the identity of the column, e.g. the materialized view matching, and
// would make the plan above be rewritten for nothing.
projectBuilder.add(new Alias(expr.getExprId(), replacedExpr, expr.getName()));
containsConstantOutput = true;
} else {
Expand All @@ -218,6 +265,7 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> aggregate, Ca
}

if (newGroupByExprs.equals(oldGroupByExprs) && newOutputExprs.equals(oldOutputExprs)) {
// the aggregate is not changed, the plan above keeps referencing the output slots of the aggregate
return aggregate;
}

Expand All @@ -238,6 +286,7 @@ public Plan visitLogicalJoin(LogicalJoin<? extends Plan, ? extends Plan> join, C
// update: BE not support nested loop mark join, so mark join condition need to be an equation,
// but constant propagation may rewrite an equal to TRUE/FALSE/NULL, so we don't rewrite mark join condition.
join = visitChildren(this, join, context);
join = (LogicalJoin<? extends Plan, ? extends Plan>) exprIdRewriter.rewriteExpr(join, replacedExprIds);
if (join.getJoinType().isAsofJoin()) {
return join;
}
Expand Down Expand Up @@ -293,15 +342,18 @@ public Plan visitLogicalSink(LogicalSink<? extends Plan> sink, CascadesContext c
// // the select will add a parent plan: result sink. the result sink contains a output slot reference, and its
// // data type is varchar(30), but if replace the slot reference with a varchar literal '1',
// // then the data type info varchar(30) will lost, because varchar literal '1' data type is always varchar(1),
// // so t's column will get a error type. so we don't rewrite logical sink then.
// // so t's column will get a error type. so we don't replace the slots of the sink by literals.
// Pair<ImmutableEqualSet<Slot>, Map<Slot, Literal>> childEqualTrait
// = getChildEqualSetAndConstants(sink, context);
// List<NamedExpression> newOutputExprs = sink.getOutputExprs().stream()
// .map(expr ->
// replaceNameExpressionConstants(expr, context, childEqualTrait.first, childEqualTrait.second))
// .collect(ImmutableList.toImmutableList());
// return newOutputExprs.equals(sink.getOutputExprs()) ? sink : sink.withOutputExprs(newOutputExprs);
return sink;
// The sink still has to follow the fresh exprIds of the constants materialized below it, otherwise it
// would reference a slot which its child does not output any more. Only the exprId is remapped here, so
// the name and the data type declared by the sink are kept.
return (LogicalSink<? extends Plan>) exprIdRewriter.rewriteExpr(sink, replacedExprIds);
}

/**
Expand Down Expand Up @@ -348,8 +400,7 @@ public Expression replaceConstantsAndRewriteExpr(LogicalPlan plan, Expression ex

// process NameExpression
private NamedExpression replaceNameExpressionConstants(NamedExpression expr, ExpressionRewriteContext context,
ImmutableEqualSet<Slot> equalSet, Map<Slot, Literal> constants) {

ImmutableEqualSet<Slot> equalSet, Map<Slot, Literal> constants, Set<ExprId> childOutputExprIds) {
// if a project item is a slot reference, and the slot equals to a constant value, don't rewrite it.
// because rule `EliminateUnnecessaryProject ` can eliminate a project when the project's output slots equal to
// its child's output slots.
Expand All @@ -369,12 +420,21 @@ private NamedExpression replaceNameExpressionConstants(NamedExpression expr, Exp
return expr;
}

// the item is an Alias which is not a plain slot and not a cast of one: replacing the constants of an
// Alias always yields an Alias again, so only its expression may change
Expression newExpr = replaceConstants(expr, false, context, equalSet, constants);
if (newExpr instanceof NamedExpression) {
if (!(newExpr instanceof Alias) || newExpr.equals(expr)
|| !childOutputExprIds.contains(expr.getExprId())) {
// unchanged, or the child does not produce the column any more, then this project item is the only
// definition of the exprId and it is kept
return (NamedExpression) newExpr;
} else {
return new Alias(expr.getExprId(), newExpr, expr.getName());
}
// the child still produces this column, so the project item would shadow it: materialize the constant with
// a fresh exprId and refresh the references above, see replacedExprIds
Alias replacedAlias = (Alias) newExpr;
ExprId freshExprId = StatementScopeIdGenerator.newExprId();
replacedExprIds.put(expr.getExprId(), freshExprId);
return new Alias(freshExprId, replacedAlias.child(0), replacedAlias.getName());
}

private Expression replaceConstants(Expression expression, boolean useInnerInfer, ExpressionRewriteContext context,
Expand Down
Loading
Loading