RewriteCteChildren.java
// 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.CascadesContext;
import org.apache.doris.nereids.annotation.DependsRules;
import org.apache.doris.nereids.jobs.JobContext;
import org.apache.doris.nereids.jobs.executor.Rewriter;
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;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.plans.LimitPhase;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.RelationId;
import org.apache.doris.nereids.trees.plans.logical.LogicalCTEAnchor;
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.LogicalFilter;
import org.apache.doris.nereids.trees.plans.logical.LogicalLimit;
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.visitor.CustomRewriter;
import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter;
import org.apache.doris.nereids.util.ExpressionUtils;
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;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
/**
* rewrite CteAnchor consumer side and producer side recursively��� all CteAnchor must at top of the plan
*/
@DependsRules({PullUpCteAnchor.class, CTEInline.class})
public class RewriteCteChildren extends DefaultPlanRewriter<CascadesContext> implements CustomRewriter {
private final List<RewriteJob> jobs;
private final boolean runCboRules;
public RewriteCteChildren(List<RewriteJob> jobs, boolean runCboRules) {
this.jobs = jobs;
this.runCboRules = runCboRules;
}
@Override
public Plan rewriteRoot(Plan plan, JobContext jobContext) {
return plan.accept(this, jobContext.getCascadesContext());
}
@Override
public Plan visit(Plan plan, CascadesContext context) {
Rewriter.getCteChildrenRewriter(context, jobs, runCboRules).execute();
return context.getRewritePlan();
}
@Override
public Plan visitLogicalCTEAnchor(LogicalCTEAnchor<? extends Plan, ? extends Plan> cteAnchor,
CascadesContext cascadesContext) {
LogicalPlan outer;
if (cascadesContext.getStatementContext().getRewrittenCteConsumer().containsKey(cteAnchor.getCteId())) {
outer = cascadesContext.getStatementContext().getRewrittenCteConsumer().get(cteAnchor.getCteId());
} else {
CascadesContext outerCascadesCtx = CascadesContext.newSubtreeContext(
Optional.empty(), cascadesContext, cteAnchor.child(1),
cascadesContext.getCurrentJobContext().getRequiredProperties());
AtomicReference<LogicalPlan> outerResult = new AtomicReference<>();
StatsDerive statsDerive = new StatsDerive(false);
cteAnchor.child(0).accept(statsDerive, new StatsDerive.DeriveContext());
outerCascadesCtx.withPlanProcess(cascadesContext.showPlanProcess(), () -> {
outerResult.set((LogicalPlan) cteAnchor.child(1).accept(this, outerCascadesCtx));
});
outer = outerResult.get();
cascadesContext.addPlanProcesses(outerCascadesCtx.getPlanProcesses());
cascadesContext.getStatementContext().getRewrittenCteConsumer().put(cteAnchor.getCteId(), outer);
}
Set<LogicalCTEConsumer> 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<ExprId, ExprId> producerExprIdRemap = cascadesContext.getStatementContext()
.getProducerExprIdRemap(cteAnchor.getCteId());
if (!producerExprIdRemap.isEmpty()) {
Map<ExprId, Slot> 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.<Plan, Void>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<LogicalCTEConsumer> collectCteConsumers(Plan outer, CTEId cteId) {
Set<LogicalCTEConsumer> cteConsumers = Sets.newHashSet();
outer.foreach(p -> {
if (p instanceof LogicalCTEConsumer) {
LogicalCTEConsumer logicalCTEConsumer = (LogicalCTEConsumer) p;
if (logicalCTEConsumer.getCteId().equals(cteId)) {
cteConsumers.add(logicalCTEConsumer);
}
}
return false;
});
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<Void> {
private final CTEId cteId;
private final Map<ExprId, ExprId> producerExprIdRemap;
private final Map<ExprId, Slot> producerOutputSlots;
CteConsumerOutputsRefresher(CTEId cteId, Map<ExprId, ExprId> producerExprIdRemap,
Map<ExprId, Slot> 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<Slot, Slot> consumerToProducerOutputMap = new LinkedHashMap<>();
Multimap<Slot, Slot> producerToConsumerOutputMap = LinkedHashMultimap.create();
for (Map.Entry<Slot, Slot> 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;
}
}
@Override
public Plan visitLogicalCTEProducer(LogicalCTEProducer<? extends Plan> cteProducer,
CascadesContext cascadesContext) {
LogicalPlan child;
if (cascadesContext.getStatementContext().getRewrittenCteProducer().containsKey(cteProducer.getCteId())) {
child = cascadesContext.getStatementContext().getRewrittenCteProducer().get(cteProducer.getCteId());
} else {
child = (LogicalPlan) cteProducer.child();
child = tryToConstructFilter(cascadesContext, cteProducer.getCteId(), child);
child = tryToConstructLimit(cascadesContext, cteProducer.getCteId(), child);
Set<Slot> producerOutputs = cascadesContext.getStatementContext()
.getCteIdToOutputIds().get(cteProducer.getCteId());
if (producerOutputs != null && producerOutputs.size() < child.getOutput().size()) {
ImmutableList.Builder<NamedExpression> projectsBuilder
= ImmutableList.builderWithExpectedSize(producerOutputs.size());
for (Slot slot : child.getOutput()) {
if (producerOutputs.contains(slot)) {
projectsBuilder.add(slot);
}
}
child = new LogicalProject<>(projectsBuilder.build(), child);
child = pushPlanUnderAnchor(child);
}
CascadesContext rewrittenCtx = CascadesContext.newSubtreeContext(
Optional.of(cteProducer.getCteId()), cascadesContext, child, PhysicalProperties.ANY);
AtomicReference<LogicalPlan> result = new AtomicReference<>(child);
rewrittenCtx.withPlanProcess(cascadesContext.showPlanProcess(), () -> {
result.set((LogicalPlan) result.get().accept(this, rewrittenCtx));
});
child = result.get();
cascadesContext.addPlanProcesses(rewrittenCtx.getPlanProcesses());
cascadesContext.getStatementContext().getRewrittenCteProducer().put(cteProducer.getCteId(), child);
}
LogicalCTEProducer<? extends Plan> rewrittenProducer = (LogicalCTEProducer<? extends Plan>) cteProducer
.withChildren(child);
cascadesContext.getStatementContext().setCteProducer(rewrittenProducer.getCteId(), rewrittenProducer);
return rewrittenProducer;
}
private LogicalPlan pushPlanUnderAnchor(LogicalPlan plan) {
if (plan.child(0) instanceof LogicalCTEAnchor) {
LogicalPlan child = (LogicalPlan) plan.withChildren(plan.child(0).child(1));
return (LogicalPlan) plan.child(0).withChildren(
plan.child(0).child(0), pushPlanUnderAnchor(child));
}
return plan;
}
private LogicalPlan tryToConstructLimit(CascadesContext cascadesContext, CTEId cteId, LogicalPlan child) {
Set<LogicalCTEConsumer> consumers = cascadesContext.getCteIdToConsumers().get(cteId);
long limit = 0;
for (LogicalCTEConsumer consumer : consumers) {
Long rowsNeeded = cascadesContext.getConsumerIdToLimitRows().get(consumer.getRelationId());
if (rowsNeeded == null) {
return child;
}
limit = Math.max(limit, rowsNeeded);
}
return pushPlanUnderAnchor(new LogicalLimit<>(limit, 0, LimitPhase.ORIGIN, child));
}
/*
* An expression can only be pushed down if it has filter expressions on all consumers that reference the slot.
* For example, let's assume a producer has two consumers, consumer1 and consumer2:
*
* filter(a > 5 and b < 1) -> consumer1
* filter(a < 8) -> consumer2
*
* In this case, the only expression that can be pushed down to the producer is filter(a > 5 or a < 8).
*/
private LogicalPlan tryToConstructFilter(CascadesContext cascadesContext, CTEId cteId, LogicalPlan child) {
Set<RelationId> consumerIds = cascadesContext.getCteIdToConsumers().get(cteId).stream()
.map(LogicalCTEConsumer::getRelationId)
.collect(Collectors.toSet());
List<Set<Expression>> filtersAboveEachConsumer = cascadesContext.getConsumerIdToFilters().entrySet().stream()
.filter(kv -> consumerIds.contains(kv.getKey()))
.map(Entry::getValue)
.collect(Collectors.toList());
Set<Expression> someone = filtersAboveEachConsumer.stream().findFirst().orElse(null);
if (someone == null) {
return child;
}
int filterSize = cascadesContext.getCteIdToConsumers().get(cteId).size();
Set<Expression> conjuncts = new HashSet<>();
for (Expression f : someone) {
int matchCount = 0;
Set<SlotReference> slots = f.collect(e -> e instanceof SlotReference);
Set<Expression> mightBeJoined = new HashSet<>();
for (Set<Expression> another : filtersAboveEachConsumer) {
if (another.equals(someone)) {
matchCount++;
continue;
}
Set<Expression> matched = new HashSet<>();
for (Expression e : another) {
Set<SlotReference> otherSlots = e.collect(ae -> ae instanceof SlotReference);
if (otherSlots.equals(slots)) {
matched.add(e);
}
}
if (!matched.isEmpty()) {
matchCount++;
}
mightBeJoined.addAll(matched);
}
if (matchCount >= filterSize) {
mightBeJoined.add(f);
conjuncts.add(ExpressionUtils.or(mightBeJoined));
}
}
if (!conjuncts.isEmpty()) {
// distinct conjuncts
ImmutableSet.Builder<Expression> newConjuncts = ImmutableSet.builderWithExpectedSize(conjuncts.size());
for (Expression conjunct : conjuncts) {
newConjuncts.addAll(ExpressionUtils.extractConjunction(conjunct));
}
LogicalPlan filter = new LogicalFilter<>(newConjuncts.build(), child);
return pushPlanUnderAnchor(filter);
}
return child;
}
}