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) {