Skip to content

[opt](nereids) Reuse identical uncorrelated scalar subqueries referencing the same CTE #66781

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LLLLLL302311 wants to merge 1 commit into
base: master
Choose a base branch
from
+197 −8
Open

[opt](nereids) Reuse identical uncorrelated scalar subqueries referencing the same CTE #66781

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 118 additions & 8 deletions fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SubqueryToApply.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -102,6 +105,7 @@ public List<Rule> buildRules() {
ImmutableSet.Builder<Expression> newConjuncts = new ImmutableSet.Builder<>();
LogicalPlan applyPlan = null;
LogicalPlan tmpPlan = (LogicalPlan) filter.child();
ScalarSubqueryReuseContext scalarSubqueryReuseContext = new ScalarSubqueryReuseContext();

List<Set<SubqueryExpr>> subqueryExprsList = collectSubquerys.subqueies;
// Subquery traversal with the conjunct of and as the granularity.
Expand All @@ -114,8 +118,9 @@ public List<Rule> 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.
Expand All @@ -130,6 +135,7 @@ public List<Rule> buildRules() {
Pair<LogicalPlan, Optional<Expression>> result = subqueryToApply(subqueryExprs.stream()
.collect(ImmutableList.toImmutableList()), tmpPlan,
context.getSubqueryToMarkJoinSlot(),
scalarSubqueryReuseContext,
ctx.cascadesContext,
Optional.of(conjunct), isMarkSlotNotNull);
applyPlan = result.first;
Expand All @@ -155,6 +161,7 @@ public List<Rule> buildRules() {
ImmutableList.Builder<NamedExpression> newProjects = new ImmutableList.Builder<>();
LogicalPlan childPlan = (LogicalPlan) project.child();
LogicalPlan applyPlan;
ScalarSubqueryReuseContext scalarSubqueryReuseContext = new ScalarSubqueryReuseContext();
for (int i = 0; i < subqueryExprsList.size(); ++i) {
Set<SubqueryExpr> subqueryExprs = subqueryExprsList.get(i);
if (subqueryExprs.isEmpty()) {
Expand All @@ -164,16 +171,17 @@ public List<Rule> 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<LogicalPlan, Optional<Expression>> 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(
Expand Down Expand Up @@ -234,7 +242,7 @@ public List<Rule> 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);
/*
Expand All @@ -257,7 +265,7 @@ public List<Rule> buildRules() {
Pair<LogicalPlan, Optional<Expression>> 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) {
Expand Down Expand Up @@ -355,7 +363,8 @@ private ImmutableList<RelatedInfo> collectRelatedInfo(List<Expression> subqueryC
private Pair<LogicalPlan, Optional<Expression>> subqueryToApply(
List<SubqueryExpr> subqueryExprs, LogicalPlan childPlan,
Map<SubqueryExpr, Optional<MarkJoinSlotReference>> subqueryToMarkJoinSlot,
CascadesContext ctx, Optional<Expression> correlatedOuterExpr, boolean isMarkJoinSlotNotNull) {
ScalarSubqueryReuseContext scalarSubqueryReuseContext, CascadesContext ctx,
Optional<Expression> correlatedOuterExpr, boolean isMarkJoinSlotNotNull) {
Pair<LogicalPlan, Optional<Expression>> tmpPlan = Pair.of(childPlan, correlatedOuterExpr);
for (int i = 0; i < subqueryExprs.size(); ++i) {
SubqueryExpr subqueryExpr = subqueryExprs.get(i);
Expand All @@ -366,6 +375,16 @@ private Pair<LogicalPlan, Optional<Expression>> 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);
Expand Down Expand Up @@ -614,11 +633,14 @@ private static class ReplaceSubquery extends DefaultExpressionRewriter<SubqueryC
private boolean isMarkJoin;

private final boolean shouldOutputMarkJoinSlot;
private final ScalarSubqueryReuseContext scalarSubqueryReuseContext;

public ReplaceSubquery(StatementContext statementContext,
boolean shouldOutputMarkJoinSlot) {
boolean shouldOutputMarkJoinSlot,
ScalarSubqueryReuseContext scalarSubqueryReuseContext) {
this.statementContext = Objects.requireNonNull(statementContext, "statementContext can't be null");
this.shouldOutputMarkJoinSlot = shouldOutputMarkJoinSlot;
this.scalarSubqueryReuseContext = scalarSubqueryReuseContext;
}

public Expression replace(Expression expression, SubqueryContext subqueryContext) {
Expand Down Expand Up @@ -669,6 +691,12 @@ public Expression visitInSubquery(InSubquery in, SubqueryContext context) {

@Override
public Expression visitScalarSubquery(ScalarSubquery scalar, SubqueryContext context) {
if (scalarSubqueryReuseContext != null) {
Expression cachedOutput = scalarSubqueryReuseContext.getOutput(scalar);
if (cachedOutput != null) {
return cachedOutput;
}
}
return scalar.getSubqueryOutput();
}

Expand All @@ -689,6 +717,88 @@ public Expression visitCompoundPredicate(CompoundPredicate compound, SubqueryCon
}
}

private static class ScalarSubqueryReuseContext {
private final Map<String, SubqueryExpr> representatives = new HashMap<>();
private final Set<String> 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<SubqueryExpr, String> 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<SubqueryExpr> 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:
Expand Down
79 changes: 79 additions & 0 deletions fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AnalyzeSubQueryTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<LogicalApply<?, ?>> 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<LogicalApply<?, ?>> 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) {
Expand Down
Loading