Skip to content

Commit 0270314

Browse files
julian-elasticjulian-elastic
committed
Address code review feedback part 2
1 parent fcbe4bc commit 0270314

File tree

2 files changed

+218
-5
lines changed

2 files changed

+218
-5
lines changed

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

Lines changed: 108 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.elasticsearch.index.IndexMode;
1414
import org.elasticsearch.test.ESTestCase;
1515
import org.elasticsearch.xpack.esql.EsqlTestUtils;
16+
import org.elasticsearch.xpack.esql.VerificationException;
1617
import org.elasticsearch.xpack.esql.analysis.Analyzer;
1718
import org.elasticsearch.xpack.esql.analysis.AnalyzerContext;
1819
import org.elasticsearch.xpack.esql.core.expression.Alias;
@@ -43,6 +44,7 @@
4344
import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
4445
import org.elasticsearch.xpack.esql.index.EsIndex;
4546
import org.elasticsearch.xpack.esql.index.IndexResolution;
47+
import org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules;
4648
import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.InferIsNotNull;
4749
import org.elasticsearch.xpack.esql.parser.EsqlParser;
4850
import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
@@ -59,10 +61,12 @@
5961
import org.elasticsearch.xpack.esql.plan.logical.local.EmptyLocalSupplier;
6062
import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
6163
import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
64+
import org.elasticsearch.xpack.esql.rule.RuleExecutor;
6265
import org.elasticsearch.xpack.esql.stats.SearchStats;
6366
import org.hamcrest.Matchers;
6467
import org.junit.BeforeClass;
6568

69+
import java.util.ArrayList;
6670
import java.util.List;
6771
import java.util.Locale;
6872
import java.util.Map;
@@ -88,6 +92,9 @@
8892
import static org.elasticsearch.xpack.esql.EsqlTestUtils.unboundLogicalOptimizerContext;
8993
import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
9094
import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
95+
import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
96+
import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.DOWN;
97+
import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.UP;
9198
import static org.hamcrest.Matchers.contains;
9299
import static org.hamcrest.Matchers.containsString;
93100
import static org.hamcrest.Matchers.equalTo;
@@ -780,7 +787,7 @@ public void testGroupingByMissingFields() {
780787
as(eval.child(), EsRelation.class);
781788
}
782789

