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.SlotReplacementMapping;
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.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.Sets;
import java.util.HashSet;
import java.util.List;
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) {
// Derive the operative slots of the producer before the statistics of the consumer are
// derived below: StatsCalculator.computeOlapScan fetches the column stats of the slots that
// are marked as operative, and without this derivation every column of every scan of the
// producer would be fetched. The derived producer is the plan whose statistics are derived
// and which is rewritten below, so that those statistics are not computed on a discarded
// copy of it.
//
// The producer subtree does not tell by itself which of its columns are read: the demand
// comes from its consumers, and a forwarding project of the producer creates none, because a
// projection of a bare slot is not an expression the derivation counts as a demand. Seed the
// slots of the producer that the estimation above this anchor uses, so that exactly those
// columns are demanded instead of being only probed from the cache.
//
// The demand is collected from the consumer side first, with a fresh derivation context: an
// anchor is always at the top of the plan (see PullUpCteAnchor), so nothing above this anchor
// demands anything and the context of the enclosing derivation would be empty anyway.
OperativeColumnDerive.DeriveContext consumerDemand = new OperativeColumnDerive.DeriveContext();
cteAnchor.child(1).accept(new OperativeColumnDerive(), consumerDemand);
OperativeColumnDerive.DeriveContext deriveContext = new OperativeColumnDerive.DeriveContext();
seedProducerOperativeSlots(cteAnchor, cascadesContext, consumerDemand, deriveContext);
Plan derivedProducer = cteAnchor.child(0).accept(new OperativeColumnDerive(), deriveContext);
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);
derivedProducer.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 = Sets.newHashSet();
outer.foreach(p -> {
if (p instanceof LogicalCTEConsumer) {
LogicalCTEConsumer logicalCTEConsumer = (LogicalCTEConsumer) p;
if (logicalCTEConsumer.getCteId().equals(cteAnchor.getCteId())) {
cteConsumers.add(logicalCTEConsumer);
}
}
return false;
});
cascadesContext.getCteIdToConsumers().put(cteAnchor.getCteId(), cteConsumers);
if (cteConsumers.isEmpty()) {
return outer;
}
Plan producer = derivedProducer.accept(this, cascadesContext);
return cteAnchor.withChildren(producer, outer);
}
@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);
List<Slot> producerOutputsBeforeRewrite = child.getOutput();
AtomicReference<LogicalPlan> result = new AtomicReference<>(child);
rewrittenCtx.withPlanProcess(cascadesContext.showPlanProcess(), () -> {
result.set((LogicalPlan) result.get().accept(this, rewrittenCtx));
});
child = result.get();
recordProducerOutputReplacements(cascadesContext, producerOutputsBeforeRewrite, child.getOutput());
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;
}
/**
* Marks the slots of the producer of {@code cteAnchor} that the estimation of its consumers uses
* as operative, i.e. the union over all the consumers of the slots of the producer they demand.
*
* <p>The producer subtree alone does not tell which of its columns are read: the demand comes
* from the consumers, each of which copies the slot of the producer it reads, and the producer
* only forwards that slot through projections, which create no demand of their own. Without this
* seeding the derivation marks the whole output of the producer as not operative, and the
* statistics derivation only probes those columns from the cache: a miss leaves them unknown and
* the estimation of the consumer, e.g. the local distinct inference of a set operation over the
* CTE, then works on unknown statistics.
*
* <p>A slot can be replaced more than once while a statement is rewritten, so the slot of the
* producer that a consumer keeps is resolved through the {@link SlotReplacementMapping} of the
* statement, which follows the whole replacement chain, and not only its first replacement.
*/
static void seedProducerOperativeSlots(LogicalCTEAnchor<?, ?> cteAnchor,
CascadesContext cascadesContext, OperativeColumnDerive.DeriveContext consumerDemand,
OperativeColumnDerive.DeriveContext deriveContext) {
SlotReplacementMapping replacements = cascadesContext.getStatementContext()
.getSlotReplacementMapping();
cteAnchor.child(1).foreach(plan -> {
if (plan instanceof LogicalCTEConsumer) {
LogicalCTEConsumer consumer = (LogicalCTEConsumer) plan;
if (consumer.getCteId().equals(cteAnchor.getCteId())) {
for (Entry<Slot, Slot> consumerToProducer : consumer.getConsumerToProducerOutputMap().entrySet()) {
if (consumerDemand.operativeSlotIds
.contains(consumerToProducer.getKey().getExprId().asInt())) {
deriveContext.addOperativeSlotId(
replacements.resolveToInt(consumerToProducer.getValue().getExprId()));
}
}
}
}
return false;
});
}
/**
* Records the slots that the rewrite of the producer replaced, so that a slot of the producer
* that a consumer still references resolves to the slot that is current in the rewritten plan.
*
* <p>The output of the producer may also be pruned by the rewrite, in which case the outputs of
* the rewritten producer are not the same columns and there is no replacement to record.
*/
private static void recordProducerOutputReplacements(CascadesContext cascadesContext,
List<Slot> outputsBeforeRewrite, List<Slot> outputsAfterRewrite) {
if (outputsBeforeRewrite.size() != outputsAfterRewrite.size()) {
return;
}
SlotReplacementMapping replacements = cascadesContext.getStatementContext()
.getSlotReplacementMapping();
for (int i = 0; i < outputsBeforeRewrite.size(); i++) {
replacements.record(outputsBeforeRewrite.get(i).getExprId(),
outputsAfterRewrite.get(i).getExprId());
}
}
}