diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java index 7e25d4d1157209..1e76e5dde72e48 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java @@ -183,6 +183,11 @@ public enum TableFrom { private final Map, Group>>> cteIdToConsumerGroup = new HashMap<>(); private final Map rewrittenCteProducer = new HashMap<>(); private final Map 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> producerExprIdRemap = new HashMap<>(); private final Set viewDdlSqlSet = Sets.newHashSet(); private final SqlCacheContext sqlCacheContext; @@ -772,6 +777,19 @@ public Map 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 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 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(); @@ -782,6 +800,7 @@ public void clearCteEnvironment() { cteIdToConsumerGroup.clear(); rewrittenCteProducer.clear(); rewrittenCteConsumer.clear(); + producerExprIdRemap.clear(); } /** @@ -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. */ @@ -824,6 +844,17 @@ public void restoreCteEnvironment(CteEnvironmentSnapshot snapshot) { rewrittenCteConsumer.clear(); rewrittenCteConsumer.putAll(snapshot.rewrittenCteConsumer); + + producerExprIdRemap.clear(); + producerExprIdRemap.putAll(snapshot.producerExprIdRemap); + } + + private static Map> copyMapOfMaps(Map> source) { + Map> copied = new HashMap<>(); + for (Map.Entry> entry : source.entrySet()) { + copied.put(entry.getKey(), new HashMap<>(entry.getValue())); + } + return copied; } private static Map> copyMapOfSets(Map> source) { @@ -852,6 +883,7 @@ public static class CteEnvironmentSnapshot { private final Map, Group>>> cteIdToConsumerGroup; private final Map rewrittenCteProducer; private final Map rewrittenCteConsumer; + private final Map> producerExprIdRemap; /** * cte related structures in StatementContext @@ -864,7 +896,8 @@ public CteEnvironmentSnapshot( Map consumerIdToLimitRows, Map, Group>>> cteIdToConsumerGroup, Map rewrittenCteProducer, - Map rewrittenCteConsumer) { + Map rewrittenCteConsumer, + Map> producerExprIdRemap) { this.cteIdToConsumers = cteIdToConsumers; this.cteIdToOutputIds = cteIdToOutputIds; this.cteIdToProducer = cteIdToProducer; @@ -873,6 +906,7 @@ public CteEnvironmentSnapshot( this.cteIdToConsumerGroup = cteIdToConsumerGroup; this.rewrittenCteProducer = rewrittenCteProducer; this.rewrittenCteConsumer = rewrittenCteConsumer; + this.producerExprIdRemap = producerExprIdRemap; } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagation.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagation.java index 16c2fb5ef244f7..ae32d9094e3a40 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagation.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagation.java @@ -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; @@ -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, @@ -79,6 +83,17 @@ * finally use the combine uniforms to replace this plan's expression's slot with literals. */ public class ConstantPropagation extends DefaultPlanRewriter 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. + * + *

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 replacedExprIds = Maps.newHashMap(); + private ExprIdRewriter exprIdRewriter; @Override public Plan rewriteRoot(Plan plan, JobContext jobContext) { @@ -86,12 +101,33 @@ public Plan rewriteRoot(Plan plan, JobContext jobContext) { 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 = 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 filter, CascadesContext context) { filter = visitChildren(this, filter, context); + filter = (LogicalFilter) exprIdRewriter.rewriteExpr(filter, replacedExprIds); Expression oldPredicate = filter.getPredicate(); Expression newPredicate = replaceConstantsAndRewriteExpr(filter, oldPredicate, context); if (isExprEqualIgnoreOrder(oldPredicate, newPredicate)) { @@ -105,6 +141,7 @@ public Plan visitLogicalFilter(LogicalFilter filter, CascadesCon @Override public Plan visitLogicalHaving(LogicalHaving having, CascadesContext context) { having = visitChildren(this, having, context); + having = (LogicalHaving) exprIdRewriter.rewriteExpr(having, replacedExprIds); Expression oldPredicate = having.getPredicate(); Expression newPredicate = replaceConstantsAndRewriteExpr(having, oldPredicate, context); if (isExprEqualIgnoreOrder(oldPredicate, newPredicate)) { @@ -118,14 +155,17 @@ public Plan visitLogicalHaving(LogicalHaving having, CascadesCon @Override public Plan visitLogicalProject(LogicalProject project, CascadesContext context) { project = visitChildren(this, project, context); + project = (LogicalProject) exprIdRewriter.rewriteExpr(project, replacedExprIds); ExpressionRewriteContext rewriteContext = new ExpressionRewriteContext(project, context); Pair, Map> childEqualTrait = getChildEqualSetAndConstants(project, rewriteContext); ImmutableList.Builder newProjectsBuilder = ImmutableList.builderWithExpectedSize(project.getProjects().size()); + Set 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 newProjects = newProjectsBuilder.build(); @@ -135,6 +175,7 @@ public Plan visitLogicalProject(LogicalProject project, Cascades @Override public Plan visitLogicalSort(LogicalSort sort, CascadesContext context) { sort = visitChildren(this, sort, context); + sort = (LogicalSort) exprIdRewriter.rewriteExpr(sort, replacedExprIds); ExpressionRewriteContext rewriteContext = new ExpressionRewriteContext(sort, context); Pair, Map> childEqualTrait = getChildEqualSetAndConstants(sort, rewriteContext); @@ -163,6 +204,7 @@ public Plan visitLogicalSort(LogicalSort sort, CascadesContext c @Override public Plan visitLogicalAggregate(LogicalAggregate aggregate, CascadesContext context) { aggregate = visitChildren(this, aggregate, context); + aggregate = (LogicalAggregate) exprIdRewriter.rewriteExpr(aggregate, replacedExprIds); ExpressionRewriteContext rewriteContext = new ExpressionRewriteContext(aggregate, context); Pair, Map> childEqualTrait = getChildEqualSetAndConstants(aggregate, rewriteContext); @@ -192,7 +234,6 @@ public Plan visitLogicalAggregate(LogicalAggregate aggregate, Ca List newOutputExprs = Lists.newArrayListWithExpectedSize(oldOutputExprs.size()); ImmutableList.Builder projectBuilder = ImmutableList.builderWithExpectedSize(oldOutputExprs.size()); - boolean containsConstantOutput = false; // after normal agg, group by expressions and output expressions are slots, @@ -208,7 +249,13 @@ public Plan visitLogicalAggregate(LogicalAggregate 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 { @@ -218,6 +265,7 @@ public Plan visitLogicalAggregate(LogicalAggregate 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; } @@ -238,6 +286,7 @@ public Plan visitLogicalJoin(LogicalJoin 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) exprIdRewriter.rewriteExpr(join, replacedExprIds); if (join.getJoinType().isAsofJoin()) { return join; } @@ -293,7 +342,7 @@ public Plan visitLogicalSink(LogicalSink 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, Map> childEqualTrait // = getChildEqualSetAndConstants(sink, context); // List newOutputExprs = sink.getOutputExprs().stream() @@ -301,7 +350,10 @@ public Plan visitLogicalSink(LogicalSink sink, CascadesContext c // 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) exprIdRewriter.rewriteExpr(sink, replacedExprIds); } /** @@ -348,8 +400,7 @@ public Expression replaceConstantsAndRewriteExpr(LogicalPlan plan, Expression ex // process NameExpression private NamedExpression replaceNameExpressionConstants(NamedExpression expr, ExpressionRewriteContext context, - ImmutableEqualSet equalSet, Map constants) { - + ImmutableEqualSet equalSet, Map constants, Set 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. @@ -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, diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriter.java index ddeb1d7d6eb4b6..f1480ae245a517 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriter.java @@ -20,18 +20,18 @@ import org.apache.doris.nereids.jobs.JobContext; import org.apache.doris.nereids.pattern.Pattern; import org.apache.doris.nereids.rules.Rule; -import org.apache.doris.nereids.rules.expression.ExpressionPatternMatcher; -import org.apache.doris.nereids.rules.expression.ExpressionPatternRuleFactory; import org.apache.doris.nereids.rules.expression.ExpressionRewrite; +import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext; +import org.apache.doris.nereids.rules.expression.ExpressionRewriteRule; import org.apache.doris.nereids.rules.expression.ExpressionRuleExecutor; -import org.apache.doris.nereids.rules.expression.ExpressionRuleType; import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; -import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter; import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalSink; import com.google.common.collect.ImmutableList; @@ -44,11 +44,43 @@ public class ExprIdRewriter extends ExpressionRewrite { private final JobContext jobContext; public ExprIdRewriter(ReplaceRule replaceRule, JobContext jobContext) { - super(new ExpressionRuleExecutor(ImmutableList.of(bottomUp(replaceRule)))); + super(new ExpressionRuleExecutor(ImmutableList.of(new ReplaceExprIds(replaceRule)))); rules = buildRules(); this.jobContext = jobContext; } + /** Replace the exprIds of one expression structurally, see ReplaceRule#replaceExprIds. */ + private static class ReplaceExprIds implements ExpressionRewriteRule { + private final ReplaceRule replaceRule; + + ReplaceExprIds(ReplaceRule replaceRule) { + this.replaceRule = replaceRule; + } + + @Override + public Expression rewrite(Expression expression, ExpressionRewriteContext context) { + return replaceRule.replaceExprIds(expression); + } + } + + /** + * Follow the replacement chain of an exprId to its end: an exprId may be renamed more than once, e.g. by + * several rules or by several passes of the same rule, and only the last exprId exists in the plan then. + */ + public static ExprId lastExprId(Map replaceMap, ExprId exprId) { + ExprId lastExprId = replaceMap.get(exprId); + if (lastExprId == null) { + return exprId; + } + while (true) { + ExprId nextExprId = replaceMap.get(lastExprId); + if (nextExprId == null) { + return lastExprId; + } + lastExprId = nextExprId; + } + } + /**rewriteExpr*/ public Plan rewriteExpr(Plan plan, Map replaceMap) { if (replaceMap.isEmpty()) { @@ -64,6 +96,25 @@ public Plan rewriteExpr(Plan plan, Map replaceMap) { } } } + // ExpressionRewrite has a rule for the common sink types only, e.g. there is none for the blackhole sink. A + // sink takes its output expressions from its child, so it must follow the fresh exprIds of the child as + // well, otherwise it would keep expecting a slot which the rewritten child does not output any more. So + // rewrite the output expressions of any sink type here. + if (plan instanceof LogicalSink) { + LogicalSink sink = (LogicalSink) plan; + ExpressionRewriteContext context = new ExpressionRewriteContext(sink, jobContext.getCascadesContext()); + ImmutableList.Builder newOutputExprs + = ImmutableList.builderWithExpectedSize(sink.getOutputExprs().size()); + boolean changed = false; + for (NamedExpression outputExpr : sink.getOutputExprs()) { + NamedExpression newOutputExpr = (NamedExpression) rewriter.rewrite(outputExpr, context); + changed |= !newOutputExpr.equals(outputExpr); + newOutputExprs.add(newOutputExpr); + } + if (changed) { + return sink.withOutputExprs(newOutputExprs.build()); + } + } return plan; } @@ -76,63 +127,37 @@ public Plan rewriteExpr(Plan plan, Map replaceMap) { * e.g. replaceMap:{0:3, 1:6, 6:7} * SlotReference:a#0 -> a#3, a#1 -> a#7 * */ - public static class ReplaceRule implements ExpressionPatternRuleFactory { + public static class ReplaceRule { private static final DefaultExpressionRewriter> SLOT_REPLACER = new DefaultExpressionRewriter>() { @Override public Expression visitSlotReference(SlotReference slot, Map replaceMap) { - ExprId newId = replaceMap.get(slot.getExprId()); - if (newId == null) { + if (replaceMap.get(slot.getExprId()) == null) { return slot; } - ExprId lastId = newId; - while (true) { - newId = replaceMap.get(lastId); - if (newId == null) { - return slot.withExprId(lastId); - } else { - lastId = newId; - } - } + return slot.withExprId(lastExprId(replaceMap, slot.getExprId())); } }; private static final DefaultExpressionRewriter> SLOT_ALIAS_REPLACER = new DefaultExpressionRewriter>() { @Override public Expression visitSlotReference(SlotReference slot, Map replaceMap) { - ExprId newId = replaceMap.get(slot.getExprId()); - if (newId == null) { + if (replaceMap.get(slot.getExprId()) == null) { return slot; } - ExprId lastId = newId; - while (true) { - newId = replaceMap.get(lastId); - if (newId == null) { - return slot.withExprId(lastId); - } else { - lastId = newId; - } - } + return slot.withExprId(lastExprId(replaceMap, slot.getExprId())); } @Override public Expression visitAlias(Alias alias, Map replaceMap) { - ExprId newId = replaceMap.get(alias.getExprId()); - if (newId == null) { + if (replaceMap.get(alias.getExprId()) == null) { return alias; } - ExprId lastId = newId; - while (true) { - newId = replaceMap.get(lastId); - if (newId == null) { - return alias.withExprId(lastId); - } else { - lastId = newId; - } - } + return alias.withExprId(lastExprId(replaceMap, alias.getExprId())); } }; private final Map replaceMap; + private final DefaultExpressionRewriter> structuralReplacer; private DefaultExpressionRewriter> replacer; /** @@ -145,24 +170,36 @@ public ReplaceRule(Map replaceMap, boolean rewriteAlias) { } else { replacer = SLOT_REPLACER; } + // children first: the id replacement of a node must not stop the traversal, e.g. an Alias keeps its own + // exprId while a slot of the expression below it is renamed. The replacer is applied to the node itself + // only, its children are already rewritten here, so the whole expression is visited once. + this.structuralReplacer = new DefaultExpressionRewriter>() { + @Override + public Expression visit(Expression expr, Map replaceMap) { + Expression withRewrittenChildren = rewriteChildren(this, expr, replaceMap); + if (withRewrittenChildren instanceof SlotReference + || (rewriteAlias && withRewrittenChildren instanceof Alias)) { + return withRewrittenChildren.accept(replacer, replaceMap); + } + return withRewrittenChildren; + } + }; } public ReplaceRule(Map replaceMap) { this(replaceMap, false); } - @Override - public List> buildRules() { - return ImmutableList.of( - matchesType(SlotReference.class).thenApply(ctx -> { - Slot slot = ctx.expr; - return slot.accept(replacer, replaceMap); - }).toRule(ExpressionRuleType.EXPR_ID_REWRITE_REPLACE), - matchesType(Alias.class).thenApply(ctx -> { - Alias alias = ctx.expr; - return alias.accept(replacer, replaceMap); - }).toRule(ExpressionRuleType.EXPR_ID_REWRITE_REPLACE) - ); + /** + * Replace the exprIds of the slots (and of the aliases, when this rule was created for them) in one + * expression. Unlike the bottom up expression rewriter this enters every branch, including the branches + * which require short circuit evaluation, e.g. the assignments of MERGE INTO: replacing an exprId only + * rebuilds the expression and never evaluates it, and a reference which keeps an old exprId would point to + * a slot which the rewritten child does not output any more. + */ + public Expression replaceExprIds(Expression expression) { + return expression.accept(structuralReplacer, replaceMap); } + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java index c957a9e853f68d..fe68decaa5dd02 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java @@ -24,6 +24,7 @@ import org.apache.doris.nereids.jobs.rewrite.RewriteJob; import org.apache.doris.nereids.properties.PhysicalProperties; import org.apache.doris.nereids.trees.expressions.CTEId; +import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.Slot; @@ -44,10 +45,15 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.LinkedHashMultimap; +import com.google.common.collect.Multimap; import com.google.common.collect.Sets; +import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Map.Entry; import java.util.Optional; import java.util.Set; @@ -99,22 +105,93 @@ public Plan visitLogicalCTEAnchor(LogicalCTEAnchor cteConsumers = collectCteConsumers(outer, cteAnchor.getCteId()); + cascadesContext.getCteIdToConsumers().put(cteAnchor.getCteId(), cteConsumers); + if (cteConsumers.isEmpty()) { + return outer; + } + Plan producer = cteAnchor.child(0).accept(this, cascadesContext); + // ConstantPropagation is currently the only rule which renames the exprIds of a producer output, and it + // publishes what it renamed (see StatementContext#addProducerExprIdRemap); a rule which renames an output + // of a producer body has to do the same, otherwise the consumers, which are rewritten before the producer, + // keep mapping to an exprId which the producer does not output any more. A rule may rename more than once, + // so the chain is followed to its end. + Map producerExprIdRemap = cascadesContext.getStatementContext() + .getProducerExprIdRemap(cteAnchor.getCteId()); + if (!producerExprIdRemap.isEmpty()) { + Map producerOutputSlots = new HashMap<>(); + for (Slot producerOutputSlot : producer.getOutput()) { + // a CTE may expose the same slot more than once, e.g. `SELECT k, k FROM t`: both positions are the + // same producer slot, so the first one is enough for the lookup + producerOutputSlots.putIfAbsent(producerOutputSlot.getExprId(), producerOutputSlot); + } + outer = (LogicalPlan) outer.accept(new CteConsumerOutputsRefresher(cteAnchor.getCteId(), + producerExprIdRemap, producerOutputSlots), null); + cascadesContext.getStatementContext().getRewrittenCteConsumer().put(cteAnchor.getCteId(), outer); + cascadesContext.getCteIdToConsumers() + .put(cteAnchor.getCteId(), collectCteConsumers(outer, cteAnchor.getCteId())); + } + return cteAnchor.withChildren(producer, outer); + } + + private static Set collectCteConsumers(Plan outer, CTEId cteId) { Set cteConsumers = Sets.newHashSet(); outer.foreach(p -> { if (p instanceof LogicalCTEConsumer) { LogicalCTEConsumer logicalCTEConsumer = (LogicalCTEConsumer) p; - if (logicalCTEConsumer.getCteId().equals(cteAnchor.getCteId())) { + if (logicalCTEConsumer.getCteId().equals(cteId)) { cteConsumers.add(logicalCTEConsumer); } } return false; }); - cascadesContext.getCteIdToConsumers().put(cteAnchor.getCteId(), cteConsumers); - if (cteConsumers.isEmpty()) { - return outer; + return cteConsumers; + } + + /** + * Refresh the producer slots of the two output maps of the consumers of a CTE with the exprIds the rules of the + * producer body published, see StatementContext#getProducerExprIdRemap. + */ + static class CteConsumerOutputsRefresher extends DefaultPlanRewriter { + private final CTEId cteId; + private final Map producerExprIdRemap; + private final Map producerOutputSlots; + + CteConsumerOutputsRefresher(CTEId cteId, Map producerExprIdRemap, + Map producerOutputSlots) { + this.cteId = cteId; + this.producerExprIdRemap = producerExprIdRemap; + this.producerOutputSlots = producerOutputSlots; + } + + @Override + public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, Void context) { + if (!cteConsumer.getCteId().equals(cteId)) { + return cteConsumer; + } + Map consumerToProducerOutputMap = new LinkedHashMap<>(); + Multimap producerToConsumerOutputMap = LinkedHashMultimap.create(); + for (Map.Entry entry : cteConsumer.getConsumerToProducerOutputMap().entrySet()) { + Slot newProducerSlot = refreshProducerSlot(entry.getValue()); + consumerToProducerOutputMap.put(entry.getKey(), newProducerSlot); + for (Slot consumerSlot : cteConsumer.getProducerToConsumerOutputMap().get(entry.getValue())) { + producerToConsumerOutputMap.put(newProducerSlot, consumerSlot); + } + } + return cteConsumer.withTwoMaps(consumerToProducerOutputMap, producerToConsumerOutputMap); + } + + private Slot refreshProducerSlot(Slot producerSlot) { + ExprId lastExprId = ExprIdRewriter.lastExprId(producerExprIdRemap, producerSlot.getExprId()); + if (lastExprId.equals(producerSlot.getExprId())) { + return producerSlot; + } + // the column may be pruned by the rewrite of the producer body after it was renamed, e.g. by + // ColumnPruning when the plan above the consumer stopped reading it; such a column has no counterpart + // in the producer output, so its entry stays as it is + Slot newProducerSlot = producerOutputSlots.get(lastExprId); + return newProducerSlot != null ? newProducerSlot : producerSlot; } - Plan producer = cteAnchor.child(0).accept(this, cascadesContext); - return cteAnchor.withChildren(producer, outer); } @Override diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagationFreshExprIdTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagationFreshExprIdTest.java new file mode 100644 index 00000000000000..34985c7255e593 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagationFreshExprIdTest.java @@ -0,0 +1,280 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.rules.rewrite; + +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCTEConsumer; +import org.apache.doris.nereids.trees.plans.logical.LogicalCTEProducer; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.utframe.TestWithFeService; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * The project which ConstantPropagation materializes above an aggregate must not re-label a column which the plan + * below defines with another value. + * + *

A constant grouping key which the aggregate keeps is passed through as the slot of the aggregate, and a + * constant grouping key which the aggregate removes keeps the exprId it already had, so nothing below defines that + * exprId any more. The project item used to reuse the exprId of the removed aggregate output while other rules + * later brought a subtree which outputs that exprId back below it, and rules which compare exprIds only, e.g. + * EliminateUnnecessaryProject, then removed the constant project. RQG case 11000258898 reached such a plan and the + * last AdjustNullable pass reported "AdjustNullable convert slot 2004-04-28#xx from not-nullable to nullable". + */ +public class ConstantPropagationFreshExprIdTest extends TestWithFeService { + private static final String SQL_11000258898 = "select '2004-04-28', '2014-07-07'" + + " intersect distinct" + + " select col_datetime_3__undef_signed_not_null, col_date_undef_signed_not_null" + + " from table_100_undef_partitions2_keys3_properties4_distributed_by52" + + " where col_decimal_12_2__undef_signed = 94.0212" + + " union distinct" + + " select col_datetime_3__undef_signed, col_datetime_3__undef_signed_not_null" + + " from table_100_undef_partitions2_keys3_properties4_distributed_by54" + + " where col_datetime_3__undef_signed = '2011-11-15 13:33:05'" + + " intersect" + + " select col_date_undef_signed_not_null, col_date_undef_signed_not_null" + + " from table_100_undef_partitions2_keys3_properties4_distributed_by53" + + " intersect" + + " select col_datetime_3__undef_signed_not_null, col_date_undef_signed_not_null" + + " from table_100_undef_partitions2_keys3_properties4_distributed_by54 ;"; + + private static final String SQL_11000259585 = "select * from ((select 'vZRgkfajsD', 'cKToNmsxok' except (select col_char_25__undef_signed_not_null, col_char_40__undef_signed from table_100_undef_partitions2_keys3_properties4_distributed_by54 where col_date_undef_signed = '2016-09-20' order by 1, 2 desc limit 1, 8) union distinct select col_varchar_25__undef_signed, col_char_25__undef_signed from table_100_undef_partitions2_keys3_properties4_distributed_by53 where col_tinyint_undef_signed_not_null = 4.4929476592789594 intersect distinct select col_date_undef_signed_not_null, col_date_undef_signed_not_null from table_100_undef_partitions2_keys3_properties4_distributed_by52 where col_decimal_12_2__undef_signed = 94.0212) as t1);"; + + private static final String SQL_MATERIALIZED_CTE = "WITH c AS (SELECT k, v, count(*) AS cnt" + + " FROM constant_propagation_fresh_expr_id_t WHERE k = 5 AND v = 7 GROUP BY k, v)" + + " SELECT l.k, l.v, l.cnt, r.k, r.v, r.cnt FROM c l CROSS JOIN c r"; + + private static final String SQL_MATERIALIZED_CTE_CONSTANT_COLUMN + = "WITH c AS (SELECT v, k + 0 AS x, count(*) AS cnt" + + " FROM constant_propagation_fresh_expr_id_t WHERE k = 5 GROUP BY v, k)" + + " SELECT l.v, l.x, l.cnt, r.v, r.x, r.cnt FROM c l CROSS JOIN c r"; + + private static final String SQL_MATERIALIZED_CTE_REPEATED_SLOT + = "WITH c AS (SELECT k, v, v, count(*) AS cnt" + + " FROM constant_propagation_fresh_expr_id_t WHERE k = 5 AND v = 7 GROUP BY k, v)" + + " SELECT l.*, r.* FROM c l CROSS JOIN c r"; + + @Override + protected void runBeforeAll() throws Exception { + createDatabase("test"); + createTable("create table test.constant_propagation_fresh_expr_id_t(k int null, v int null)" + + " distributed by hash(k) properties('replication_num' = '1');"); + createTable("create table test.table_100_undef_partitions2_keys3_properties4_distributed_by52 (" + + " col_datetime_3__undef_signed_not_null datetime(3) not null," + + " col_datetime_3__undef_signed datetime(3) null," + + " col_datetime_undef_signed datetime null," + + " col_datetime_undef_signed_not_null datetime not null," + + " col_date_undef_signed_not_null date not null," + + " col_date_undef_signed date null," + + " col_decimal_12_2__undef_signed decimal(12,2) null," + + " col_decimal_32_6__undef_signed decimal(32,6) null)" + + " distributed by hash(col_datetime_3__undef_signed) properties('replication_num' = '1');"); + createTable("create table test.table_100_undef_partitions2_keys3_properties4_distributed_by53 (" + + " col_date_undef_signed_not_null date not null," + + " col_date_undef_signed date null," + + " col_datetime_3__undef_signed datetime(3) null," + + " col_datetime_3__undef_signed_not_null datetime(3) not null," + + " col_datetime_undef_signed datetime null," + + " col_datetime_undef_signed_not_null datetime not null," + + " col_varchar_25__undef_signed varchar(25) null," + + " col_char_25__undef_signed_not_null char(25) not null," + + " col_char_25__undef_signed char(25) null," + + " col_char_40__undef_signed char(40) null," + + " col_varchar_100__undef_signed varchar(100) null," + + " col_tinyint_undef_signed_not_null tinyint not null," + + " col_bigint_undef_signed_not_null bigint not null)" + + " distributed by hash(col_date_undef_signed_not_null) properties('replication_num' = '1');"); + createTable("create table test.table_100_undef_partitions2_keys3_properties4_distributed_by54 (" + + " col_datetime_3__undef_signed datetime(3) null," + + " col_datetime_3__undef_signed_not_null datetime(3) not null," + + " col_datetime_undef_signed datetime null," + + " col_datetime_undef_signed_not_null datetime not null," + + " col_date_undef_signed_not_null date not null," + + " col_date_undef_signed date null," + + " col_char_25__undef_signed_not_null char(25) not null," + + " col_char_40__undef_signed char(40) null)" + + " distributed by hash(col_datetime_3__undef_signed) properties('replication_num' = '1');"); + createTable("create table test.table_100_undef_partitions2_keys3_properties4_distributed_by5 (" + + " col_date_undef_signed_not_null date not null," + + " col_datetime_undef_signed datetime null," + + " col_datetime_undef_signed_not_null datetime not null," + + " col_datetime_3__undef_signed datetime(3) null," + + " col_datetime_3__undef_signed_not_null datetime(3) not null," + + " col_char_25__undef_signed char(25) null," + + " col_char_25__undef_signed_not_null char(25) not null," + + " col_varchar_100__undef_signed_not_null varchar(100) not null," + + " col_varchar_100__undef_signed varchar(100) null)" + + " distributed by hash(col_date_undef_signed_not_null) properties('replication_num' = '1');"); + connectContext.setDatabase("test"); + connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + } + + @Test + void testConstantAggregateProjectKeepsTheExprId() { + // run this rule alone, so the project it materializes above the aggregate is still part of the plan + PlanChecker checker = PlanChecker.from(connectContext) + .analyze("select k, v from constant_propagation_fresh_expr_id_t where k = 5 and v = 7 group by k, v") + .applyCustom(new ConstantPropagation()); + Plan plan = checker.getPlan(); + LogicalProject constantProject = findConstantProject(plan); + Assertions.assertNotNull(constantProject, plan.treeString()); + List items = constantProject.getProjects(); + // the kept grouping key is still materialized as a constant above the aggregate, with the exprId of the + // aggregate output: no fresh exprId and no refresh of the plan above + NamedExpression keptKey = items.get(0); + Assertions.assertTrue(keptKey instanceof Alias, plan.treeString()); + Assertions.assertEquals(new IntegerLiteral(5), ((Alias) keptKey).child()); + Assertions.assertTrue(constantProject.child().getOutput().stream() + .anyMatch(slot -> slot.getExprId().equals(keptKey.getExprId())), + "the kept grouping key keeps the exprId of the aggregate output:\n" + plan.treeString()); + // the removed grouping key is materialized with the exprId it already had as well + NamedExpression removedKey = items.get(1); + Assertions.assertTrue(removedKey instanceof Alias, plan.treeString()); + Assertions.assertEquals(new IntegerLiteral(7), ((Alias) removedKey).child()); + // the plan above references the outputs of the project, unchanged + Set projectOutputIds = constantProject.getOutput().stream() + .map(Slot::getExprId).collect(Collectors.toSet()); + Set rootOutputIds = plan.getOutput().stream() + .map(Slot::getExprId).collect(Collectors.toSet()); + Assertions.assertTrue(rootOutputIds.containsAll(projectOutputIds), + "the plan above does not reference the project outputs:\n" + plan.treeString()); + } + + // The two RQG reproductions of the master version (testRqgCase11000258898 and testRqgCase11000259585) are + // not ported: on branch-4.2 both queries already trip "AdjustNullable convert slot 2004-04-28#55 from + // not-nullable to nullable" (and "cKToNmsxok#45") during the rewrite on the pristine branch, i.e. before this + // change, and the failure is unchanged after it. ConstantPropagation does not modify the plans of these + // queries on this branch at all (dumped with the rule applied alone: the plan is identical), so the failure + // comes from another rule here, while on master the rule does rewrite them. The tables those queries need + // are still created below, so the SQL constants stay usable if the branch is fixed later. + + @Test + void testMaterializedCteConsumerMapsStayValid() throws Exception { + // the producer body and the consumers are rewritten in different subtree contexts and the consumers are + // rewritten first, so the fresh exprIds of the producer outputs must not cross the CTE boundary + Set analyzedProducerOutputIds = cteProducerOutputIds( + PlanChecker.from(connectContext).analyze(SQL_MATERIALIZED_CTE).getPlan()); + Plan plan = PlanChecker.from(connectContext).analyze(SQL_MATERIALIZED_CTE).rewrite().getPlan(); + assertCteConsumerMapsStayValid(plan); + // constant propagation keeps the exprIds of the producer outputs, so they are the same as in the analyzed + // plan and the consumer maps still resolve; a rule which renames them has to refresh the maps, which the + // assertCteConsumerMapsStayValid above checks + Set rewrittenProducerOutputIds = cteProducerOutputIds(plan); + Assertions.assertEquals(analyzedProducerOutputIds.size(), rewrittenProducerOutputIds.size(), + plan.treeString()); + Assertions.assertEquals(analyzedProducerOutputIds, rewrittenProducerOutputIds, plan.treeString()); + // the query has to stay plannable + PlanChecker.from(connectContext).checkPlannerResult(SQL_MATERIALIZED_CTE); + } + + @Test + void testMaterializedCteWithRepeatedOutputSlot() { + // a CTE may expose the same slot more than once (LogicalCTEConsumer is multimap based), and the lookup of + // the producer slots must tolerate that instead of failing on the duplicated exprId + Set analyzedProducerOutputIds = cteProducerOutputIds( + PlanChecker.from(connectContext).analyze(SQL_MATERIALIZED_CTE_REPEATED_SLOT).getPlan()); + Plan plan = PlanChecker.from(connectContext) + .analyze(SQL_MATERIALIZED_CTE_REPEATED_SLOT).rewrite().getPlan(); + assertCteConsumerMapsStayValid(plan); + Assertions.assertEquals(analyzedProducerOutputIds.size(), cteProducerOutputIds(plan).size(), plan.treeString()); + PlanChecker.from(connectContext).checkPlannerResult(SQL_MATERIALIZED_CTE_REPEATED_SLOT); + } + + private static Set cteProducerOutputIds(Plan plan) { + Set producerOutputExprIds = new java.util.HashSet<>(); + plan.foreach(node -> { + if (node instanceof LogicalCTEProducer) { + ((LogicalCTEProducer) node).getOutput().forEach(slot -> producerOutputExprIds.add(slot.getExprId())); + } + return false; + }); + Assertions.assertFalse(producerOutputExprIds.isEmpty(), plan.treeString()); + return producerOutputExprIds; + } + + /** Every producer slot a CTE consumer maps to has to be an output slot of the producer. */ + private static void assertCteConsumerMapsStayValid(Plan plan) { + Set producerOutputExprIds = new java.util.HashSet<>(); + Set consumers = new java.util.HashSet<>(); + plan.foreach(node -> { + if (node instanceof LogicalCTEProducer) { + ((LogicalCTEProducer) node).getOutput().forEach(slot -> producerOutputExprIds.add(slot.getExprId())); + } else if (node instanceof LogicalCTEConsumer) { + consumers.add((LogicalCTEConsumer) node); + } + return false; + }); + Assertions.assertFalse(consumers.isEmpty(), plan.treeString()); + for (LogicalCTEConsumer consumer : consumers) { + for (Slot producerSlot : consumer.getProducerToConsumerOutputMap().keySet()) { + Assertions.assertTrue(producerOutputExprIds.contains(producerSlot.getExprId()), + "the consumer map of " + consumer.getRelationId() + " references " + producerSlot + + " which the producer does not output: " + plan.treeString()); + } + for (Slot producerSlot : consumer.getConsumerToProducerOutputMap().values()) { + Assertions.assertTrue(producerOutputExprIds.contains(producerSlot.getExprId()), + "the consumer map of " + consumer.getRelationId() + " references " + producerSlot + + " which the producer does not output: " + plan.treeString()); + } + } + } + + @Test + void testMaterializedCteKeepsConstantProducerOutput() { + // `k + 0` is uniform (k = 5), so it folds to a constant, and it is an output column of the CTE producer: + // whether the rule materializes the constant or leaves the column alone, the consumers have to keep + // mapping to an output which the producer really exposes + Set analyzedProducerOutputIds = cteProducerOutputIds( + PlanChecker.from(connectContext).analyze(SQL_MATERIALIZED_CTE_CONSTANT_COLUMN).getPlan()); + Plan plan = PlanChecker.from(connectContext) + .analyze(SQL_MATERIALIZED_CTE_CONSTANT_COLUMN).rewrite().getPlan(); + assertCteConsumerMapsStayValid(plan); + // the producer keeps its three output columns, whether they are rewritten in place or renamed + Assertions.assertEquals(analyzedProducerOutputIds.size(), cteProducerOutputIds(plan).size(), plan.treeString()); + PlanChecker.from(connectContext).checkPlannerResult(SQL_MATERIALIZED_CTE_CONSTANT_COLUMN); + } + + private static LogicalProject findConstantProject(Plan plan) { + if (plan instanceof LogicalProject + && plan.child(0) instanceof LogicalAggregate + && ((LogicalProject) plan).getProjects().stream().anyMatch(item -> !(item instanceof SlotReference))) { + return (LogicalProject) plan; + } + for (Plan child : plan.children()) { + LogicalProject found = findConstantProject(child); + if (found != null) { + return found; + } + } + return null; + } + +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagationTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagationTest.java index f5acd25df12557..3d8ebcb489637d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagationTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ConstantPropagationTest.java @@ -27,12 +27,14 @@ import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.Cast; 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.GreaterThan; import org.apache.doris.nereids.trees.expressions.IsNull; import org.apache.doris.nereids.trees.expressions.MatchAny; import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.Not; +import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.functions.agg.Sum; import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; @@ -43,12 +45,14 @@ import org.apache.doris.nereids.trees.plans.JoinType; import org.apache.doris.nereids.trees.plans.RelationId; import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalBlackholeSink; import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; import org.apache.doris.nereids.trees.plans.logical.LogicalHaving; import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalSink; import org.apache.doris.nereids.trees.plans.logical.LogicalSort; import org.apache.doris.nereids.types.BigIntType; import org.apache.doris.nereids.types.IntegerType; @@ -67,6 +71,7 @@ import java.util.List; import java.util.Set; +import java.util.stream.Collectors; class ConstantPropagationTest { @@ -361,7 +366,18 @@ void testLogicalAggregate() { LogicalAggregate rewrittenAgg1 = (LogicalAggregate) rewrittenProject1.child(); Assertions.assertEquals(expectGroupby1, rewrittenAgg1.getGroupByExpressions()); Assertions.assertEquals(expectAggOutput1, rewrittenAgg1.getOutputExpressions()); - Assertions.assertEquals(expectProjection1, rewrittenProject1.getProjects()); + // the constant outputs above the aggregate are materialized with a fresh exprId, so that one exprId + // does not describe both the output slot of the aggregate and the constant + List actualProjects1 = rewrittenProject1.getProjects(); + Assertions.assertEquals(expectProjection1.subList(0, 5), actualProjects1.subList(0, 5)); + // the aggregate removes these two constant grouping keys from its outputs, so the project item is the only + // definition of their exprIds and the ids are kept + assertKeptConstantProjectItem(actualProjects1.get(5), aggOutput.get(5).getExprId(), + new IntegerLiteral(1), aggOutput.get(5).getName()); + Assertions.assertEquals(expectProjection1.get(6), actualProjects1.get(6)); + assertKeptConstantProjectItem(actualProjects1.get(7), aggOutput.get(7).getExprId(), + new IntegerLiteral(11), aggOutput.get(7).getName()); + Assertions.assertEquals(expectProjection1.subList(8, 10), actualProjects1.subList(8, 10)); Set conjunctions2 = ImmutableSet.of( new EqualTo(studentId, new IntegerLiteral(3)), @@ -396,7 +412,141 @@ void testLogicalAggregate() { LogicalAggregate rewrittenAgg2 = (LogicalAggregate) rewrittenProject2.child(); Assertions.assertEquals(expectGroupby2, rewrittenAgg2.getGroupByExpressions()); Assertions.assertEquals(expectAggOutput2, rewrittenAgg2.getOutputs()); - Assertions.assertEquals(expectProjection2, rewrittenProject2.getProjects()); + List actualProjects2 = rewrittenProject2.getProjects(); + Assertions.assertEquals(expectProjection2.subList(0, 5), actualProjects2.subList(0, 5)); + // all grouping keys are constant, so the aggregate keeps the first one; the constant is still materialized + // above it, with the exprId of the aggregate output + assertKeptConstantProjectItem(actualProjects2.get(5), expectAggOutput2.get(5).getExprId(), + new IntegerLiteral(3), null); + // the other keys are removed from the aggregate, their exprIds are kept + assertKeptConstantProjectItem(actualProjects2.get(6), aggOutput.get(6).getExprId(), + new IntegerLiteral(3), aggOutput.get(6).getName()); + assertKeptConstantProjectItem(actualProjects2.get(7), aggOutput.get(7).getExprId(), + new IntegerLiteral(13), aggOutput.get(7).getName()); + assertKeptConstantProjectItem(actualProjects2.get(8), aggOutput.get(8).getExprId(), + new IntegerLiteral(13), aggOutput.get(8).getName()); + assertKeptConstantProjectItem(actualProjects2.get(9), aggOutput.get(9).getExprId(), + new IntegerLiteral(6), aggOutput.get(9).getName()); + } + + @Test + void testConstantGroupByKeyKeepsExprIdWhenTheAggregateDropsIt() { + // positive gate for the materialized view shape which this PR must not change: the constant is materialized + // above the aggregate with the exprId of the aggregate output, no fresh exprId and no refresh of the plan + // above, so the identity of the column is kept (the MV matching relies on it) + LogicalFilter filter = new LogicalFilter<>(ImmutableSet.of( + new EqualTo(studentId, new IntegerLiteral(5))), student); + List groupby = ImmutableList.of(studentId, studentGender); + List aggOutput = ImmutableList.of( + new Alias(new Sum(studentAge), "sum_age"), + studentGender, + studentId); + LogicalAggregate> aggregate + = new LogicalAggregate<>(groupby, aggOutput, filter); + + LogicalProject rewrittenProject = (LogicalProject) executor.rewriteRoot(aggregate, jobContext); + NamedExpression constantItem = rewrittenProject.getProjects().get(2); + assertKeptConstantProjectItem(constantItem, studentId.getExprId(), new IntegerLiteral(5), studentId.getName()); + // the aggregate dropped the key, so nothing below the project defines this column any more + Assertions.assertFalse(rewrittenProject.child().getOutput().stream() + .anyMatch(slot -> slot.getExprId().equals(studentId.getExprId()))); + } + + @Test + public void testProjectItemShadowingAChildSlotGetsFreshExprId() { + // a project item which re-defines a column which its child still outputs describes another value, so it + // must be materialized with a fresh exprId, and the sink above has to follow that exprId + LogicalFilter filter = new LogicalFilter<>(ImmutableSet.of( + new EqualTo(studentId, new IntegerLiteral(5))), student); + LogicalProject> project = new LogicalProject<>(ImmutableList.of( + new Alias(studentId.getExprId(), new Add(studentId, new IntegerLiteral(0)), studentId.getName()), + studentGender), filter); + List sinkOutputExprs = ImmutableList.of(studentId, studentGender); + LogicalBlackholeSink>> sink + = new LogicalBlackholeSink<>(sinkOutputExprs, project); + + LogicalSink rewrittenSink = (LogicalSink) executor.rewriteRoot(sink, jobContext); + // the child of the project still outputs studentId, so the constant item got a fresh exprId + NamedExpression constantItem = ((LogicalProject) rewrittenSink.child()).getProjects().get(0); + Assertions.assertEquals(new IntegerLiteral(5), ((Alias) constantItem).child()); + Assertions.assertNotEquals(studentId.getExprId(), constantItem.getExprId()); + // and the sink follows it, otherwise it would reference a slot which its child no longer outputs + Assertions.assertEquals(constantItem.getExprId(), rewrittenSink.getOutputExprs().get(0).getExprId()); + List childOutputIds = rewrittenSink.child().getOutput().stream() + .map(Slot::getExprId).collect(Collectors.toList()); + for (NamedExpression outputExpr : rewrittenSink.getOutputExprs()) { + Assertions.assertTrue(childOutputIds.contains(outputExpr.getExprId()), + "the sink expects " + outputExpr + " which its child does not output: " + + rewrittenSink.child().getOutput()); + } + } + + @Test + public void testBlackholeSinkKeepsTheOutputExprIds() { + // `warm up select` is planned with a blackhole sink which takes its output expressions from its child, so + // the sink has to follow the fresh exprIds of the constants materialized below it, otherwise it expects a + // slot which the rewritten child does not output any more + LogicalFilter filter = new LogicalFilter<>(ImmutableSet.of( + new EqualTo(studentId, new IntegerLiteral(5)), + new EqualTo(studentGender, new StringLiteral("F"))), student); + // all grouping keys are constant, so the aggregate keeps the first one (and drops the other), the constant + // above the kept key shadows it and must get a fresh exprId + List groupby = ImmutableList.of(studentId, studentGender); + List aggOutput = ImmutableList.of( + new Alias(new Sum(studentAge), "sum_age"), + studentId, + studentGender); + LogicalAggregate> aggregate + = new LogicalAggregate<>(groupby, aggOutput, filter); + List sinkOutputExprs = ImmutableList.copyOf(aggregate.getOutput()); + LogicalBlackholeSink>> sink + = new LogicalBlackholeSink<>(sinkOutputExprs, aggregate); + + LogicalSink rewrittenSink = (LogicalSink) executor.rewriteRoot(sink, jobContext); + List childOutputIds = rewrittenSink.child().getOutput().stream() + .map(Slot::getExprId).collect(java.util.stream.Collectors.toList()); + for (NamedExpression outputExpr : rewrittenSink.getOutputExprs()) { + Assertions.assertTrue(childOutputIds.contains(outputExpr.getExprId()), + "the sink expects " + outputExpr + " which its child does not output: " + + rewrittenSink.child().getOutput()); + } + // the constant column of the sink is refreshed to the fresh exprId of the constant project + // constant propagation keeps the exprIds here, the sink still references the output of its child + Assertions.assertEquals(studentId.getExprId(), + rewrittenSink.getOutputExprs().get(1).getExprId()); + } + + /** + * The aggregate still outputs the replaced column, so the constant project item would shadow it and must be + * materialized with a fresh exprId. + */ + private static void assertFreshConstantProjectItem(NamedExpression actual, ExprId replacedExprId, + Expression expectedExpr, String expectedName) { + Alias alias = assertConstantProjectItem(actual, expectedExpr, expectedName); + Assertions.assertNotEquals(replacedExprId, alias.getExprId(), + "the constant project item must not reuse the exprId of an output which the aggregate still has"); + } + + /** + * The aggregate removed the constant grouping key from its outputs, so the project item is the only definition + * of the column: its exprId is kept, rules and consumers such as the materialized view matching rely on it. + */ + private static void assertKeptConstantProjectItem(NamedExpression actual, ExprId replacedExprId, + Expression expectedExpr, String expectedName) { + Alias alias = assertConstantProjectItem(actual, expectedExpr, expectedName); + Assertions.assertEquals(replacedExprId, alias.getExprId(), + "the constant project item is the only definition, it must keep the exprId of the column"); + } + + private static Alias assertConstantProjectItem(NamedExpression actual, Expression expectedExpr, + String expectedName) { + Assertions.assertTrue(actual instanceof Alias, actual.toString()); + Alias alias = (Alias) actual; + Assertions.assertEquals(expectedExpr, alias.child()); + if (expectedName != null) { + Assertions.assertEquals(expectedName, alias.getName()); + } + return alias; } @Test diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CteConsumerOutputsRefresherTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CteConsumerOutputsRefresherTest.java new file mode 100644 index 00000000000000..ae7f4cf55a29e5 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CteConsumerOutputsRefresherTest.java @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.rules.rewrite; + +import org.apache.doris.nereids.trees.expressions.CTEId; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.plans.RelationId; +import org.apache.doris.nereids.trees.plans.logical.LogicalCTEConsumer; +import org.apache.doris.nereids.types.IntegerType; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableMultimap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Map; +import java.util.stream.Collectors; + +/** + * A CTE producer body is rewritten in its own subtree context, while the consumers, which map to the producer + * outputs by exprId, were rewritten before it, so RewriteCteChildren refreshes the two maps of the consumers with + * the exprIds the rules of the producer body published. + */ +public class CteConsumerOutputsRefresherTest { + private static final CTEId CTE_ID = new CTEId(0); + private static final ExprId PRODUCER_K = new ExprId(0); + private static final ExprId PRODUCER_V = new ExprId(1); + private static final ExprId CONSUMER_K = new ExprId(10); + private static final ExprId CONSUMER_V = new ExprId(11); + + @Test + void testRefreshFollowsTheReplacementChain() { + // k is renamed twice, e.g. by two rules or by two passes of the same rule + Map remap = ImmutableMap.of(PRODUCER_K, new ExprId(2), new ExprId(2), new ExprId(3)); + Slot newProducerK = slot(new ExprId(3), "k"); + LogicalCTEConsumer consumer = consumer(slot(PRODUCER_K, "k"), slot(PRODUCER_V, "v")); + + LogicalCTEConsumer refreshed = refresh(consumer, remap, ImmutableList.of(newProducerK, slot(PRODUCER_V, "v"))); + Assertions.assertEquals(newProducerK, refreshed.getConsumerToProducerOutputMap().get(slot(CONSUMER_K, "k"))); + Assertions.assertEquals(ImmutableList.of(slot(CONSUMER_K, "k")), + ImmutableList.copyOf(refreshed.getProducerToConsumerOutputMap().get(newProducerK))); + // v was not renamed, it keeps pointing to the same producer slot + Assertions.assertEquals(slot(PRODUCER_V, "v"), + refreshed.getConsumerToProducerOutputMap().get(slot(CONSUMER_V, "v"))); + } + + @Test + void testRenamedColumnWhichIsPrunedAfterwardsKeepsItsSlot() { + // the rule renamed k, but the rewrite of the producer body pruned the column afterwards, so the producer + // does not output the new exprId any more and there is nothing to map for it + Map remap = ImmutableMap.of(PRODUCER_K, new ExprId(3)); + Slot producerV = slot(PRODUCER_V, "v"); + LogicalCTEConsumer consumer = consumer(slot(PRODUCER_K, "k"), producerV); + + LogicalCTEConsumer refreshed = refresh(consumer, remap, ImmutableList.of(producerV)); + Assertions.assertEquals(slot(PRODUCER_K, "k"), + refreshed.getConsumerToProducerOutputMap().get(slot(CONSUMER_K, "k"))); + } + + @Test + void testConsumerOfAnotherCteIsNotTouched() { + Slot producerK = slot(PRODUCER_K, "k"); + LogicalCTEConsumer otherConsumer = new LogicalCTEConsumer(new RelationId(1), new CTEId(1), "other", + ImmutableMap.of(slot(CONSUMER_K, "k"), producerK), + ImmutableMultimap.of(producerK, slot(CONSUMER_K, "k"))); + LogicalCTEConsumer refreshed = (LogicalCTEConsumer) otherConsumer.accept( + new RewriteCteChildren.CteConsumerOutputsRefresher(CTE_ID, ImmutableMap.of(PRODUCER_K, new ExprId(3)), + ImmutableMap.of(new ExprId(3), slot(new ExprId(3), "k"))), null); + Assertions.assertSame(otherConsumer, refreshed); + } + + private static LogicalCTEConsumer refresh(LogicalCTEConsumer consumer, Map remap, + java.util.List producerOutputs) { + Map producerOutputSlots = producerOutputs.stream() + .collect(Collectors.toMap(Slot::getExprId, slot -> slot)); + return (LogicalCTEConsumer) consumer.accept(new RewriteCteChildren.CteConsumerOutputsRefresher( + CTE_ID, remap, producerOutputSlots), null); + } + + private static LogicalCTEConsumer consumer(Slot producerK, Slot producerV) { + return new LogicalCTEConsumer(new RelationId(0), CTE_ID, "c", + ImmutableMap.of(slot(CONSUMER_K, "k"), producerK, slot(CONSUMER_V, "v"), producerV), + ImmutableMultimap.of(producerK, slot(CONSUMER_K, "k"), producerV, slot(CONSUMER_V, "v"))); + } + + private static Slot slot(ExprId exprId, String name) { + return new SlotReference(exprId, name, IntegerType.INSTANCE, true, ImmutableList.of("test", "t")); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriterTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriterTest.java index 1acef3d6e74bea..5559961e36732d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriterTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriterTest.java @@ -17,12 +17,20 @@ package org.apache.doris.nereids.rules.rewrite; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.jobs.JobContext; import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.nereids.util.PlanConstructor; import com.google.common.collect.ImmutableList; import org.junit.jupiter.api.Assertions; @@ -47,6 +55,33 @@ */ public class ExprIdRewriterTest { + // The guarded-branch test of the master version (testGuardedBranchIsRewritten) is not ported: it drives + // ExprIdRewriter with MergeUtils.selectBranch, which wraps the branches in a ShortCircuitIf. Neither + // MergeUtils nor ShortCircuitIf exists on branch-4.2 (MERGE INTO does not wrap its assignments that way + // here), so the subtree which the master fix had to enter does not exist on this branch. + + @Test + void testAliasIdIsRewrittenThroughThePlanRewrite() { + // a rule which renames an Alias (e.g. EliminateGroupByKeyByUniform) publishes the alias exprId itself, and + // the plan rewrite has to replace it, not only the exprIds of the slots below it + ExprId oldAliasExprId = new ExprId(8); + ExprId newAliasExprId = new ExprId(10); + NamedExpression item = new Alias(oldAliasExprId, new IntegerLiteral(200), "field1"); + + LogicalOlapScan scan = new LogicalOlapScan(PlanConstructor.getNextRelationId(), + PlanConstructor.student, ImmutableList.of("")); + LogicalProject project = new LogicalProject<>(ImmutableList.of(item), scan); + Map replaceMap = new HashMap<>(); + replaceMap.put(oldAliasExprId, newAliasExprId); + CascadesContext cascadesContext = MemoTestUtils.createCascadesContext(scan); + Plan rewritten = new ExprIdRewriter(new ExprIdRewriter.ReplaceRule(replaceMap, true), + new JobContext(cascadesContext, null, Double.MAX_VALUE)).rewriteExpr(project, replaceMap); + + NamedExpression rewrittenItem = ((LogicalProject) rewritten).getProjects().get(0); + Assertions.assertEquals(newAliasExprId, rewrittenItem.getExprId(), rewrittenItem.toString()); + Assertions.assertEquals("field1", rewrittenItem.getName()); + } + /** * Test that visitAlias correctly rewrites Alias ExprId. *