Skip to content

Commit

Permalink
[fix](Nereids) producer to consumer should be multimap in cte (#39850)
Browse files Browse the repository at this point in the history 
because consumer could refer multi times for one producer' slot, so
producer to consumer slot map should be multimap
  • Loading branch information
@morrySnow
morrySnow authored Aug 23, 2024
1 parent 2c8d65b commit d3b48a4
Show file tree
Hide file tree
Showing 10 changed files with 100 additions and 64 deletions.
19 changes: 13 additions & 6 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

Expand Down Expand Up @@ -618,8 +619,8 @@ public Map<RelationId, Set<Expression>> getConsumerIdToFilters() {
return this.statementContext.getConsumerIdToFilters();
}

public void addCTEConsumerGroup(CTEId cteId, Group g, Map<Slot, Slot> producerSlotToConsumerSlot) {
List<Pair<Map<Slot, Slot>, Group>> consumerGroups =
public void addCTEConsumerGroup(CTEId cteId, Group g, Multimap<Slot, Slot> producerSlotToConsumerSlot) {
List<Pair<Multimap<Slot, Slot>, Group>> consumerGroups =
this.statementContext.getCteIdToConsumerGroup().computeIfAbsent(cteId, k -> new ArrayList<>());
consumerGroups.add(Pair.of(producerSlotToConsumerSlot, g));
}
Expand All @@ -628,12 +629,18 @@ public void addCTEConsumerGroup(CTEId cteId, Group g, Map<Slot, Slot> producerSl
* Update CTE consumer group as producer's stats update
*/
public void updateConsumerStats(CTEId cteId, Statistics statistics) {
List<Pair<Map<Slot, Slot>, Group>> consumerGroups = this.statementContext.getCteIdToConsumerGroup().get(cteId);
for (Pair<Map<Slot, Slot>, Group> p : consumerGroups) {
Map<Slot, Slot> producerSlotToConsumerSlot = p.first;
List<Pair<Multimap<Slot, Slot>, Group>> consumerGroups
= this.statementContext.getCteIdToConsumerGroup().get(cteId);
for (Pair<Multimap<Slot, Slot>, Group> p : consumerGroups) {
Multimap<Slot, Slot> producerSlotToConsumerSlot = p.first;
Statistics updatedConsumerStats = new StatisticsBuilder(statistics).build();
for (Entry<Expression, ColumnStatistic> entry : statistics.columnStatistics().entrySet()) {
updatedConsumerStats.addColumnStats(producerSlotToConsumerSlot.get(entry.getKey()), entry.getValue());
if (!(entry.getKey() instanceof Slot)) {
continue;
}
for (Slot consumer : producerSlotToConsumerSlot.get((Slot) entry.getKey())) {
updatedConsumerStats.addColumnStats(consumer, entry.getValue());
}
}
p.value().setStatistics(updatedConsumerStats);
}
Expand Down
5 changes: 3 additions & 2 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
Expand Down Expand Up @@ -115,7 +116,7 @@ public class StatementContext implements Closeable {
private final Map<CTEId, Set<Slot>> cteIdToOutputIds = new HashMap<>();
private final Map<RelationId, Set<Expression>> consumerIdToFilters = new HashMap<>();
// Used to update consumer's stats
private final Map<CTEId, List<Pair<Map<Slot, Slot>, Group>>> cteIdToConsumerGroup = new HashMap<>();
private final Map<CTEId, List<Pair<Multimap<Slot, Slot>, Group>>> cteIdToConsumerGroup = new HashMap<>();
private final Map<CTEId, LogicalPlan> rewrittenCteProducer = new HashMap<>();
private final Map<CTEId, LogicalPlan> rewrittenCteConsumer = new HashMap<>();
private final Set<String> viewDdlSqlSet = Sets.newHashSet();
Expand Down Expand Up @@ -357,7 +358,7 @@ public Map<PlaceholderId, Expression> getIdToPlaceholderRealExpr() {
return idToPlaceholderRealExpr;
}

public Map<CTEId, List<Pair<Map<Slot, Slot>, Group>>> getCteIdToConsumerGroup() {
public Map<CTEId, List<Pair<Multimap<Slot, Slot>, Group>>> getCteIdToConsumerGroup() {
return cteIdToConsumerGroup;
}

Expand Down
7 changes: 2 additions & 5 deletions ...e-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1205,12 +1205,9 @@ public PlanFragment visitPhysicalCTEConsumer(PhysicalCTEConsumer cteConsumer,
for (Slot producerSlot : cteProducer.getOutput()) {
SlotRef slotRef = context.findSlotRef(producerSlot.getExprId());
tupleDescriptor = slotRef.getDesc().getParent();
Slot consumerSlot = cteConsumer.getProducerToConsumerSlotMap().get(producerSlot);
// consumerSlot could be null if we prune partial consumers' columns
if (consumerSlot == null) {
continue;
for (Slot consumerSlot : cteConsumer.getProducerToConsumerSlotMap().get(producerSlot)) {
context.addExprIdSlotRefPair(consumerSlot.getExprId(), slotRef);
}
context.addExprIdSlotRefPair(consumerSlot.getExprId(), slotRef);
}
CTEScanNode cteScanNode = new CTEScanNode(tupleDescriptor);
context.getRuntimeTranslator().ifPresent(runtimeFilterTranslator ->
Expand Down
15 changes: 9 additions & 6 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,10 @@
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.LinkedHashMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;

import java.util.LinkedHashMap;
import java.util.List;
Expand Down Expand Up @@ -261,14 +263,15 @@ public Plan visitLogicalPartitionTopN(LogicalPartitionTopN<? extends Plan> parti
@Override
public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, Map<ExprId, Slot> replaceMap) {
Map<Slot, Slot> consumerToProducerOutputMap = new LinkedHashMap<>();
Map<Slot, Slot> producerToConsumerOutputMap = new LinkedHashMap<>();
Multimap<Slot, Slot> producerToConsumerOutputMap = LinkedHashMultimap.create();
for (Slot producerOutputSlot : cteConsumer.getConsumerToProducerOutputMap().values()) {
Slot newProducerOutputSlot = updateExpression(producerOutputSlot, replaceMap);
Slot newConsumerOutputSlot = cteConsumer.getProducerToConsumerOutputMap().get(producerOutputSlot)
.withNullable(newProducerOutputSlot.nullable());
producerToConsumerOutputMap.put(newProducerOutputSlot, newConsumerOutputSlot);
consumerToProducerOutputMap.put(newConsumerOutputSlot, newProducerOutputSlot);
replaceMap.put(newConsumerOutputSlot.getExprId(), newConsumerOutputSlot);
for (Slot consumerOutputSlot : cteConsumer.getProducerToConsumerOutputMap().get(producerOutputSlot)) {
Slot newConsumerOutputSlot = consumerOutputSlot.withNullable(newProducerOutputSlot.nullable());
producerToConsumerOutputMap.put(newProducerOutputSlot, newConsumerOutputSlot);
consumerToProducerOutputMap.put(newConsumerOutputSlot, newProducerOutputSlot);
replaceMap.put(newConsumerOutputSlot.getExprId(), newConsumerOutputSlot);
}
}
return cteConsumer.withTwoMaps(consumerToProducerOutputMap, producerToConsumerOutputMap);
}
Expand Down
8 changes: 4 additions & 4 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/OrExpansion.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -206,11 +206,11 @@ public Plan visitLogicalJoin(LogicalJoin<? extends Plan, ? extends Plan> join, O
private Map<Slot, Slot> constructReplaceMap(LogicalCTEConsumer leftConsumer, Map<Slot, Slot> leftCloneToLeft,
LogicalCTEConsumer rightConsumer, Map<Slot, Slot> rightCloneToRight) {
Map<Slot, Slot> replaced = new HashMap<>();
for (Entry<Slot, Slot> entry : leftConsumer.getProducerToConsumerOutputMap().entrySet()) {
replaced.put(leftCloneToLeft.get(entry.getKey()), entry.getValue());
for (Entry<Slot, Slot> entry : leftConsumer.getConsumerToProducerOutputMap().entrySet()) {
replaced.put(leftCloneToLeft.get(entry.getValue()), entry.getKey());
}
for (Entry<Slot, Slot> entry : rightConsumer.getProducerToConsumerOutputMap().entrySet()) {
replaced.put(rightCloneToRight.get(entry.getKey()), entry.getValue());
for (Entry<Slot, Slot> entry : rightConsumer.getConsumerToProducerOutputMap().entrySet()) {
replaced.put(rightCloneToRight.get(entry.getValue()), entry.getKey());
}
return replaced;
}
Expand Down
4 changes: 3 additions & 1 deletion fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/VariantSubPathPruning.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,10 @@

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.LinkedHashMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;

import java.util.Collections;
Expand Down Expand Up @@ -394,7 +396,7 @@ public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, Context cont
return cteConsumer;
}
Map<Slot, Slot> consumerToProducerOutputMap = Maps.newHashMap();
Map<Slot, Slot> producerToConsumerOutputMap = Maps.newHashMap();
Multimap<Slot, Slot> producerToConsumerOutputMap = LinkedHashMultimap.create();
Map<Slot, Map<List<String>, SlotReference>> oriSlotToSubPathToSlot = Maps.newHashMap();
for (Map.Entry<Slot, Slot> consumerToProducer : cteConsumer.getConsumerToProducerOutputMap().entrySet()) {
Slot consumer = consumerToProducer.getKey();
Expand Down
4 changes: 3 additions & 1 deletion fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,8 @@

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 java.util.LinkedHashMap;
import java.util.List;
Expand Down Expand Up @@ -397,7 +399,7 @@ public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, DeepCopierCo
return context.getRelationReplaceMap().get(cteConsumer.getRelationId());
}
Map<Slot, Slot> consumerToProducerOutputMap = new LinkedHashMap<>();
Map<Slot, Slot> producerToConsumerOutputMap = new LinkedHashMap<>();
Multimap<Slot, Slot> producerToConsumerOutputMap = LinkedHashMultimap.create();
for (Slot consumerOutput : cteConsumer.getOutput()) {
Slot newOutput = (Slot) ExpressionDeepCopier.INSTANCE.deepCopy(consumerOutput, context);
consumerToProducerOutputMap.put(newOutput, cteConsumer.getProducerSlot(consumerOutput));
Expand Down
57 changes: 28 additions & 29 deletions ...e-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,10 @@

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.LinkedHashMultimap;
import com.google.common.collect.Multimap;

import java.util.LinkedHashMap;
import java.util.List;
Expand All @@ -50,20 +54,15 @@ public class LogicalCTEConsumer extends LogicalRelation implements BlockFuncDeps
private final String name;
private final CTEId cteId;
private final Map<Slot, Slot> consumerToProducerOutputMap;
private final Map<Slot, Slot> producerToConsumerOutputMap;
private final Multimap<Slot, Slot> producerToConsumerOutputMap;

/**
* Logical CTE consumer.
*/
public LogicalCTEConsumer(RelationId relationId, CTEId cteId, String name,
Map<Slot, Slot> consumerToProducerOutputMap, Map<Slot, Slot> producerToConsumerOutputMap) {
super(relationId, PlanType.LOGICAL_CTE_CONSUMER, Optional.empty(), Optional.empty());
this.cteId = Objects.requireNonNull(cteId, "cteId should not null");
this.name = Objects.requireNonNull(name, "name should not null");
this.consumerToProducerOutputMap = Objects.requireNonNull(consumerToProducerOutputMap,
"consumerToProducerOutputMap should not null");
this.producerToConsumerOutputMap = Objects.requireNonNull(producerToConsumerOutputMap,
"producerToConsumerOutputMap should not null");
Map<Slot, Slot> consumerToProducerOutputMap, Multimap<Slot, Slot> producerToConsumerOutputMap) {
this(relationId, cteId, name, consumerToProducerOutputMap, producerToConsumerOutputMap,
Optional.empty(), Optional.empty());
}

/**
Expand All @@ -73,24 +72,31 @@ public LogicalCTEConsumer(RelationId relationId, CTEId cteId, String name, Logic
super(relationId, PlanType.LOGICAL_CTE_CONSUMER, Optional.empty(), Optional.empty());
this.cteId = Objects.requireNonNull(cteId, "cteId should not null");
this.name = Objects.requireNonNull(name, "name should not null");
this.consumerToProducerOutputMap = new LinkedHashMap<>();
this.producerToConsumerOutputMap = new LinkedHashMap<>();
initOutputMaps(producerPlan);
ImmutableMap.Builder<Slot, Slot> cToPBuilder = ImmutableMap.builder();
ImmutableMultimap.Builder<Slot, Slot> pToCBuilder = ImmutableMultimap.builder();
List<Slot> producerOutput = producerPlan.getOutput();
for (Slot producerOutputSlot : producerOutput) {
Slot consumerSlot = generateConsumerSlot(this.name, producerOutputSlot);
cToPBuilder.put(consumerSlot, producerOutputSlot);
pToCBuilder.put(producerOutputSlot, consumerSlot);
}
consumerToProducerOutputMap = cToPBuilder.build();
producerToConsumerOutputMap = pToCBuilder.build();
}

/**
* Logical CTE consumer.
*/
public LogicalCTEConsumer(RelationId relationId, CTEId cteId, String name,
Map<Slot, Slot> consumerToProducerOutputMap, Map<Slot, Slot> producerToConsumerOutputMap,
Map<Slot, Slot> consumerToProducerOutputMap, Multimap<Slot, Slot> producerToConsumerOutputMap,
Optional<GroupExpression> groupExpression, Optional<LogicalProperties> logicalProperties) {
super(relationId, PlanType.LOGICAL_CTE_CONSUMER, groupExpression, logicalProperties);
this.cteId = Objects.requireNonNull(cteId, "cteId should not null");
this.name = Objects.requireNonNull(name, "name should not null");
this.consumerToProducerOutputMap = Objects.requireNonNull(consumerToProducerOutputMap,
"consumerToProducerOutputMap should not null");
this.producerToConsumerOutputMap = Objects.requireNonNull(producerToConsumerOutputMap,
"producerToConsumerOutputMap should not null");
this.consumerToProducerOutputMap = ImmutableMap.copyOf(Objects.requireNonNull(consumerToProducerOutputMap,
"consumerToProducerOutputMap should not null"));
this.producerToConsumerOutputMap = ImmutableMultimap.copyOf(Objects.requireNonNull(producerToConsumerOutputMap,
"producerToConsumerOutputMap should not null"));
}

/**
Expand All @@ -107,20 +113,11 @@ public static SlotReference generateConsumerSlot(String cteName, Slot producerOu
slotRef != null ? Optional.of(slotRef.getInternalName()) : Optional.empty());
}

private void initOutputMaps(LogicalPlan childPlan) {
List<Slot> producerOutput = childPlan.getOutput();
for (Slot producerOutputSlot : producerOutput) {
Slot consumerSlot = generateConsumerSlot(this.name, producerOutputSlot);
producerToConsumerOutputMap.put(producerOutputSlot, consumerSlot);
consumerToProducerOutputMap.put(consumerSlot, producerOutputSlot);
}
}

public Map<Slot, Slot> getConsumerToProducerOutputMap() {
return consumerToProducerOutputMap;
}

public Map<Slot, Slot> getProducerToConsumerOutputMap() {
public Multimap<Slot, Slot> getProducerToConsumerOutputMap() {
return producerToConsumerOutputMap;
}

Expand All @@ -129,7 +126,8 @@ public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
return visitor.visitLogicalCTEConsumer(this, context);
}

public Plan withTwoMaps(Map<Slot, Slot> consumerToProducerOutputMap, Map<Slot, Slot> producerToConsumerOutputMap) {
public Plan withTwoMaps(Map<Slot, Slot> consumerToProducerOutputMap,
Multimap<Slot, Slot> producerToConsumerOutputMap) {
return new LogicalCTEConsumer(relationId, cteId, name,
consumerToProducerOutputMap, producerToConsumerOutputMap);
}
Expand Down Expand Up @@ -162,7 +160,8 @@ public List<Slot> computeOutput() {
@Override
public Plan pruneOutputs(List<NamedExpression> prunedOutputs) {
Map<Slot, Slot> consumerToProducerOutputMap = new LinkedHashMap<>(this.consumerToProducerOutputMap.size());
Map<Slot, Slot> producerToConsumerOutputMap = new LinkedHashMap<>(this.consumerToProducerOutputMap.size());
Multimap<Slot, Slot> producerToConsumerOutputMap = LinkedHashMultimap.create(
this.consumerToProducerOutputMap.size(), this.consumerToProducerOutputMap.size());
for (Entry<Slot, Slot> consumerToProducerSlot : this.consumerToProducerOutputMap.entrySet()) {
if (prunedOutputs.contains(consumerToProducerSlot.getKey())) {
consumerToProducerOutputMap.put(consumerToProducerSlot.getKey(), consumerToProducerSlot.getValue());
Expand Down
Loading

0 comments on commit d3b48a4

Please sign in to comment.