Skip to content

Commit 0e049ae

Browse files
carlosdelestcarlosdelest
committed
Separate push down limits to knn into its own rule class
1 parent 967a415 commit 0e049ae

File tree

4 files changed

+97
-42
lines changed

4 files changed

+97
-42
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEnrich;
4545
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEval;
4646
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownRegexExtract;
47+
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushLimitToKnn;
4748
import org.elasticsearch.xpack.esql.optimizer.rules.logical.RemoveStatsOverride;
4849
import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateAggExpressionWithEval;
4950
import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateNestedExpressionWithEval;
@@ -193,6 +194,7 @@ protected static Batch<LogicalPlan> operators() {
193194
new PruneFilters(),
194195
new PruneColumns(),
195196
new PruneLiteralsInOrderBy(),
197+
new PushLimitToKnn(),
196198
new PushDownAndCombineLimits(),
197199
new PushDownAndCombineFilters(),
198200
new PushDownAndCombineSample(),

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

Lines changed: 0 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -7,17 +7,13 @@
77

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

10-
import org.elasticsearch.xpack.esql.core.expression.Expression;
11-
import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
1210
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
1311
import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
1412
import org.elasticsearch.xpack.esql.plan.logical.Enrich;
1513
import org.elasticsearch.xpack.esql.plan.logical.Eval;
16-
import org.elasticsearch.xpack.esql.plan.logical.Filter;
1714
import org.elasticsearch.xpack.esql.plan.logical.Limit;
1815
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
1916
import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
20-
import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
2117
import org.elasticsearch.xpack.esql.plan.logical.Project;
2218
import org.elasticsearch.xpack.esql.plan.logical.RegexExtract;
2319
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
@@ -68,23 +64,6 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
6864
}
6965
}
7066
}
71-
if (unary instanceof Filter filter) {
72-
Expression limitAppliedExpression = limitFilterExpressions(filter.condition(), limit, ctx);
73-
if (limitAppliedExpression.equals(filter.condition()) == false) {
74-
return limit.replaceChild(filter.with(limitAppliedExpression));
75-
}
76-
}
77-
else if (unary instanceof OrderBy orderBy) {
78-
return limit.replaceChild(orderBy.transformDown(lp -> {
79-
if (lp instanceof Filter filter) {
80-
Expression limitAppliedExpression = limitFilterExpressions(filter.condition(), limit, ctx);
81-
if (limitAppliedExpression.equals(filter.condition()) == false) {
82-
return filter.with(limitAppliedExpression);
83-
}
84-
}
85-
return lp;
86-
}));
87-
}
8867
} else if (limit.child() instanceof Join join && join.config().type() == JoinTypes.LEFT) {
8968
// Left joins increase the number of rows if any join key has multiple matches from the right hand side.
9069
// Therefore, we cannot simply push down the limit - but we can add another limit before the join.
@@ -94,19 +73,6 @@ else if (unary instanceof OrderBy orderBy) {
9473
return limit;
9574
}
9675

97-
/**
98-
* Applies a limit to the filter expressions of a condition. Some filter expressions, such as KNN function,
99-
* can be optimized by applying the limit directly to them.
100-
*/
101-
private Expression limitFilterExpressions(Expression condition, Limit limit, LogicalOptimizerContext ctx) {
102-
return condition.transformDown(exp -> {
103-
if (exp instanceof Knn knn) {
104-
return knn.replaceLimit((int) limit.limit().fold(ctx.foldCtx()));
105-
}
106-
return exp;
107-
});
108-
}
109-
11076
/**
11177
* Checks the existence of another 'visible' Limit, that exists behind an operation that doesn't produce output more data than
11278
* its input (that is not a relation/source nor aggregation).

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
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.xpack.esql.core.expression.Expression;
11+
import org.elasticsearch.xpack.esql.core.util.Holder;
12+
import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
13+
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
14+
import org.elasticsearch.xpack.esql.plan.logical.Filter;
15+
import org.elasticsearch.xpack.esql.plan.logical.Limit;
16+
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
17+
18+
public class PushLimitToKnn extends OptimizerRules.ParameterizedOptimizerRule<Limit, LogicalOptimizerContext> {
19+
20+
public PushLimitToKnn() {
21+
super(OptimizerRules.TransformDirection.DOWN);
22+
}
23+
24+
@Override
25+
public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
26+
var currentLimit = (int) limit.limit().fold(ctx.foldCtx());
27+
Holder<Integer> currentLimitHolder = new Holder<>(currentLimit);
28+
return limit.transformDown(plan -> {
29+
if (plan instanceof Filter filter) {
30+
Expression limitAppliedExpression = limitFilterExpressions(filter.condition(), limit, ctx);
31+
if (limitAppliedExpression.equals(filter.condition()) == false) {
32+
return filter.with(limitAppliedExpression);
33+
}
34+
} else if (plan instanceof Limit descendantLimit) {
35+
var newLimit = (int) descendantLimit.limit().fold(ctx.foldCtx());
36+
if (newLimit < currentLimitHolder.get()) {
37+
currentLimitHolder.set(newLimit);
38+
}
39+
return descendantLimit;
40+
}
41+
42+
return plan;
43+
});
44+
}
45+
46+
/**
47+
* Applies a limit to the filter expressions of a condition. Some filter expressions, such as KNN function,
48+
* can be optimized by applying the limit directly to them.
49+
*/
50+
private Expression limitFilterExpressions(Expression condition, Limit limit, LogicalOptimizerContext ctx) {
51+
return condition.transformDown(exp -> {
52+
if (exp instanceof Knn knn) {
53+
return knn.replaceLimit((int) limit.limit().fold(ctx.foldCtx()));
54+
}
55+
return exp;
56+
});
57+
}
58+
}

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

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2115,25 +2115,54 @@ public void testKnnWithoExplicitLimitSorted() {
21152115
assertThat(knnQuery.k(), is(10));
21162116
}
21172117

