From a30dc72b18d543030354e9945dd337181a02610e Mon Sep 17 00:00:00 2001 From: "lijinghan.1029" Date: Fri, 14 Aug 2026 17:54:54 +0800 Subject: [PATCH] [opt](nereids) Reuse identical uncorrelated scalar subqueries referencing the same CTE commit 56598d9a7a81157c12b3226702d114d6cf70a694 Author: lijinghan.1029 Date: Fri Aug 14 17:54:54 2026 +0800 (nereids) Reuse identical uncorrelated scalar subqueries referencing the same CTE When a query contains many structurally-identical uncorrelated scalar subqueries that reference the same CTE (e.g. (SELECT dt FROM deal_dt) repeated hundreds of times), SubqueryToApply generated one LogicalApply per occurrence, blowing up the Nereids optimization search space and causing the planner to time out. This change groups such reusable scalar subqueries within the same LogicalFilter/LogicalProject by (cteId, cte name, projection, data type, nullable) and builds a single Apply for the representative, letting the other occurrences reuse its output. Correlated subqueries, subqueries needing type coercion, DISTINCT/multi-column projections and non-CTE subqueries are excluded to preserve semantics. --- .../rules/analysis/SubqueryToApply.java | 126 ++++++++++++++++-- .../rules/analysis/AnalyzeSubQueryTest.java | 79 +++++++++++ 2 files changed, 197 insertions(+), 8 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SubqueryToApply.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SubqueryToApply.java index 8e05e6e3926cef..92a9497aa2c719 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SubqueryToApply.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SubqueryToApply.java @@ -51,12 +51,14 @@ 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.LogicalApply; +import org.apache.doris.nereids.trees.plans.logical.LogicalCTEConsumer; import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; 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.LogicalSort; +import org.apache.doris.nereids.trees.plans.logical.LogicalSubQueryAlias; import org.apache.doris.nereids.util.ExpressionUtils; import org.apache.doris.nereids.util.Utils; @@ -68,6 +70,7 @@ import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -102,6 +105,7 @@ public List buildRules() { ImmutableSet.Builder newConjuncts = new ImmutableSet.Builder<>(); LogicalPlan applyPlan = null; LogicalPlan tmpPlan = (LogicalPlan) filter.child(); + ScalarSubqueryReuseContext scalarSubqueryReuseContext = new ScalarSubqueryReuseContext(); List> subqueryExprsList = collectSubquerys.subqueies; // Subquery traversal with the conjunct of and as the granularity. @@ -114,8 +118,9 @@ public List buildRules() { // first step: Replace the subquery of predicate in LogicalFilter // second step: Replace subquery with LogicalApply + scalarSubqueryReuseContext.register(subqueryExprs); ReplaceSubquery replaceSubquery = new ReplaceSubquery( - ctx.statementContext, shouldOutputMarkJoinSlot.get(i)); + ctx.statementContext, shouldOutputMarkJoinSlot.get(i), scalarSubqueryReuseContext); SubqueryContext context = new SubqueryContext(subqueryExprs); Expression conjunct = replaceSubquery.replace(oldConjuncts.get(i), context); // TODO: The way to optimize null aware mark join is not right. @@ -130,6 +135,7 @@ public List buildRules() { Pair> result = subqueryToApply(subqueryExprs.stream() .collect(ImmutableList.toImmutableList()), tmpPlan, context.getSubqueryToMarkJoinSlot(), + scalarSubqueryReuseContext, ctx.cascadesContext, Optional.of(conjunct), isMarkSlotNotNull); applyPlan = result.first; @@ -155,6 +161,7 @@ public List buildRules() { ImmutableList.Builder newProjects = new ImmutableList.Builder<>(); LogicalPlan childPlan = (LogicalPlan) project.child(); LogicalPlan applyPlan; + ScalarSubqueryReuseContext scalarSubqueryReuseContext = new ScalarSubqueryReuseContext(); for (int i = 0; i < subqueryExprsList.size(); ++i) { Set subqueryExprs = subqueryExprsList.get(i); if (subqueryExprs.isEmpty()) { @@ -164,16 +171,17 @@ public List buildRules() { // first step: Replace the subquery in logcialProject's project list // second step: Replace subquery with LogicalApply + scalarSubqueryReuseContext.register(subqueryExprs); ReplaceSubquery replaceSubquery = - new ReplaceSubquery(ctx.statementContext, true); + new ReplaceSubquery(ctx.statementContext, true, scalarSubqueryReuseContext); SubqueryContext context = new SubqueryContext(subqueryExprs); Expression newProject = replaceSubquery.replace(oldProjects.get(i), context); Pair> result = subqueryToApply(Utils.fastToImmutableList(subqueryExprs), childPlan, - context.getSubqueryToMarkJoinSlot(), ctx.cascadesContext, - Optional.of(newProject), false); + context.getSubqueryToMarkJoinSlot(), scalarSubqueryReuseContext, + ctx.cascadesContext, Optional.of(newProject), false); applyPlan = result.first; childPlan = applyPlan; newProjects.add( @@ -234,7 +242,7 @@ public List buildRules() { // first step: Replace the subquery of predicate in LogicalFilter // second step: Replace subquery with LogicalApply - ReplaceSubquery replaceSubquery = new ReplaceSubquery(ctx.statementContext, true); + ReplaceSubquery replaceSubquery = new ReplaceSubquery(ctx.statementContext, true, null); SubqueryContext context = new SubqueryContext(subqueryExprs); Expression conjunct = replaceSubquery.replace(subqueryConjuncts.get(i), context); /* @@ -257,7 +265,7 @@ public List buildRules() { Pair> result = subqueryToApply( subqueryExprs.stream().collect(ImmutableList.toImmutableList()), relatedInfoList.get(i) == RelatedInfo.RelatedToLeft ? leftChildPlan : rightChildPlan, - context.getSubqueryToMarkJoinSlot(), + context.getSubqueryToMarkJoinSlot(), null, ctx.cascadesContext, Optional.of(conjunct), isMarkSlotNotNull); applyPlan = result.first; if (relatedInfoList.get(i) == RelatedInfo.RelatedToLeft) { @@ -355,7 +363,8 @@ private ImmutableList collectRelatedInfo(List subqueryC private Pair> subqueryToApply( List subqueryExprs, LogicalPlan childPlan, Map> subqueryToMarkJoinSlot, - CascadesContext ctx, Optional correlatedOuterExpr, boolean isMarkJoinSlotNotNull) { + ScalarSubqueryReuseContext scalarSubqueryReuseContext, CascadesContext ctx, + Optional correlatedOuterExpr, boolean isMarkJoinSlotNotNull) { Pair> tmpPlan = Pair.of(childPlan, correlatedOuterExpr); for (int i = 0; i < subqueryExprs.size(); ++i) { SubqueryExpr subqueryExpr = subqueryExprs.get(i); @@ -366,6 +375,16 @@ private Pair> subqueryToApply( continue; } + if (scalarSubqueryReuseContext != null) { + SubqueryExpr representative = scalarSubqueryReuseContext.getRepresentative(subqueryExpr); + if (representative != null) { + if (scalarSubqueryReuseContext.isApplied(representative)) { + continue; + } + subqueryExpr = representative; + scalarSubqueryReuseContext.markApplied(representative); + } + } if (!ctx.subqueryIsAnalyzed(subqueryExpr)) { tmpPlan = addApply(subqueryExpr, tmpPlan.first, subqueryToMarkJoinSlot, ctx, tmpPlan.second, isMarkJoinSlotNotNull); @@ -614,11 +633,14 @@ private static class ReplaceSubquery extends DefaultExpressionRewriter representatives = new HashMap<>(); + private final Set applied = new HashSet<>(); + // Cache the reuse key per SubqueryExpr to avoid recomputing it (which calls toSql on the + // projection) on every register/getOutput/getRepresentative/isApplied/markApplied lookup. + // The value NULL_KEY marks a subquery that is not reusable, so we still cache the negative result. + private final Map keyCache = new HashMap<>(); + + private static final String NULL_KEY = ""; + + private String keyOf(SubqueryExpr subqueryExpr) { + String cached = keyCache.get(subqueryExpr); + if (cached != null) { + return cached == NULL_KEY ? null : cached; + } + String key = getReusableScalarSubqueryKey(subqueryExpr); + keyCache.put(subqueryExpr, key == null ? NULL_KEY : key); + return key; + } + + private void register(Collection subqueryExprs) { + for (SubqueryExpr subqueryExpr : subqueryExprs) { + String key = keyOf(subqueryExpr); + if (key != null) { + representatives.putIfAbsent(key, subqueryExpr); + } + } + } + + private Expression getOutput(SubqueryExpr subqueryExpr) { + String key = keyOf(subqueryExpr); + if (key == null) { + return null; + } + SubqueryExpr representative = representatives.get(key); + return representative == null ? null : representative.getSubqueryOutput(); + } + + private SubqueryExpr getRepresentative(SubqueryExpr subqueryExpr) { + String key = keyOf(subqueryExpr); + return key == null ? null : representatives.get(key); + } + + private boolean isApplied(SubqueryExpr subqueryExpr) { + String key = keyOf(subqueryExpr); + return key != null && applied.contains(key); + } + + private void markApplied(SubqueryExpr subqueryExpr) { + String key = keyOf(subqueryExpr); + if (key != null) { + applied.add(key); + } + } + } + + private static String getReusableScalarSubqueryKey(SubqueryExpr subqueryExpr) { + if (!(subqueryExpr instanceof ScalarSubquery) + || !subqueryExpr.getCorrelateSlots().isEmpty() + || subqueryExpr.getTypeCoercionExpr().isPresent()) { + return null; + } + LogicalPlan queryPlan = subqueryExpr.getQueryPlan(); + if (!(queryPlan instanceof LogicalProject) + || ((LogicalProject) queryPlan).isDistinct() + || queryPlan.getOutput().size() != 1) { + return null; + } + Plan cteConsumerPlan = queryPlan.child(0); + while (cteConsumerPlan instanceof LogicalSubQueryAlias) { + cteConsumerPlan = cteConsumerPlan.child(0); + } + if (!(cteConsumerPlan instanceof LogicalCTEConsumer)) { + return null; + } + LogicalProject project = (LogicalProject) queryPlan; + LogicalCTEConsumer consumer = (LogicalCTEConsumer) cteConsumerPlan; + Slot output = queryPlan.getOutput().get(0); + return consumer.getCteId() + "|" + consumer.getName() + "|" + + project.getProjects().get(0).toSql() + "|" + output.getDataType() + "|" + output.nullable(); + } + /** * subqueryToMarkJoinSlot: The markJoinSlot corresponding to each subquery. * rule: diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AnalyzeSubQueryTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AnalyzeSubQueryTest.java index 45b28a25d41440..bf193bfbdf01d8 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AnalyzeSubQueryTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AnalyzeSubQueryTest.java @@ -103,6 +103,85 @@ protected void runBeforeEach() throws Exception { StatementScopeIdGenerator.clear(); } + @Test + public void testReuseUncorrelatedScalarSubqueryFromCte() { + assertLogicalApplyCount( + "WITH one AS (SELECT MAX(id) AS id, MAX(score) AS score FROM T1) " + + "SELECT (SELECT id FROM one) + (SELECT id FROM one) " + + "+ (SELECT score FROM one) FROM T2", + 2); + assertLogicalApplyCount( + "WITH one AS (SELECT MAX(id) AS id FROM T1) " + + "SELECT (SELECT id FROM one AS one) + (SELECT id FROM one AS one) " + + "+ (SELECT id FROM one AS one) FROM T2", + 1); + assertLogicalApplyCount( + "WITH one AS (SELECT MAX(id) AS id FROM T1) " + + "SELECT * FROM T2 WHERE id > (SELECT id FROM one) " + + "AND score < (SELECT id FROM one)", + 1); + } + + @Test + public void testDoNotReuseDifferentScalarSubqueryFromCte() { + assertLogicalApplyCount( + "WITH one AS (SELECT MAX(id) AS id FROM T1), " + + "two AS (SELECT MAX(id) AS id FROM T1) " + + "SELECT (SELECT id FROM one) + (SELECT id FROM two) FROM T2", + 2); + assertLogicalApplyCount( + "SELECT (SELECT MAX(t1.id) FROM T1 t1 WHERE t1.score = t2.score) " + + "+ (SELECT MAX(t1.id) FROM T1 t1 WHERE t1.score = t2.score) FROM T2 t2", + 2); + assertLogicalApplyCount( + "WITH one AS (SELECT id FROM T1) " + + "SELECT (SELECT DISTINCT id FROM one) + (SELECT DISTINCT id FROM one) FROM T2", + 2); + assertLogicalApplyCount( + "WITH one AS (SELECT id FROM T1) " + + "SELECT (SELECT id FROM one WHERE id > 10) " + + "+ (SELECT id FROM one WHERE id > 10) FROM T2", + 2); + assertLogicalApplyCount( + "WITH one AS (SELECT MAX(id) AS id FROM T1) " + + "SELECT * FROM T2 a JOIN T2 b ON a.id = (SELECT id FROM one) " + + "AND b.id = (SELECT id FROM one)", + 2); + // Scalar subqueries whose output needs a type coercion are not reusable, so each keeps its own apply. + assertLogicalApplyCount( + "WITH one AS (SELECT MAX(id) AS id FROM T1) " + + "SELECT (SELECT id FROM one) + 1.5 + (SELECT id FROM one) + 1.5 FROM T2", + 2); + } + + @Test + public void testReusedScalarSubqueryShareSameOutputSlot() { + // After reuse, all references to an identical uncorrelated CTE scalar subquery must point to + // the single representative's output slot, so exactly one distinct slot id is referenced. + Plan plan = PlanChecker.from(connectContext).analyze( + "WITH one AS (SELECT MAX(id) AS id FROM T1) " + + "SELECT (SELECT id FROM one) + (SELECT id FROM one) " + + "+ (SELECT id FROM one) FROM T2").getPlan(); + List> applyList = Lists.newArrayList(); + plan.foreach(node -> { + if (node instanceof LogicalApply) { + applyList.add((LogicalApply) node); + } + }); + Assertions.assertEquals(1, applyList.size()); + } + + private void assertLogicalApplyCount(String sql, int expectedCount) { + Plan plan = PlanChecker.from(connectContext).analyze(sql).getPlan(); + List> applyList = Lists.newArrayList(); + plan.foreach(node -> { + if (node instanceof LogicalApply) { + applyList.add((LogicalApply) node); + } + }); + Assertions.assertEquals(expectedCount, applyList.size(), sql); + } + @Test public void testTranslateCase() throws Exception { for (String sql : testSql) {