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 6f4b97b9403f69..6da14c14bcfc6d 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 @@ -54,6 +54,7 @@ import org.apache.doris.nereids.trees.plans.RelationId; import org.apache.doris.nereids.trees.plans.TableId; 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.LogicalPlan; import org.apache.doris.nereids.util.RelationUtil; import org.apache.doris.qe.ConnectContext; @@ -156,7 +157,7 @@ public enum TableFrom { private final Map> cteIdToConsumers = new HashMap<>(); private final Map> cteIdToOutputIds = new HashMap<>(); - private final Map cteIdToProducerStats = new HashMap<>(); + private final Map> cteIdToProducer = new HashMap<>(); private final Map> consumerIdToFilters = new HashMap<>(); // Used to update consumer's stats @@ -638,7 +639,7 @@ public CteEnvironmentSnapshot cacheCteEnvironment() { return new CteEnvironmentSnapshot( copyMapOfSets(cteIdToConsumers), copyMapOfSets(cteIdToOutputIds), - new HashMap<>(cteIdToProducerStats), + new HashMap<>(cteIdToProducer), copyMapOfSets(consumerIdToFilters), copyMapOfLists(cteIdToConsumerGroup), new HashMap<>(rewrittenCteProducer), @@ -653,8 +654,8 @@ public void restoreCteEnvironment(CteEnvironmentSnapshot snapshot) { cteIdToOutputIds.clear(); cteIdToOutputIds.putAll(snapshot.cteIdToOutputIds); - cteIdToProducerStats.clear(); - cteIdToProducerStats.putAll(snapshot.cteIdToProducerStats); + cteIdToProducer.clear(); + cteIdToProducer.putAll(snapshot.cteIdToProducer); consumerIdToFilters.clear(); consumerIdToFilters.putAll(snapshot.consumerIdToFilters); @@ -689,7 +690,7 @@ private static Map> copyMapOfLists(Map> source) { public static class CteEnvironmentSnapshot { private final Map> cteIdToConsumers; private final Map> cteIdToOutputIds; - private final Map cteIdToProducerStats; + private final Map> cteIdToProducer; private final Map> consumerIdToFilters; private final Map, Group>>> cteIdToConsumerGroup; private final Map rewrittenCteProducer; @@ -701,14 +702,14 @@ public static class CteEnvironmentSnapshot { public CteEnvironmentSnapshot( Map> cteIdToConsumers, Map> cteIdToOutputIds, - Map cteIdToProducerStats, + Map> cteIdToProducer, Map> consumerIdToFilters, Map, Group>>> cteIdToConsumerGroup, Map rewrittenCteProducer, Map rewrittenCteConsumer) { this.cteIdToConsumers = cteIdToConsumers; this.cteIdToOutputIds = cteIdToOutputIds; - this.cteIdToProducerStats = cteIdToProducerStats; + this.cteIdToProducer = cteIdToProducer; this.consumerIdToFilters = consumerIdToFilters; this.cteIdToConsumerGroup = cteIdToConsumerGroup; this.rewrittenCteProducer = rewrittenCteProducer; @@ -1111,12 +1112,12 @@ public boolean isPrepareStage() { return prepareStage; } - public Statistics getProducerStatsByCteId(CTEId id) { - return cteIdToProducerStats.get(id); + public LogicalCTEProducer getCteProducerByCteId(CTEId id) { + return cteIdToProducer.get(id); } - public void setProducerStats(CTEId id, Statistics stats) { - cteIdToProducerStats.put(id, stats); + public void setCteProducer(CTEId id, LogicalCTEProducer producer) { + cteIdToProducer.put(id, producer); } public void setIsInsert(boolean isInsert) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AnalyzeCTE.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AnalyzeCTE.java index 4287e8ca7dea2c..c49b3f01c51039 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AnalyzeCTE.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AnalyzeCTE.java @@ -109,7 +109,9 @@ private Pair>> analyzeCte( aliasQuery.withChildren(ImmutableList.of(analyzedCtePlan)); outerCteCtx = new CTEContext(cteId, logicalSubQueryAlias, outerCteCtx); outerCteCtx.setAnalyzedPlan(logicalSubQueryAlias); - cteProducerPlans.add(new LogicalCTEProducer<>(cteId, logicalSubQueryAlias)); + LogicalCTEProducer cteProducer = new LogicalCTEProducer<>(cteId, logicalSubQueryAlias); + cascadesContext.getStatementContext().setCteProducer(cteId, cteProducer); + cteProducerPlans.add(cteProducer); } return Pair.of(outerCteCtx, cteProducerPlans); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/OrExpansion.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/OrExpansion.java index 4cb07fd44fd89b..2048ea8f50e14d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/OrExpansion.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/OrExpansion.java @@ -141,10 +141,12 @@ public Plan visitLogicalJoin(LogicalJoin join, O .deepCopy((LogicalPlan) join.left(), new DeepCopierContext()); LogicalCTEProducer leftProducer = new LogicalCTEProducer<>( ctx.statementContext.getNextCTEId(), leftClone); + ctx.statementContext.setCteProducer(leftProducer.getCteId(), leftProducer); LogicalPlan rightClone = LogicalPlanDeepCopier.INSTANCE .deepCopy((LogicalPlan) join.right(), new DeepCopierContext()); LogicalCTEProducer rightProducer = new LogicalCTEProducer<>( ctx.statementContext.getNextCTEId(), rightClone); + ctx.statementContext.setCteProducer(rightProducer.getCteId(), rightProducer); Map leftCloneToLeft = new HashMap<>(); for (int i = 0; i < leftClone.getOutput().size(); i++) { leftCloneToLeft.put(leftClone.getOutput().get(i), (join.left()).getOutput().get(i)); 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 31aeb0be1bafe7..ac8ddbfeada8b6 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 @@ -53,9 +53,10 @@ import java.util.stream.Collectors; /** - * rewrite CteAnchor consumer side and producer side recursively， all CteAnchor must at top of the plan + * rewrite CteAnchor consumer side and producer side recursively， all CteAnchor + * must at top of the plan */ -@DependsRules({PullUpCteAnchor.class, CTEInline.class}) +@DependsRules({ PullUpCteAnchor.class, CTEInline.class }) public class RewriteCteChildren extends DefaultPlanRewriter implements CustomRewriter { private final List jobs; @@ -127,8 +128,8 @@ public Plan visitLogicalCTEProducer(LogicalCTEProducer cteProduc Set producerOutputs = cascadesContext.getStatementContext() .getCteIdToOutputIds().get(cteProducer.getCteId()); if (producerOutputs.size() < child.getOutput().size()) { - ImmutableList.Builder projectsBuilder - = ImmutableList.builderWithExpectedSize(producerOutputs.size()); + ImmutableList.Builder projectsBuilder = ImmutableList + .builderWithExpectedSize(producerOutputs.size()); for (Slot slot : child.getOutput()) { if (producerOutputs.contains(slot)) { projectsBuilder.add(slot); @@ -147,7 +148,10 @@ public Plan visitLogicalCTEProducer(LogicalCTEProducer cteProduc cascadesContext.addPlanProcesses(rewrittenCtx.getPlanProcesses()); cascadesContext.getStatementContext().getRewrittenCteProducer().put(cteProducer.getCteId(), child); } - return cteProducer.withChildren(child); + LogicalCTEProducer rewrittenProducer = (LogicalCTEProducer) cteProducer + .withChildren(child); + cascadesContext.getStatementContext().setCteProducer(rewrittenProducer.getCteId(), rewrittenProducer); + return rewrittenProducer; } private LogicalPlan pushPlanUnderAnchor(LogicalPlan plan) { @@ -160,13 +164,16 @@ private LogicalPlan pushPlanUnderAnchor(LogicalPlan plan) { } /* - * 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: + * 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). + * 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 consumerIds = cascadesContext.getCteIdToConsumers().get(cteId).stream() diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SplitMultiDistinctStrategy.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SplitMultiDistinctStrategy.java index 8fc150c705590a..4363e7d33ff92f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SplitMultiDistinctStrategy.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SplitMultiDistinctStrategy.java @@ -18,7 +18,6 @@ package org.apache.doris.nereids.rules.rewrite; import org.apache.doris.nereids.rules.rewrite.DistinctAggStrategySelector.DistinctSelectorContext; -import org.apache.doris.nereids.rules.rewrite.StatsDerive.DeriveContext; import org.apache.doris.nereids.trees.copier.DeepCopierContext; import org.apache.doris.nereids.trees.copier.LogicalPlanDeepCopier; import org.apache.doris.nereids.trees.expressions.Alias; @@ -61,8 +60,7 @@ public static Plan rewrite(LogicalAggregate agg, DistinctSelecto LogicalCTEProducer producer = new LogicalCTEProducer<>(ctx.statementContext.getNextCTEId(), cloneAgg.child()); ctx.cteProducerList.add(producer); - StatsDerive derive = new StatsDerive(false); - producer.accept(derive, new DeriveContext()); + ctx.statementContext.setCteProducer(producer.getCteId(), producer); Map originToProducerSlot = new HashMap<>(); for (int i = 0; i < agg.child().getOutput().size(); ++i) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/StatsDerive.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/StatsDerive.java index e3e06c23e6b9a6..cc7240a98bf367 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/StatsDerive.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/StatsDerive.java @@ -72,8 +72,8 @@ public class StatsDerive extends PlanVisitor window, Deriv @Override public Statistics visitLogicalCTEProducer(LogicalCTEProducer cteProducer, DeriveContext context) { + // Fallback registration: ensure cteId -> producer mapping is always available + // even if some upstream rewrite path misses explicit registration. + ConnectContext.get().getStatementContext().setCteProducer(cteProducer.getCteId(), cteProducer); Statistics prodStats = cteProducer.child().accept(this, context); StatisticsBuilder builder = new StatisticsBuilder(prodStats); builder.setWidthInJoinCluster(1); Statistics stats = builder.build(); cteProducer.setStatistics(stats); - ConnectContext.get().getStatementContext().setProducerStats(cteProducer.getCteId(), stats); return stats; } @Override public Statistics visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, DeriveContext context) { CTEId cteId = cteConsumer.getCteId(); - Statistics prodStats = ConnectContext.get().getStatementContext().getProducerStatsByCteId(cteId); - Preconditions.checkArgument(prodStats != null, String.format("Stats for CTE: %s not found", cteId)); + LogicalCTEProducer cteProducer = ConnectContext.get().getStatementContext() + .getCteProducerByCteId(cteId); + Preconditions.checkState(cteProducer != null, + String.format("CTE producer for CTE: %s not found", cteId)); + Statistics prodStats = cteProducer.getStats(); + if (prodStats == null || deepDerive) { + prodStats = cteProducer.accept(this, context); + } Statistics consumerStats = new Statistics(prodStats.getRowCount(), 1, new HashMap<>()); for (Slot slot : cteConsumer.getOutput()) { Slot prodSlot = cteConsumer.getProducerSlot(slot); @@ -404,8 +412,3 @@ public Statistics visitLogicalRelation(LogicalRelation relation, DeriveContext c return stats; } } - - - - -