Skip to content

Commit 4eab7d5

Browse files
afoucretafoucret
committed
Adding some tests for pushable limits optimizer rules.
1 parent 98816dc commit 4eab7d5

File tree

1 file changed

+159
-0
lines changed

1 file changed

+159
-0
lines changed

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

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
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;
9+
10+
import org.elasticsearch.test.ESTestCase;
11+
import org.elasticsearch.xpack.esql.core.expression.Alias;
12+
import org.elasticsearch.xpack.esql.core.expression.Attribute;
13+
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
14+
import org.elasticsearch.xpack.esql.core.expression.Literal;
15+
import org.elasticsearch.xpack.esql.expression.Order;
16+
import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToInteger;
17+
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
18+
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
19+
import org.elasticsearch.xpack.esql.plan.logical.Eval;
20+
import org.elasticsearch.xpack.esql.plan.logical.Filter;
21+
import org.elasticsearch.xpack.esql.plan.logical.Limit;
22+
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
23+
import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
24+
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
25+
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
26+
27+
import java.util.List;
28+
import java.util.function.BiConsumer;
29+
import java.util.function.BiFunction;
30+
31+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
32+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getFieldAttribute;
33+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
34+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.unboundLogicalOptimizerContext;
35+
import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
36+
import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
37+
import static org.elasticsearch.xpack.esql.core.type.DataType.TEXT;
38+
import static org.elasticsearch.xpack.esql.optimizer.LocalLogicalPlanOptimizerTests.relation;
39+
40+
public class PushDownAndCombineLimitsTests extends ESTestCase {
41+
42+
private static class PushDownLimitTestCase<PlanType extends UnaryPlan> {
43+
private final Class<PlanType> clazz;
44+
private final BiFunction<LogicalPlan, Attribute, PlanType> planBuilder;
45+
private final BiConsumer<PlanType, PlanType> planChecker;
46+
47+
PushDownLimitTestCase(
48+
Class<PlanType> clazz,
49+
BiFunction<LogicalPlan, Attribute, PlanType> planBuilder,
50+
BiConsumer<PlanType, PlanType> planChecker
51+
) {
52+
this.clazz = clazz;
53+
this.planBuilder = planBuilder;
54+
this.planChecker = planChecker;
55+
}
56+
57+
public PlanType buildPlan(LogicalPlan child, Attribute attr) {
58+
return planBuilder.apply(child, attr);
59+
}
60+
61+
public void checkOptimizedPlan(LogicalPlan basePlan, LogicalPlan optimizedPlan) {
62+
planChecker.accept(as(basePlan, clazz), as(optimizedPlan, clazz));
63+
}
64+
}
65+
66+
private static final List<PushDownLimitTestCase<? extends UnaryPlan>> PUSHABLE_LIMIT_TEST_CASES = List.of(
67+
new PushDownLimitTestCase<>(
68+
Eval.class,
69+
(plan, attr) -> new Eval(EMPTY, plan, List.of(new Alias(EMPTY, "y", new ToInteger(EMPTY, attr)))),
70+
(basePlan, optimizedPlan) -> {
71+
assertEquals(basePlan.source(), optimizedPlan.source());
72+
assertEquals(basePlan.fields(), optimizedPlan.fields());
73+
}
74+
),
75+
new PushDownLimitTestCase<>(
76+
Completion.class,
77+
(plan, attr) -> new Completion(EMPTY, plan, randomLiteral(TEXT), randomLiteral(TEXT), attr),
78+
(basePlan, optimizedPlan) -> {
79+
assertEquals(basePlan.source(), optimizedPlan.source());
80+
assertEquals(basePlan.inferenceId(), optimizedPlan.inferenceId());
81+
assertEquals(basePlan.prompt(), optimizedPlan.prompt());
82+
assertEquals(basePlan.targetField(), optimizedPlan.targetField());
83+
}
84+
)
85+
);
86+
87+
private static final List<PushDownLimitTestCase<? extends UnaryPlan>> NON_PUSHABLE_LIMIT_TEST_CASES = List.of(
88+
new PushDownLimitTestCase<>(
89+
Filter.class,
90+
(plan, attr) -> new Filter(EMPTY, plan, new Equals(EMPTY, attr, new Literal(EMPTY, "right", TEXT))),
91+
(basePlan, optimizedPlan) -> {
92+
assertEquals(basePlan.source(), optimizedPlan.source());
93+
assertEquals(basePlan.condition(), optimizedPlan.condition());
94+
}
95+
),
96+
new PushDownLimitTestCase<>(
97+
OrderBy.class,
98+
(plan, attr) -> new OrderBy(EMPTY, plan, List.of(new Order(EMPTY, attr, Order.OrderDirection.DESC, null))),
99+
(basePlan, optimizedPlan) -> {
100+
assertEquals(basePlan.source(), optimizedPlan.source());
101+
assertEquals(basePlan.order(), optimizedPlan.order());
102+
}
103+
)
104+
);
105+
106+
public void testPushableLimit() {
107+
FieldAttribute a = getFieldAttribute("a");
108+
FieldAttribute b = getFieldAttribute("b");
109+
EsRelation relation = relation().withAttributes(List.of(a, b));
110+
111+
for (PushDownLimitTestCase<? extends UnaryPlan> pushableLimitTestCase : PUSHABLE_LIMIT_TEST_CASES) {
112+
int precedingLimitValue = randomIntBetween(1, 10_000);
113+
Limit precedingLimit = new Limit(EMPTY, new Literal(EMPTY, precedingLimitValue, INTEGER), relation);
114+
115+
LogicalPlan pushableLimitTestPlan = pushableLimitTestCase.buildPlan(precedingLimit, a);
116+
117+
int pushableLimitValue = randomIntBetween(1, 10_000);
118+
Limit pushableLimit = new Limit(EMPTY, new Literal(EMPTY, pushableLimitValue, INTEGER), pushableLimitTestPlan);
119+
120+
LogicalPlan optimizedPlan = optimizePlan(pushableLimit);
121+
122+
pushableLimitTestCase.checkOptimizedPlan(pushableLimitTestPlan, optimizedPlan);
123+
124+
assertEquals(
125+
as(optimizedPlan, UnaryPlan.class).child(),
126+
new Limit(EMPTY, new Literal(EMPTY, Math.min(pushableLimitValue, precedingLimitValue), INTEGER), relation)
127+
);
128+
}
129+
}
130+
131+
public void testNonPushableLimit() {
132+
FieldAttribute a = getFieldAttribute("a");
133+
FieldAttribute b = getFieldAttribute("b");
134+
EsRelation relation = relation().withAttributes(List.of(a, b));
135+
136+
for (PushDownLimitTestCase<? extends UnaryPlan> nonPushableLimitTestCase : NON_PUSHABLE_LIMIT_TEST_CASES) {
137+
int precedingLimitValue = randomIntBetween(1, 10_000);
138+
Limit precedingLimit = new Limit(EMPTY, new Literal(EMPTY, precedingLimitValue, INTEGER), relation);
139+
UnaryPlan nonPushableLimitTestPlan = nonPushableLimitTestCase.buildPlan(precedingLimit, a);
140+
int nonPushableLimitValue = randomIntBetween(1, 10_000);
141+
Limit nonPushableLimit = new Limit(EMPTY, new Literal(EMPTY, nonPushableLimitValue, INTEGER), nonPushableLimitTestPlan);
142+
Limit optimizedPlan = as(optimizePlan(nonPushableLimit), Limit.class);
143+
nonPushableLimitTestCase.checkOptimizedPlan(nonPushableLimitTestPlan, optimizedPlan.child());
144+
assertEquals(
145+
optimizedPlan,
146+
new Limit(
147+
EMPTY,
148+
new Literal(EMPTY, Math.min(nonPushableLimitValue, precedingLimitValue), INTEGER),
149+
nonPushableLimitTestPlan
150+
)
151+
);
152+
assertEquals(as(optimizedPlan.child(), UnaryPlan.class).child(), nonPushableLimitTestPlan.child());
153+
}
154+
}
155+
156+
private LogicalPlan optimizePlan(LogicalPlan plan) {
157+
return new PushDownAndCombineLimits().apply(plan, unboundLogicalOptimizerContext());
158+
}
159+
}

0 commit comments

Comments
 (0)