Skip to content

Commit fd206a4

Browse files
bpinteabpintea
authored
ESQL: Add optimization to purge join on null merge key (elastic#127583) (elastic#128733)
This adds a new logical optimization rule to purge a Join in case the merge key(s) are null. The null detection is based on recognizing a tree pattern where the join sits atop a project and/or eval (possibly a few nodes deep) which contains a reference to a `null`, reference which matches the join key. It works at coordinator planning level, but it's most useful locally, after insertions of `nulls` in the plan on detecting missing fields. The Join is substituted with a projection with the same attributes as the join, atop an eval with all join's right fields aliased to null. Closes elastic#125577. (cherry picked from commit 21fe40a)
1 parent b82ebdf commit fd206a4

File tree

11 files changed

+430
-61
lines changed

11 files changed

+430
-61
lines changed

‎docs/changelog/127583.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
pr: 127583
2+
summary: Add optimization to purge join on null merge key
3+
area: ES|QL
4+
type: enhancement
5+
issues:
6+
- 125577

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/Alias.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ public String toString() {
138138

139139
@Override
140140
public String nodeString() {
141-
return child.nodeString() + " AS " + name();
141+
return child.nodeString() + " AS " + name() + "#" + id();
142142
}
143143

144144
/**

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1636,3 +1636,22 @@ max:long
16361636
3450233
16371637
8268153
16381638
;
1639+
1640+
nullifiedJoinKeyToPurgeTheJoin
1641+
required_capability: join_lookup_v12
1642+
1643+
FROM employees
1644+
| RENAME languages AS language_code
1645+
| SORT emp_no, language_code
1646+
| LIMIT 4
1647+
| EVAL language_code = TO_INTEGER(NULL)
1648+
| LOOKUP JOIN languages_lookup ON language_code
1649+
| KEEP emp_no, language_code, language_name
1650+
;
1651+
1652+
emp_no:integer | language_code:integer | language_name:keyword
1653+
10001 |null |null
1654+
10002 |null |null
1655+
10003 |null |null
1656+
10004 |null |null
1657+
;

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import org.elasticsearch.xpack.esql.optimizer.rules.logical.SubstituteSurrogatePlans;
6161
import org.elasticsearch.xpack.esql.optimizer.rules.logical.SubstituteSurrogates;
6262
import org.elasticsearch.xpack.esql.optimizer.rules.logical.TranslateMetricsAggregate;
63+
import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.PruneLeftJoinOnNullMatchingField;
6364
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
6465
import org.elasticsearch.xpack.esql.rule.ParameterizedRuleExecutor;
6566

@@ -188,7 +189,8 @@ protected static Batch<LogicalPlan> operators() {
188189
new PushDownEnrich(),
189190
new PushDownAndCombineOrderBy(),
190191
new PruneRedundantOrderBy(),
191-
new PruneRedundantSortClauses()
192+
new PruneRedundantSortClauses(),
193+
new PruneLeftJoinOnNullMatchingField()
192194
);
193195
}
194196

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/RuleUtils.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.optimizer.rules;
9+
10+
import org.elasticsearch.common.util.Maps;
11+
import org.elasticsearch.core.Tuple;
12+
import org.elasticsearch.xpack.esql.core.expression.Alias;
13+
import org.elasticsearch.xpack.esql.core.expression.Attribute;
14+
import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
15+
import org.elasticsearch.xpack.esql.core.expression.Expression;
16+
import org.elasticsearch.xpack.esql.core.expression.Literal;
17+
import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
18+
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
19+
import org.elasticsearch.xpack.esql.core.type.DataType;
20+
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
21+
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
22+
23+
import java.util.ArrayList;
24+
import java.util.List;
25+
import java.util.Map;
26+
import java.util.function.Predicate;
27+
28+
public final class RuleUtils {
29+
30+
private RuleUtils() {}
31+
32+
/**
33+
* Returns a tuple of two lists:
34+
* 1. A list of aliases to null literals for those data types in the {@param outputAttributes} that {@param shouldBeReplaced}.
35+
* 2. A list of named expressions where attributes that match the predicate are replaced with their corresponding null alias.
36+
*
37+
* @param outputAttributes The original output attributes.
38+
* @param shouldBeReplaced A predicate to determine which attributes should be replaced with null aliases.
39+
*/
40+
public static Tuple<List<Alias>, List<NamedExpression>> aliasedNulls(
41+
List<Attribute> outputAttributes,
42+
Predicate<Attribute> shouldBeReplaced
43+
) {
44+
Map<DataType, Alias> nullLiterals = Maps.newLinkedHashMapWithExpectedSize(DataType.types().size());
45+
List<NamedExpression> newProjections = new ArrayList<>(outputAttributes.size());
46+
for (Attribute attr : outputAttributes) {
47+
NamedExpression projection;
48+
if (shouldBeReplaced.test(attr)) {
49+
DataType dt = attr.dataType();
50+
Alias nullAlias = nullLiterals.get(dt);
51+
// save the first field as null (per datatype)
52+
if (nullAlias == null) {
53+
// Keep the same id so downstream query plans don't need updating
54+
// NOTE: THIS IS BRITTLE AND CAN LEAD TO BUGS.
55+
// In case some optimizer rule or so inserts a plan node that requires the field BEFORE the Eval that we're adding
56+
// on top of the EsRelation, this can trigger a field extraction in the physical optimizer phase, causing wrong
57+
// layouts due to a duplicate name id.
58+
// If someone reaches here AGAIN when debugging e.g. ClassCastExceptions NPEs from wrong layouts, we should probably
59+
// give up on this approach and instead insert EvalExecs in InsertFieldExtraction.
60+
Alias alias = new Alias(attr.source(), attr.name(), Literal.of(attr, null), attr.id());
61+
nullLiterals.put(dt, alias);
62+
projection = alias.toAttribute();
63+
}
64+
// otherwise point to it since this avoids creating field copies
65+
else {
66+
projection = new Alias(attr.source(), attr.name(), nullAlias.toAttribute(), attr.id());
67+
}
68+
} else {
69+
projection = attr;
70+
}
71+
newProjections.add(projection);
72+
}
73+
74+
return new Tuple<>(new ArrayList<>(nullLiterals.values()), newProjections);
75+
}
76+
77+
/**
78+
* Collects references to foldables from the given logical plan, returning an {@link AttributeMap} that maps
79+
* foldable aliases to their corresponding literal values.
80+
*
81+
* @param plan The logical plan to analyze.
82+
* @param ctx The optimizer context providing fold context.
83+
* @return An {@link AttributeMap} containing foldable references and their literal values.
84+
*/
85+
public static AttributeMap<Expression> foldableReferences(LogicalPlan plan, LogicalOptimizerContext ctx) {
86+
AttributeMap.Builder<Expression> collectRefsBuilder = AttributeMap.builder();
87+
88+
// collect aliases bottom-up
89+
plan.forEachExpressionUp(Alias.class, a -> {
90+
var c = a.child();
91+
boolean shouldCollect = c.foldable();
92+
// try to resolve the expression based on an existing foldables
93+
if (shouldCollect == false) {
94+
c = c.transformUp(ReferenceAttribute.class, r -> collectRefsBuilder.build().resolve(r, r));
95+
shouldCollect = c.foldable();
96+
}
97+
if (shouldCollect) {
98+
collectRefsBuilder.put(a.toAttribute(), Literal.of(ctx.foldCtx(), c));
99+
}
100+
});
101+
102+
return collectRefsBuilder.build();
103+
}
104+
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PropagateEvalFoldables.java‎

Lines changed: 4 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,11 @@
77

88
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
99

10-
import org.elasticsearch.xpack.esql.core.expression.Alias;
1110
import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
1211
import org.elasticsearch.xpack.esql.core.expression.Expression;
13-
import org.elasticsearch.xpack.esql.core.expression.Literal;
1412
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
1513
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
14+
import org.elasticsearch.xpack.esql.optimizer.rules.RuleUtils;
1615
import org.elasticsearch.xpack.esql.plan.logical.Eval;
1716
import org.elasticsearch.xpack.esql.plan.logical.Filter;
1817
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
@@ -26,24 +25,8 @@ public final class PropagateEvalFoldables extends ParameterizedRule<LogicalPlan,
2625

2726
@Override
2827
public LogicalPlan apply(LogicalPlan plan, LogicalOptimizerContext ctx) {
29-
AttributeMap.Builder<Expression> collectRefsBuilder = AttributeMap.builder();
30-
31-
java.util.function.Function<ReferenceAttribute, Expression> replaceReference = r -> collectRefsBuilder.build().resolve(r, r);
32-
33-
// collect aliases bottom-up
34-
plan.forEachExpressionUp(Alias.class, a -> {
35-
var c = a.child();
36-
boolean shouldCollect = c.foldable();
37-
// try to resolve the expression based on an existing foldables
38-
if (shouldCollect == false) {
39-
c = c.transformUp(ReferenceAttribute.class, replaceReference);
40-
shouldCollect = c.foldable();
41-
}
42-
if (shouldCollect) {
43-
collectRefsBuilder.put(a.toAttribute(), Literal.of(ctx.foldCtx(), c));
44-
}
45-
});
46-
if (collectRefsBuilder.isEmpty()) {
28+
AttributeMap<Expression> collectRefs = RuleUtils.foldableReferences(plan, ctx);
29+
if (collectRefs.isEmpty()) {
4730
return plan;
4831
}
4932

@@ -52,7 +35,7 @@ public LogicalPlan apply(LogicalPlan plan, LogicalOptimizerContext ctx) {
5235
// TODO: also allow aggregates once aggs on constants are supported.
5336
// C.f. https://github.com/elastic/elasticsearch/issues/100634
5437
if (p instanceof Filter || p instanceof Eval) {
55-
p = p.transformExpressionsOnly(ReferenceAttribute.class, replaceReference);
38+
p = p.transformExpressionsOnly(ReferenceAttribute.class, r -> collectRefs.resolve(r, r));
5639
}
5740
return p;
5841
});

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/PruneLeftJoinOnNullMatchingField.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.optimizer.rules.logical.local;
9+
10+
import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
11+
import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
12+
import org.elasticsearch.xpack.esql.core.expression.Expression;
13+
import org.elasticsearch.xpack.esql.core.expression.Expressions;
14+
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
15+
import org.elasticsearch.xpack.esql.optimizer.rules.RuleUtils;
16+
import org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules;
17+
import org.elasticsearch.xpack.esql.plan.logical.Eval;
18+
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
19+
import org.elasticsearch.xpack.esql.plan.logical.Project;
20+
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
21+
22+
import static org.elasticsearch.xpack.esql.core.expression.Expressions.isGuaranteedNull;
23+
import static org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.LEFT;
24+
25+
/**
26+
* The rule checks if the join's performed on a field which is aliased to null (in type or value); if that's the case, it prunes the join,
27+
* replacing it with an Eval - returning aliases to null for all the fields added in by the right side of the Join - plus a Project on top
28+
* of it. The rule can apply on the coordinator already, but it's more likely to be effective on the data nodes, where null aliasing is
29+
* inserted due to locally missing fields. This rule relies on that behavior -- see {@link ReplaceFieldWithConstantOrNull}.
30+
*/
31+
public class PruneLeftJoinOnNullMatchingField extends OptimizerRules.ParameterizedOptimizerRule<Join, LogicalOptimizerContext> {
32+
33+
public PruneLeftJoinOnNullMatchingField() {
34+
super(OptimizerRules.TransformDirection.DOWN);
35+
}
36+
37+
@Override
38+
protected LogicalPlan rule(Join join, LogicalOptimizerContext ctx) {
39+
LogicalPlan plan = join;
40+
if (join.config().type() == LEFT) { // other types will have different replacement logic
41+
AttributeMap<Expression> attributeMap = RuleUtils.foldableReferences(join, ctx);
42+
43+
for (var attr : AttributeSet.of(join.config().matchFields())) {
44+
var resolved = attributeMap.resolve(attr);
45+
if (resolved != null && isGuaranteedNull(resolved)) {
46+
plan = replaceJoin(join);
47+
break;
48+
}
49+
}
50+
}
51+
return plan;
52+
}
53+
54+
private static LogicalPlan replaceJoin(Join join) {
55+
var joinRightOutput = join.rightOutputFields();
56+
// can be empty when the join key is null and the rest of the right side entries pruned (such as by an agg)
57+
if (joinRightOutput.isEmpty()) {
58+
return join.left();
59+
}
60+
var aliasedNulls = RuleUtils.aliasedNulls(joinRightOutput, a -> true);
61+
var eval = new Eval(join.source(), join.left(), aliasedNulls.v1());
62+
return new Project(
63+
join.source(),
64+
eval,
65+
Join.computeOutput(join.left().output(), Expressions.asAttributes(aliasedNulls.v2()), join.config())
66+
);
67+
}
68+
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceFieldWithConstantOrNull.java‎

Lines changed: 8 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -7,17 +7,14 @@
77

88
package org.elasticsearch.xpack.esql.optimizer.rules.logical.local;
99

10-
import org.elasticsearch.common.util.Maps;
1110
import org.elasticsearch.index.IndexMode;
12-
import org.elasticsearch.xpack.esql.core.expression.Alias;
1311
import org.elasticsearch.xpack.esql.core.expression.Attribute;
1412
import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
1513
import org.elasticsearch.xpack.esql.core.expression.Expression;
1614
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
1715
import org.elasticsearch.xpack.esql.core.expression.Literal;
18-
import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
19-
import org.elasticsearch.xpack.esql.core.type.DataType;
2016
import org.elasticsearch.xpack.esql.optimizer.LocalLogicalOptimizerContext;
17+
import org.elasticsearch.xpack.esql.optimizer.rules.RuleUtils;
2118
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
2219
import org.elasticsearch.xpack.esql.plan.logical.Eval;
2320
import org.elasticsearch.xpack.esql.plan.logical.Filter;
@@ -28,7 +25,6 @@
2825
import org.elasticsearch.xpack.esql.plan.logical.TopN;
2926
import org.elasticsearch.xpack.esql.rule.ParameterizedRule;
3027

31-
import java.util.ArrayList;
3228
import java.util.HashMap;
3329
import java.util.List;
3430
import java.util.Map;
@@ -90,42 +86,18 @@ private LogicalPlan replaceWithNullOrConstant(
9086
// \_Eval[field1 = null, field3 = null]
9187
// \_EsRelation[field1, field2, field3]
9288
List<Attribute> relationOutput = relation.output();
93-
Map<DataType, Alias> nullLiterals = Maps.newLinkedHashMapWithExpectedSize(DataType.types().size());
94-
List<NamedExpression> newProjections = new ArrayList<>(relationOutput.size());
95-
for (int i = 0, size = relationOutput.size(); i < size; i++) {
96-
Attribute attr = relationOutput.get(i);
97-
NamedExpression projection;
98-
if (attr instanceof FieldAttribute f && shouldBeRetained.test(f) == false) {
99-
DataType dt = f.dataType();
100-
Alias nullAlias = nullLiterals.get(dt);
101-
// save the first field as null (per datatype)
102-
if (nullAlias == null) {
103-
// Keep the same id so downstream query plans don't need updating
104-
// NOTE: THIS IS BRITTLE AND CAN LEAD TO BUGS.
105-
// In case some optimizer rule or so inserts a plan node that requires the field BEFORE the Eval that we're adding
106-
// on top of the EsRelation, this can trigger a field extraction in the physical optimizer phase, causing wrong
107-
// layouts due to a duplicate name id.
108-
// If someone reaches here AGAIN when debugging e.g. ClassCastExceptions NPEs from wrong layouts, we should probably
109-
// give up on this approach and instead insert EvalExecs in InsertFieldExtraction.
110-
Alias alias = new Alias(f.source(), f.name(), Literal.of(f, null), f.id());
111-
nullLiterals.put(dt, alias);
112-
projection = alias.toAttribute();
113-
}
114-
// otherwise point to it since this avoids creating field copies
115-
else {
116-
projection = new Alias(f.source(), f.name(), nullAlias.toAttribute(), f.id());
117-
}
118-
} else {
119-
projection = attr;
120-
}
121-
newProjections.add(projection);
122-
}
89+
var aliasedNulls = RuleUtils.aliasedNulls(
90+
relationOutput,
91+
attr -> attr instanceof FieldAttribute f && shouldBeRetained.test(f) == false
92+
);
93+
var nullLiterals = aliasedNulls.v1();
94+
var newProjections = aliasedNulls.v2();
12395

12496
if (nullLiterals.size() == 0) {
12597
return plan;
12698
}
12799

128-
Eval eval = new Eval(plan.source(), relation, new ArrayList<>(nullLiterals.values()));
100+
Eval eval = new Eval(plan.source(), relation, nullLiterals);
129101
// This projection is redundant if there's another projection downstream (and no commands depend on the order until we hit it).
130102
return new Project(plan.source(), eval, newProjections);
131103
}

0 commit comments

Comments
 (0)