783-
public void testPlanSanityCheck() throws Exception {
790+
public void testVerifierOnMissingReferences() throws Exception {
784791
var plan = localPlan("""
785792
from test
786793
| stats a = min(salary) by emp_no
@@ -806,6 +813,106 @@ public void testPlanSanityCheck() throws Exception {
806813
assertThat(e.getMessage(), containsString(" optimized incorrectly due to missing references [salary"));
807814
}
808815

816+
private LocalLogicalPlanOptimizer getCustomRulesLocalLogicalPlanOptimizer(List<RuleExecutor.Batch<LogicalPlan>> batches) {
817+
LocalLogicalOptimizerContext context = new LocalLogicalOptimizerContext(
818+
EsqlTestUtils.TEST_CFG,
819+
FoldContext.small(),
820+
TEST_SEARCH_STATS
821+
);
822+
LocalLogicalPlanOptimizer customOptimizer = new LocalLogicalPlanOptimizer(context) {
823+
@Override
824+
protected List<Batch<LogicalPlan>> batches() {
825+
return batches;
826+
}
827+
};
828+
return customOptimizer;
829+
}
830+
831+
public void testVerifierOnAdditionalAttributeAdded() throws Exception {
832+
var plan = localPlan("""
833+
from test
834+
| stats a = min(salary) by emp_no
835+
""");
836+
837+
var limit = as(plan, Limit.class);
838+
var aggregate = as(limit.child(), Aggregate.class);
839+
var min = as(Alias.unwrap(aggregate.aggregates().get(0)), Min.class);
840+
var salary = as(min.field(), NamedExpression.class);
841+
assertThat(salary.name(), is("salary"));
842+
843+
// use a custom rule that adds another output attribute
844+
var customRuleBatch = new RuleExecutor.Batch<>(
845+
"CustomRuleBatch",
846+
RuleExecutor.Limiter.ONCE,
847+
new OptimizerRules.ParameterizedOptimizerRule<Aggregate, LocalLogicalOptimizerContext>(UP) {
848+
static Integer appliedCount = 0;
849+
850+
@Override
851+
protected LogicalPlan rule(Aggregate plan, LocalLogicalOptimizerContext context) {
852+
// This rule adds a missing attribute to the plan output
853+
// We only want to apply it once, so we use a static counter
854+
if (appliedCount == 0) {
855+
appliedCount++;
856+
Literal additionalLiteral = new Literal(Source.EMPTY, "additional literal", INTEGER);
857+
return new Eval(plan.source(), plan, List.of(new Alias(Source.EMPTY, "additionalAttribute", additionalLiteral)));
858+
}
859+
return plan;
860+
}
861+
862+
}
863+
);
864+
LocalLogicalPlanOptimizer customRulesLocalLogicalPlanOptimizer = getCustomRulesLocalLogicalPlanOptimizer(List.of(customRuleBatch));
865+
Exception e = expectThrows(VerificationException.class, () -> customRulesLocalLogicalPlanOptimizer.localOptimize(plan));
866+
assertThat(e.getMessage(), containsString("Output has changed from"));
867+
assertThat(e.getMessage(), containsString("additionalAttribute"));
868+
}
869+
870+
public void testVerifierOnAttributeDatatypeChanged(){
871+
var plan = localPlan("""
872+
from test
873+
| stats a = min(salary) by emp_no
874+
""");
875+
876+
var limit = as(plan, Limit.class);
877+
var aggregate = as(limit.child(), Aggregate.class);
878+
var min = as(Alias.unwrap(aggregate.aggregates().get(0)), Min.class);
879+
var salary = as(min.field(), NamedExpression.class);
880+
assertThat(salary.name(), is("salary"));
881+
882+
// use a custom rule that changes the datatype of an output attribute
883+
var customRuleBatch = new RuleExecutor.Batch<>(
884+
"CustomRuleBatch",
885+
RuleExecutor.Limiter.ONCE,
886+
new OptimizerRules.ParameterizedOptimizerRule<LogicalPlan, LocalLogicalOptimizerContext>(DOWN) {
887+
static Integer appliedCount = 0;
888+
889+
@Override
890+
protected LogicalPlan rule(LogicalPlan plan, LocalLogicalOptimizerContext context) {
891+
// We only want to apply it once, so we use a static counter
892+
if (appliedCount == 0) {
893+
appliedCount++;
894+
Limit limit = as(plan, Limit.class);
895+
Limit newLimit = new Limit(plan.source(), limit.limit(), limit.child()) {
896+
@Override
897+
public List<Attribute> output() {
898+
List<Attribute> oldOutput = super.output();
899+
List<Attribute> newOutput = new ArrayList<>(oldOutput);
900+
newOutput.set(0, oldOutput.get(0).withDataType(DataType.DATETIME));
901+
return newOutput;
902+
}
903+
};
904+
return newLimit;
905+
}
906+
return plan;
907+
}
908+
909+
}
910+
);
911+
LocalLogicalPlanOptimizer customRulesLocalLogicalPlanOptimizer = getCustomRulesLocalLogicalPlanOptimizer(List.of(customRuleBatch));
912+
Exception e = expectThrows(VerificationException.class, () -> customRulesLocalLogicalPlanOptimizer.localOptimize(plan));
913+
assertThat(e.getMessage(), containsString("Output has changed from"));
914+
}
915+
809916
private IsNotNull isNotNull(Expression field) {
810917
return new IsNotNull(EMPTY, field);
811918
}

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

Lines changed: 110 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,14 @@
3434
import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
3535
import org.elasticsearch.xpack.esql.EsqlTestUtils;
3636
import org.elasticsearch.xpack.esql.EsqlTestUtils.TestSearchStats;
37+
import org.elasticsearch.xpack.esql.VerificationException;
3738
import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
3839
import org.elasticsearch.xpack.esql.analysis.Analyzer;
3940
import org.elasticsearch.xpack.esql.analysis.AnalyzerContext;
4041
import org.elasticsearch.xpack.esql.analysis.EnrichResolution;
4142
import org.elasticsearch.xpack.esql.analysis.Verifier;
4243
import org.elasticsearch.xpack.esql.core.expression.Alias;
44+
import org.elasticsearch.xpack.esql.core.expression.Attribute;
4345
import org.elasticsearch.xpack.esql.core.expression.Expression;
4446
import org.elasticsearch.xpack.esql.core.expression.Expressions;
4547
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
@@ -98,6 +100,7 @@
98100
import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
99101
import org.elasticsearch.xpack.esql.querydsl.query.SingleValueQuery;
100102
import org.elasticsearch.xpack.esql.rule.Rule;
103+
import org.elasticsearch.xpack.esql.rule.RuleExecutor;
101104
import org.elasticsearch.xpack.esql.session.Configuration;
102105
import org.elasticsearch.xpack.esql.stats.SearchContextStats;
103106
import org.elasticsearch.xpack.esql.stats.SearchStats;
@@ -2084,20 +2087,123 @@ public void testVerifierOnMissingReferences() throws Exception {
20842087
// We want to verify that the localOptimize detects the missing attribute.
20852088
// However, it also throws an error in one of the rules before we get to the verifier.
20862089
// So we use an implementation of LocalPhysicalPlanOptimizer that does not have any rules.
2090+
LocalPhysicalPlanOptimizer optimizerWithNoRules = getCustomRulesLocalPhysicalPlanOptimizer(List.of());
2091+
Exception e = expectThrows(IllegalStateException.class, () -> optimizerWithNoRules.localOptimize(topNExec));
2092+
assertThat(e.getMessage(), containsString(" optimized incorrectly due to missing references [missing attr"));
2093+
}
2094+
2095+
private LocalPhysicalPlanOptimizer getCustomRulesLocalPhysicalPlanOptimizer(List<RuleExecutor.Batch<PhysicalPlan>> batches) {
20872096
LocalPhysicalOptimizerContext context = new LocalPhysicalOptimizerContext(
20882097
new EsqlFlags(true),
20892098
config,
20902099
FoldContext.small(),
20912100
SearchStats.EMPTY
20922101
);
2093-
LocalPhysicalPlanOptimizer optimizerWithNoopExecute = new LocalPhysicalPlanOptimizer(context) {
2102+
LocalPhysicalPlanOptimizer localPhysicalPlanOptimizer = new LocalPhysicalPlanOptimizer(context) {
20942103
@Override
20952104
protected List<Batch<PhysicalPlan>> batches() {
2096-
return List.of();
2105+
return batches;
20972106
}
20982107
};
2099-
Exception e = expectThrows(IllegalStateException.class, () -> optimizerWithNoopExecute.localOptimize(topNExec));
2100-
assertThat(e.getMessage(), containsString(" optimized incorrectly due to missing references [missing attr"));
2108+
return localPhysicalPlanOptimizer;
2109+
}
2110+
2111+
public void testVerifierOnAdditionalAttributeAdded() throws Exception {
2112+
2113+
PhysicalPlan plan = plannerOptimizer.plan("""
2114+
from test
2115+
| stats a = min(salary) by emp_no
2116+
""");
2117+
2118+
var limit = as(plan, LimitExec.class);
2119+
var aggregate = as(limit.child(), AggregateExec.class);
2120+
var min = as(Alias.unwrap(aggregate.aggregates().get(0)), Min.class);
2121+
var salary = as(min.field(), NamedExpression.class);
2122+
assertThat(salary.name(), is("salary"));
2123+
2124+
// use a custom rule that adds another output attribute
2125+
var customRuleBatch = new RuleExecutor.Batch<>(
2126+
"CustomRuleBatch",
2127+
RuleExecutor.Limiter.ONCE,
2128+
new PhysicalOptimizerRules.ParameterizedOptimizerRule<PhysicalPlan, LocalPhysicalOptimizerContext>() {
2129+
static Integer appliedCount = 0;
2130+
2131+
@Override
2132+
public PhysicalPlan rule(PhysicalPlan plan, LocalPhysicalOptimizerContext context) {
2133+
// This rule adds a missing attribute to the plan output
2134+
// We only want to apply it once, so we use a static counter
2135+
if (appliedCount == 0) {
2136+
appliedCount++;
2137+
Literal additionalLiteral = new Literal(Source.EMPTY, "additional literal", INTEGER);
2138+
return new EvalExec(
2139+
plan.source(),
2140+
plan,
2141+
List.of(new Alias(Source.EMPTY, "additionalAttribute", additionalLiteral))
2142+
);
2143+
}
2144+
return plan;
2145+
}
2146+
}
2147+
);
2148+
LocalPhysicalPlanOptimizer customRulesLocalPhysicalPlanOptimizer = getCustomRulesLocalPhysicalPlanOptimizer(
2149+
List.of(customRuleBatch)
2150+
);
2151+
Exception e = expectThrows(VerificationException.class, () -> customRulesLocalPhysicalPlanOptimizer.localOptimize(plan));
2152+
assertThat(e.getMessage(), containsString("Output has changed from"));
2153+
assertThat(e.getMessage(), containsString("additionalAttribute"));
2154+
}
2155+
2156+
public void testVerifierOnAttributeDatatypeChanged() throws Exception {
2157+
2158+
PhysicalPlan plan = plannerOptimizer.plan("""
2159+
from test
2160+
| stats a = min(salary) by emp_no
2161+
""");
2162+
2163+
var limit = as(plan, LimitExec.class);
2164+
var aggregate = as(limit.child(), AggregateExec.class);
2165+
var min = as(Alias.unwrap(aggregate.aggregates().get(0)), Min.class);
2166+
var salary = as(min.field(), NamedExpression.class);
2167+
assertThat(salary.name(), is("salary"));
2168+
2169+
// use a custom rule that changes the datatype of an output attribute
2170+
var customRuleBatch = new RuleExecutor.Batch<>(
2171+
"CustomRuleBatch",
2172+
RuleExecutor.Limiter.ONCE,
2173+
new PhysicalOptimizerRules.ParameterizedOptimizerRule<PhysicalPlan, LocalPhysicalOptimizerContext>() {
2174+
static Integer appliedCount = 0;
2175+
2176+
@Override
2177+
public PhysicalPlan rule(PhysicalPlan plan, LocalPhysicalOptimizerContext context) {
2178+
// We only want to apply it once, so we use a static counter
2179+
if (appliedCount == 0) {
2180+
appliedCount++;
2181+
LimitExec limit = as(plan, LimitExec.class);
2182+
LimitExec newLimit = new LimitExec(
2183+
plan.source(),
2184+
limit.child(),
2185+
new Literal(Source.EMPTY, 1000, INTEGER),
2186+
randomEstimatedRowSize()
2187+
) {
2188+
@Override
2189+
public List<Attribute> output() {
2190+
List<Attribute> oldOutput = super.output();
2191+
List<Attribute> newOutput = new ArrayList<>(oldOutput);
2192+
newOutput.set(0, oldOutput.get(0).withDataType(DataType.DATETIME));
2193+
return newOutput;
2194+
}
2195+
};
2196+
return newLimit;
2197+
}
2198+
return plan;
2199+
}
2200+
}
2201+
);
2202+
LocalPhysicalPlanOptimizer customRulesLocalPhysicalPlanOptimizer = getCustomRulesLocalPhysicalPlanOptimizer(
2203+
List.of(customRuleBatch)
2204+
);
2205+
Exception e = expectThrows(VerificationException.class, () -> customRulesLocalPhysicalPlanOptimizer.localOptimize(plan));
2206+
assertThat(e.getMessage(), containsString("Output has changed from"));
21012207
}
21022208

21032209
private boolean isMultiTypeEsField(Expression e) {

0 commit comments

Comments
 (0)