|
|
@@ -63,6 +63,8 @@ import org.elasticsearch.xpack.esql.expression.function.fulltext.Kql;
|
|
|
import org.elasticsearch.xpack.esql.expression.function.fulltext.Match;
|
|
|
import org.elasticsearch.xpack.esql.expression.function.fulltext.MatchOperator;
|
|
|
import org.elasticsearch.xpack.esql.expression.function.fulltext.QueryString;
|
|
|
+import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
|
|
|
+import org.elasticsearch.xpack.esql.expression.predicate.logical.And;
|
|
|
import org.elasticsearch.xpack.esql.expression.predicate.logical.Or;
|
|
|
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.GreaterThan;
|
|
|
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.GreaterThanOrEqual;
|
|
|
@@ -108,6 +110,7 @@ import org.junit.Before;
|
|
|
|
|
|
import java.io.IOException;
|
|
|
import java.util.ArrayList;
|
|
|
+import java.util.Arrays;
|
|
|
import java.util.Collection;
|
|
|
import java.util.List;
|
|
|
import java.util.Locale;
|
|
|
@@ -1371,7 +1374,7 @@ public class LocalPhysicalPlanOptimizerTests extends MapperServiceTestCase {
|
|
|
|
|
|
public void testKnnOptionsPushDown() {
|
|
|
assumeTrue("dense_vector capability not available", EsqlCapabilities.Cap.DENSE_VECTOR_FIELD_TYPE.isEnabled());
|
|
|
- assumeTrue("knn capability not available", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
|
|
|
+ assumeTrue("knn capability not available", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
|
|
|
String query = """
|
|
|
from test
|
|
|
@@ -1836,6 +1839,308 @@ public class LocalPhysicalPlanOptimizerTests extends MapperServiceTestCase {
|
|
|
aggExec.forEachDown(EsQueryExec.class, esQueryExec -> { assertNull(esQueryExec.query()); });
|
|
|
}
|
|
|
|
|
|
+ public void testKnnPrefilters() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where knn(dense_vector, [0, 1, 2], 10) and integer > 10
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var field = as(project.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(field.child(), EsQueryExec.class);
|
|
|
+ QueryBuilder expectedFilterQueryBuilder = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(rangeQuery("integer").gt(10)),
|
|
|
+ "integer",
|
|
|
+ new Source(2, 45, "integer > 10")
|
|
|
+ );
|
|
|
+ KnnVectorQueryBuilder expectedKnnQueryBuilder = new KnnVectorQueryBuilder(
|
|
|
+ "dense_vector",
|
|
|
+ new float[] { 0, 1, 2 },
|
|
|
+ 10,
|
|
|
+ null,
|
|
|
+ null,
|
|
|
+ null
|
|
|
+ ).addFilterQuery(expectedFilterQueryBuilder);
|
|
|
+ var expectedQuery = boolQuery().must(expectedKnnQueryBuilder).must(expectedFilterQueryBuilder);
|
|
|
+ assertEquals(expectedQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testKnnPrefiltersWithMultipleFilters() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where knn(dense_vector, [0, 1, 2], 10)
|
|
|
+ | where integer > 10
|
|
|
+ | where keyword == "test"
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var field = as(project.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(field.child(), EsQueryExec.class);
|
|
|
+ var integerFilter = wrapWithSingleQuery(query, unscore(rangeQuery("integer").gt(10)), "integer", new Source(3, 8, "integer > 10"));
|
|
|
+ var keywordFilter = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(termQuery("keyword", "test")),
|
|
|
+ "keyword",
|
|
|
+ new Source(4, 8, "keyword == \"test\"")
|
|
|
+ );
|
|
|
+ QueryBuilder expectedFilterQueryBuilder = boolQuery().must(integerFilter).must(keywordFilter);
|
|
|
+ KnnVectorQueryBuilder expectedKnnQueryBuilder = new KnnVectorQueryBuilder(
|
|
|
+ "dense_vector",
|
|
|
+ new float[] { 0, 1, 2 },
|
|
|
+ 10,
|
|
|
+ null,
|
|
|
+ null,
|
|
|
+ null
|
|
|
+ ).addFilterQuery(expectedFilterQueryBuilder);
|
|
|
+ var expectedQuery = boolQuery().must(expectedKnnQueryBuilder).must(integerFilter).must(keywordFilter);
|
|
|
+ assertEquals(expectedQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testPushDownConjunctionsToKnnPrefilter() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where knn(dense_vector, [0, 1, 2], 10) and integer > 10
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var field = as(project.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(field.child(), EsQueryExec.class);
|
|
|
+
|
|
|
+ // The filter condition should be pushed down to both the KNN query and the main query
|
|
|
+ QueryBuilder expectedFilterQueryBuilder = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(rangeQuery("integer").gt(10)),
|
|
|
+ "integer",
|
|
|
+ new Source(2, 45, "integer > 10")
|
|
|
+ );
|
|
|
+
|
|
|
+ KnnVectorQueryBuilder expectedKnnQueryBuilder = new KnnVectorQueryBuilder(
|
|
|
+ "dense_vector",
|
|
|
+ new float[] { 0, 1, 2 },
|
|
|
+ 10,
|
|
|
+ null,
|
|
|
+ null,
|
|
|
+ null
|
|
|
+ ).addFilterQuery(expectedFilterQueryBuilder);
|
|
|
+
|
|
|
+ var expectedQuery = boolQuery().must(expectedKnnQueryBuilder).must(expectedFilterQueryBuilder);
|
|
|
+
|
|
|
+ assertEquals(expectedQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testPushDownNegatedConjunctionsToKnnPrefilter() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where knn(dense_vector, [0, 1, 2], 10) and NOT integer > 10
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var field = as(project.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(field.child(), EsQueryExec.class);
|
|
|
+
|
|
|
+ // The filter condition should be pushed down to both the KNN query and the main query
|
|
|
+ QueryBuilder expectedFilterQueryBuilder = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(boolQuery().mustNot(unscore(rangeQuery("integer").gt(10)))),
|
|
|
+ "integer",
|
|
|
+ new Source(2, 45, "NOT integer > 10")
|
|
|
+ );
|
|
|
+
|
|
|
+ KnnVectorQueryBuilder expectedKnnQueryBuilder = new KnnVectorQueryBuilder(
|
|
|
+ "dense_vector",
|
|
|
+ new float[] { 0, 1, 2 },
|
|
|
+ 10,
|
|
|
+ null,
|
|
|
+ null,
|
|
|
+ null
|
|
|
+ ).addFilterQuery(expectedFilterQueryBuilder);
|
|
|
+
|
|
|
+ var expectedQuery = boolQuery().must(expectedKnnQueryBuilder).must(expectedFilterQueryBuilder);
|
|
|
+
|
|
|
+ assertEquals(expectedQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testNotPushDownDisjunctionsToKnnPrefilter() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where knn(dense_vector, [0, 1, 2], 10) or integer > 10
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var field = as(project.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(field.child(), EsQueryExec.class);
|
|
|
+
|
|
|
+ // The disjunction should not be pushed down to the KNN query
|
|
|
+ KnnVectorQueryBuilder knnQueryBuilder = new KnnVectorQueryBuilder("dense_vector", new float[] { 0, 1, 2 }, 10, null, null, null);
|
|
|
+ QueryBuilder rangeQueryBuilder = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(rangeQuery("integer").gt(10)),
|
|
|
+ "integer",
|
|
|
+ new Source(2, 44, "integer > 10")
|
|
|
+ );
|
|
|
+
|
|
|
+ var expectedQuery = boolQuery().should(knnQueryBuilder).should(rangeQueryBuilder);
|
|
|
+
|
|
|
+ assertEquals(expectedQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testNotPushDownKnnWithNonPushablePrefilters() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where ((knn(dense_vector, [0, 1, 2], 10) AND integer > 10) and ((keyword == "test") or length(text) > 10))
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var field = as(project.child(), FieldExtractExec.class);
|
|
|
+ var secondLimit = as(field.child(), LimitExec.class);
|
|
|
+ var filter = as(secondLimit.child(), FilterExec.class);
|
|
|
+ var and = as(filter.condition(), And.class);
|
|
|
+ var knn = as(and.left(), Knn.class);
|
|
|
+ assertEquals("(keyword == \"test\") or length(text) > 10", knn.filterExpressions().get(0).toString());
|
|
|
+ assertEquals("integer > 10", knn.filterExpressions().get(1).toString());
|
|
|
+
|
|
|
+ var fieldExtract = as(filter.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(fieldExtract.child(), EsQueryExec.class);
|
|
|
+
|
|
|
+ // The query should only contain the pushable condition
|
|
|
+ QueryBuilder integerGtQuery = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(rangeQuery("integer").gt(10)),
|
|
|
+ "integer",
|
|
|
+ new Source(2, 47, "integer > 10")
|
|
|
+ );
|
|
|
+
|
|
|
+ assertEquals(integerGtQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testPushDownComplexNegationsToKnnPrefilter() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where ((knn(dense_vector, [0, 1, 2], 10) or NOT integer > 10)
|
|
|
+ and NOT ((keyword == "test") or knn(dense_vector, [4, 5, 6], 10)))
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var fieldExtract = as(project.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(fieldExtract.child(), EsQueryExec.class);
|
|
|
+
|
|
|
+ QueryBuilder notKeywordQuery = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(boolQuery().mustNot(unscore(termQuery("keyword", "test")))),
|
|
|
+ "keyword",
|
|
|
+ new Source(3, 12, "keyword == \"test\"")
|
|
|
+ );
|
|
|
+ QueryBuilder notKeywordFilter = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(boolQuery().mustNot(unscore(termQuery("keyword", "test")))),
|
|
|
+ "keyword",
|
|
|
+ new Source(3, 6, "NOT ((keyword == \"test\") or knn(dense_vector, [4, 5, 6], 10))")
|
|
|
+ );
|
|
|
+
|
|
|
+ QueryBuilder notIntegerGt10 = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(boolQuery().mustNot(unscore(rangeQuery("integer").gt(10)))),
|
|
|
+ "integer",
|
|
|
+ new Source(2, 46, "NOT integer > 10")
|
|
|
+ );
|
|
|
+
|
|
|
+ KnnVectorQueryBuilder firstKnn = new KnnVectorQueryBuilder("dense_vector", new float[] { 0, 1, 2 }, 10, null, null, null);
|
|
|
+ KnnVectorQueryBuilder secondKnn = new KnnVectorQueryBuilder("dense_vector", new float[] { 4, 5, 6 }, 10, null, null, null);
|
|
|
+
|
|
|
+ firstKnn.addFilterQuery(notKeywordFilter);
|
|
|
+ secondKnn.addFilterQuery(notIntegerGt10);
|
|
|
+
|
|
|
+ // Build the main boolean query structure
|
|
|
+ BoolQueryBuilder expectedQuery = boolQuery().must(notKeywordQuery) // NOT (keyword == "test")
|
|
|
+ .must(unscore(boolQuery().mustNot(secondKnn)))
|
|
|
+ .must(boolQuery().should(firstKnn).should(notIntegerGt10));
|
|
|
+
|
|
|
+ assertEquals(expectedQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testMultipleKnnQueriesInPrefilters() {
|
|
|
+ assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
|
|
|
+
|
|
|
+ String query = """
|
|
|
+ from test
|
|
|
+ | where ((knn(dense_vector, [0, 1, 2], 10) or integer > 10) and ((keyword == "test") or knn(dense_vector, [4, 5, 6], 10)))
|
|
|
+ """;
|
|
|
+ var plan = plannerOptimizer.plan(query, IS_SV_STATS, makeAnalyzer("mapping-all-types.json"));
|
|
|
+
|
|
|
+ var limit = as(plan, LimitExec.class);
|
|
|
+ var exchange = as(limit.child(), ExchangeExec.class);
|
|
|
+ var project = as(exchange.child(), ProjectExec.class);
|
|
|
+ var field = as(project.child(), FieldExtractExec.class);
|
|
|
+ var queryExec = as(field.child(), EsQueryExec.class);
|
|
|
+
|
|
|
+ KnnVectorQueryBuilder firstKnnQuery = new KnnVectorQueryBuilder("dense_vector", new float[] { 0, 1, 2 }, 10, null, null, null);
|
|
|
+ // Integer range query (right side of first OR)
|
|
|
+ QueryBuilder integerRangeQuery = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(rangeQuery("integer").gt(10)),
|
|
|
+ "integer",
|
|
|
+ new Source(2, 46, "integer > 10")
|
|
|
+ );
|
|
|
+
|
|
|
+ // Second KNN query (right side of second OR)
|
|
|
+ KnnVectorQueryBuilder secondKnnQuery = new KnnVectorQueryBuilder("dense_vector", new float[] { 4, 5, 6 }, 10, null, null, null);
|
|
|
+
|
|
|
+ // Keyword term query (left side of second OR)
|
|
|
+ QueryBuilder keywordQuery = wrapWithSingleQuery(
|
|
|
+ query,
|
|
|
+ unscore(termQuery("keyword", "test")),
|
|
|
+ "keyword",
|
|
|
+ new Source(2, 66, "keyword == \"test\"")
|
|
|
+ );
|
|
|
+
|
|
|
+ // First OR (knn1 OR integer > 10)
|
|
|
+ var firstOr = boolQuery().should(firstKnnQuery).should(integerRangeQuery);
|
|
|
+ // Second OR (keyword == "test" OR knn2)
|
|
|
+ var secondOr = boolQuery().should(keywordQuery).should(secondKnnQuery);
|
|
|
+ firstKnnQuery.addFilterQuery(keywordQuery);
|
|
|
+ secondKnnQuery.addFilterQuery(integerRangeQuery);
|
|
|
+
|
|
|
+ // Top-level AND combining both ORs
|
|
|
+ var expectedQuery = boolQuery().must(firstOr).must(secondOr);
|
|
|
+ assertEquals(expectedQuery.toString(), queryExec.query().toString());
|
|
|
+ }
|
|
|
+
|
|
|
public void testParallelizeTimeSeriesPlan() {
|
|
|
assumeTrue("requires snapshot builds", Build.current().isSnapshot());
|
|
|
var query = "TS k8s | STATS max(rate(network.total_bytes_in)) BY bucket(@timestamp, 1h)";
|
|
|
@@ -2234,4 +2539,33 @@ public class LocalPhysicalPlanOptimizerTests extends MapperServiceTestCase {
|
|
|
return "qstr(\"" + fieldName() + ": " + queryString() + "\")";
|
|
|
}
|
|
|
}
|
|
|
+
|
|
|
+ private class KnnFunctionTestCase extends FullTextFunctionTestCase {
|
|
|
+
|
|
|
+ final int k;
|
|
|
+
|
|
|
+ KnnFunctionTestCase() {
|
|
|
+ super(Knn.class, "dense_vector", randomVector());
|
|
|
+ k = randomIntBetween(1, 10);
|
|
|
+ }
|
|
|
+
|
|
|
+ private static Object randomVector() {
|
|
|
+ int numDims = randomIntBetween(10, 20);
|
|
|
+ float[] vector = new float[numDims];
|
|
|
+ for (int i = 0; i < numDims; i++) {
|
|
|
+ vector[i] = randomFloat();
|
|
|
+ }
|
|
|
+ return vector;
|
|
|
+ }
|
|
|
+
|
|
|
+ @Override
|
|
|
+ public QueryBuilder queryBuilder() {
|
|
|
+ return new KnnVectorQueryBuilder(fieldName(), (float[]) queryString(), k, null, null, null);
|
|
|
+ }
|
|
|
+
|
|
|
+ @Override
|
|
|
+ public String esqlQuery() {
|
|
|
+ return "knn(" + fieldName() + ", " + Arrays.toString(((float[]) queryString())) + ", " + k + ")";
|
|
|
+ }
|
|
|
+ }
|
|
|
}
|
Carlos Delgado