Skip to content

Commit edfb33e

Browse files
dnhatndnhatn
committed
local
1 parent c99b9a9 commit edfb33e

File tree

5 files changed

+70
-38
lines changed

5 files changed

+70
-38
lines changed

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@
3030
* This class is part of the planner. Data node level logical optimizations. At this point we have access to
3131
* {@link org.elasticsearch.xpack.esql.stats.SearchStats} which provides access to metadata about the index.
3232
*
33-
* <p>NB: This class also reapplies all the rules from {@link LogicalPlanOptimizer#operators()} and {@link LogicalPlanOptimizer#cleanup()}
33+
* <p>NB: This class also reapplies all the rules from {@link LogicalPlanOptimizer#operators(boolean)}
34+
* and {@link LogicalPlanOptimizer#cleanup()}
3435
*/
3536
public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<LogicalPlan, LocalLogicalOptimizerContext> {
3637

@@ -58,8 +59,8 @@ protected List<Batch<LogicalPlan>> batches() {
5859

5960
@SuppressWarnings("unchecked")
6061
private static Batch<LogicalPlan> localOperators() {
61-
var operators = operators();
62-
var rules = operators().rules();
62+
var operators = operators(true);
63+
var rules = operators.rules();
6364
List<Rule<?, LogicalPlan>> newRules = new ArrayList<>(rules.length);
6465

6566
// apply updates to existing rules that have different applicability locally

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -82,22 +82,22 @@
8282
* <li>The {@link LogicalPlanOptimizer#substitutions()} phase rewrites things to expand out shorthand in the syntax. For example,
8383
* a nested expression embedded in a stats gets replaced with an eval followed by a stats, followed by another eval. This phase
8484
* also applies surrogates, such as replacing an average with a sum divided by a count.</li>
85-
* <li>{@link LogicalPlanOptimizer#operators()} (NB: The word "operator" is extremely overloaded and referrers to many different
85+
* <li>{@link LogicalPlanOptimizer#operators(boolean)} (NB: The word "operator" is extremely overloaded and referrers to many different
8686
* things.) transform the tree in various different ways. This includes folding (i.e. computing constant expressions at parse
8787
* time), combining expressions, dropping redundant clauses, and some normalization such as putting literals on the right whenever
8888
* possible. These rules are run in a loop until none of the rules make any changes to the plan (there is also a safety shut off
8989
* after many iterations, although hitting that is considered a bug)</li>
9090
* <li>{@link LogicalPlanOptimizer#cleanup()} Which can replace sorts+limit with a TopN</li>
9191
* </ul>
9292
*
93-
* <p>Note that the {@link LogicalPlanOptimizer#operators()} and {@link LogicalPlanOptimizer#cleanup()} steps are reapplied at the
93+
* <p>Note that the {@link LogicalPlanOptimizer#operators(boolean)} and {@link LogicalPlanOptimizer#cleanup()} steps are reapplied at the
9494
* {@link LocalLogicalPlanOptimizer} layer.</p>
9595
*/
9696
public class LogicalPlanOptimizer extends ParameterizedRuleExecutor<LogicalPlan, LogicalOptimizerContext> {
9797

9898
private static final List<RuleExecutor.Batch<LogicalPlan>> RULES = List.of(
9999
substitutions(),
100-
operators(),
100+
operators(false),
101101
new Batch<>("Skip Compute", new SkipQueryOnLimitZero()),
102102
cleanup(),
103103
new Batch<>("Set as Optimized", Limiter.ONCE, new SetAsOptimized())
@@ -160,10 +160,10 @@ protected static Batch<LogicalPlan> substitutions() {
160160
);
161161
}
162162

163-
protected static Batch<LogicalPlan> operators() {
163+
protected static Batch<LogicalPlan> operators(boolean local) {
164164
return new Batch<>(
165165
"Operator Optimization",
166-
new CombineProjections(),
166+
new CombineProjections(local),
167167
new CombineEvals(),
168168
new PruneEmptyPlans(),
169169
new PropagateEmptyRelation(),

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

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

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

10+
import org.elasticsearch.common.util.Maps;
1011
import org.elasticsearch.core.Tuple;
1112
import org.elasticsearch.xpack.esql.core.expression.Alias;
1213
import org.elasticsearch.xpack.esql.core.expression.Attribute;
@@ -15,11 +16,13 @@
1516
import org.elasticsearch.xpack.esql.core.expression.Literal;
1617
import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
1718
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
19+
import org.elasticsearch.xpack.esql.core.type.DataType;
1820
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
1921
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
2022

2123
import java.util.ArrayList;
2224
import java.util.List;
25+
import java.util.Map;
2326
import java.util.function.Predicate;
2427

2528
public final class RuleUtils {
@@ -38,20 +41,37 @@ public static Tuple<List<Alias>, List<NamedExpression>> aliasedNulls(
3841
List<Attribute> outputAttributes,
3942
Predicate<Attribute> shouldBeReplaced
4043
) {
44+
Map<DataType, Alias> nullLiterals = Maps.newLinkedHashMapWithExpectedSize(DataType.types().size());
4145
List<NamedExpression> newProjections = new ArrayList<>(outputAttributes.size());
42-
List<Alias> nullLiterals = new ArrayList<>(outputAttributes.size());
4346
for (Attribute attr : outputAttributes) {
4447
NamedExpression projection;
4548
if (shouldBeReplaced.test(attr)) {
46-
Alias alias = new Alias(attr.source(), attr.name(), Literal.of(attr, null), attr.id());
47-
nullLiterals.add(alias);
48-
projection = alias.toAttribute();
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+
}
4968
} else {
5069
projection = attr;
5170
}
5271
newProjections.add(projection);
5372
}
54-
return new Tuple<>(nullLiterals, newProjections);
73+
74+
return new Tuple<>(new ArrayList<>(nullLiterals.values()), newProjections);
5575
}
5676

5777
/**

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

Lines changed: 23 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,12 @@
2727
import java.util.List;
2828

2929
public final class CombineProjections extends OptimizerRules.OptimizerRule<UnaryPlan> {
30+
// don't drop groupings from a local plan, as the layout has already been agreed upon
31+
private final boolean local;
3032

31-
public CombineProjections() {
33+
public CombineProjections(boolean local) {
3234
super(OptimizerRules.TransformDirection.UP);
35+
this.local = local;
3336
}
3437

3538
@Override
@@ -54,33 +57,32 @@ protected LogicalPlan rule(UnaryPlan plan) {
5457
// project can be fully removed
5558
if (newAggs != null) {
5659
var newGroups = replacePrunedAliasesUsedInGroupBy(a.groupings(), aggs, newAggs);
57-
plan = a.with(newGroups, newAggs);
60+
if (local == false || newGroups.size() == a.groupings().size()) {
61+
plan = a.with(newGroups, newAggs);
62+
}
5863
}
5964
}
6065
return plan;
6166
}
6267

6368
// Agg with underlying Project (group by on sub-queries)
64-
if (plan instanceof Aggregate a) {
65-
if (child instanceof Project p) {
66-
var groupings = a.groupings();
67-
List<NamedExpression> groupingAttrs = new ArrayList<>(a.groupings().size());
68-
for (Expression grouping : groupings) {
69-
if (grouping instanceof Attribute attribute) {
70-
groupingAttrs.add(attribute);
71-
} else if (grouping instanceof Alias as && as.child() instanceof GroupingFunction.NonEvaluatableGroupingFunction) {
72-
groupingAttrs.add(as);
73-
} else {
74-
// After applying ReplaceAggregateNestedExpressionWithEval,
75-
// evaluatable groupings can only contain attributes.
76-
throw new EsqlIllegalArgumentException("Expected an Attribute, got {}", grouping);
77-
}
69+
if (plan instanceof Aggregate a && child instanceof Project p) {
70+
var groupings = a.groupings();
71+
List<NamedExpression> groupingAttrs = new ArrayList<>(a.groupings().size());
72+
for (Expression grouping : groupings) {
73+
if (grouping instanceof Attribute attribute) {
74+
groupingAttrs.add(attribute);
75+
} else if (grouping instanceof Alias as && as.child() instanceof GroupingFunction.NonEvaluatableGroupingFunction) {
76+
groupingAttrs.add(as);
77+
} else {
78+
// After applying ReplaceAggregateNestedExpressionWithEval,
79+
// evaluatable groupings can only contain attributes.
80+
throw new EsqlIllegalArgumentException("Expected an Attribute, got {}", grouping);
7881
}
79-
plan = a.with(
80-
p.child(),
81-
combineUpperGroupingsAndLowerProjections(groupingAttrs, p.projections()),
82-
combineProjections(a.aggregates(), p.projections())
83-
);
82+
}
83+
List<Expression> newGroupings = combineUpperGroupingsAndLowerProjections(groupingAttrs, p.projections());
84+
if (local == false || newGroupings.size() == a.groupings().size()) {
85+
plan = a.with(p.child(), newGroupings, combineProjections(a.aggregates(), p.projections()));
8486
}
8587
}
8688

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
4646
import org.elasticsearch.xpack.esql.core.expression.Literal;
4747
import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
48+
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
4849
import org.elasticsearch.xpack.esql.core.tree.Source;
4950
import org.elasticsearch.xpack.esql.core.type.DataType;
5051
import org.elasticsearch.xpack.esql.core.type.EsField;
@@ -132,7 +133,6 @@
132133
import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS;
133134
import static org.elasticsearch.xpack.esql.plan.physical.EsStatsQueryExec.StatsType;
134135
import static org.hamcrest.Matchers.contains;
135-
import static org.hamcrest.Matchers.containsInAnyOrder;
136136
import static org.hamcrest.Matchers.equalTo;
137137
import static org.hamcrest.Matchers.hasSize;
138138
import static org.hamcrest.Matchers.instanceOf;
@@ -925,12 +925,21 @@ public void testMissingFieldsDoNotGetExtracted() {
925925
"salary"
926926
)
927927
);
928+
// emp_no
929+
assertThat(projections.get(1), instanceOf(ReferenceAttribute.class));
930+
// first_name
931+
assertThat(projections.get(2), instanceOf(ReferenceAttribute.class));
932+
933+
// last_name --> first_name
934+
var nullAlias = Alias.unwrap(projections.get(8));
935+
assertThat(Expressions.name(nullAlias), is("first_name"));
936+
// salary --> emp_no
937+
nullAlias = Alias.unwrap(projections.get(10));
938+
assertThat(Expressions.name(nullAlias), is("emp_no"));
939+
// check field extraction is skipped and that evaled fields are not extracted anymore
928940
var field = as(project.child(), FieldExtractExec.class);
929941
var fields = field.attributesToExtract();
930942
assertThat(Expressions.names(fields), contains("_meta_field", "gender", "hire_date", "job", "job.raw", "languages", "long_noidx"));
931-
var eval = as(field.child(), EvalExec.class);
932-
List<String> nullFields = Expressions.names(eval.fields());
933-
assertThat(nullFields, containsInAnyOrder("first_name", "last_name", "emp_no", "salary"));
934943
}
935944

936945
/*

0 commit comments

Comments
 (0)