2118-
public void testKnnWithExplicitLimit() {
2118+
public void testKnnWithoExplicitLimitSortedAndCommandsInBetween() {
21192119
var query = """
2120-
from test
2120+
from test metadata _score
21212121
| where knn(dense_vector, [0, 1, 2])
2122+
| sort _score desc
2123+
| keep _score
21222124
| limit 10
21232125
""";
21242126
var analyzer = makeAnalyzer("mapping-all-types.json");
21252127
var plan = plannerOptimizer.plan(query, IS_SV_STATS, analyzer);
2126-
var limitExec = as(plan, LimitExec.class);
2127-
assertThat(limitExec.limit().fold(FoldContext.small()), is(10));
2128-
var exchangeExec = as(limitExec.child(), ExchangeExec.class);
2129-
var projectExec = as(exchangeExec.child(), ProjectExec.class);
2130-
var fieldExtractExec = as(projectExec.child(), FieldExtractExec.class);
2131-
var queryExec = as(fieldExtractExec.child(), EsQueryExec.class);
2128+
var projectExec = as(plan, ProjectExec.class);
2129+
var topNExec = as(projectExec.child(), TopNExec.class);
2130+
assertThat(topNExec.limit().fold(FoldContext.small()), is(10));
2131+
var exchangeExec = as(topNExec.child(), ExchangeExec.class);
2132+
var projectExec2 = as(exchangeExec.child(), ProjectExec.class);
2133+
var queryExec = as(projectExec2.child(), EsQueryExec.class);
21322134
assertThat(queryExec.limit().fold(FoldContext.small()), is(10));
21332135
var knnQuery = as(queryExec.query(), KnnVectorQueryBuilder.class);
21342136
assertThat(knnQuery.k(), is(10));
21352137
}
21362138

2139+
public void testKnnWithExplicitLimit() {
2140+
var query = """
2141+
from test metadata _score
2142+
| where knn(dense_vector, [0, 1, 2]) or match(text, "blue")
2143+
| sort _score desc, text asc
2144+
| keep text, dense_vector
2145+
| limit 140
2146+
""";
2147+
var analyzer = makeAnalyzer("mapping-all-types.json");
2148+
var plan = plannerOptimizer.plan(query, IS_SV_STATS, analyzer);
2149+
var projectExec = as(plan, ProjectExec.class);
2150+
var topNExec = as(projectExec.child(), TopNExec.class);
2151+
assertThat(topNExec.limit().fold(FoldContext.small()), is(140));
2152+
var exchangeExec = as(topNExec.child(), ExchangeExec.class);
2153+
var projectExec2 = as(exchangeExec.child(), ProjectExec.class);
2154+
var fieldExtractExec = as(projectExec2.child(), FieldExtractExec.class);
2155+
var topNExec2 = as(fieldExtractExec.child(), TopNExec.class);
2156+
assertThat(topNExec2.limit().fold(FoldContext.small()), is(140));
2157+
var fieldExtractExec2 = as(topNExec2.child(), FieldExtractExec.class);
2158+
var queryExec = as(fieldExtractExec2.child(), EsQueryExec.class);
2159+
assertNull(queryExec.limit());
2160+
var expectedQuery = boolQuery()
2161+
.should(new KnnVectorQueryBuilder("dense_vector", new float[] { 0f, 1f, 2f }, 140, null, null, null))
2162+
.should(matchQuery("text", "blue").lenient(true));
2163+
assertEquals(expectedQuery.toString(), queryExec.query().toString());
2164+
}
2165+
21372166
public void testKnnWithExplicitLimitAndExistingTopK() {
21382167
var query = """
21392168
from test

0 commit comments

Comments
 (0)