ESQL - KNN function uses prefilters when pushed down to Lucene

x-pack/plugin/esql/qa/testFixtures/src/main/resources/knn-function.csv-spec

@@ -3,7 +3,7 @@
 # top-n query at the shard level 
 
 knnSearch
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 // tag::knn-function[]
 from colors metadata _score 
@@ -30,7 +30,7 @@ chartreuse | [127.0, 255.0, 0.0]
 ;
 
 knnSearchWithSimilarityOption
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score 
 | where knn(rgb_vector, [255,192,203], 140, {"similarity": 40})
@@ -46,14 +46,13 @@ wheat      | [245.0, 222.0, 179.0]
 ;
 
 knnHybridSearch
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score 
-| where match(color, "blue") or knn(rgb_vector, [65,105,225], 140)
+| where match(color, "blue") or knn(rgb_vector, [65,105,225], 10)
 | where primary == true
 | sort _score desc, color asc
 | keep color, rgb_vector
-| limit 10
 ;
 
 color:text | rgb_vector:dense_vector
@@ -68,21 +67,45 @@ red        | [255.0, 0.0, 0.0]
 yellow     | [255.0, 255.0, 0.0]
 ;
 
-knnWithMultipleFunctions
-required_capability: knn_function_v2
+knnWithPrefilter
+required_capability: knn_function_v3
 
 from colors metadata _score
-| where knn(rgb_vector, [128,128,0], 140) and match(color, "olive") 
+| where knn(rgb_vector, [128,128,0], 10) and (match(color, "olive") or match(color, "green")) 
 | sort _score desc, color asc
 | keep color, rgb_vector
 ;
 
 color:text       | rgb_vector:dense_vector
 olive            | [128.0, 128.0, 0.0]
+green            | [0.0, 128.0, 0.0]
+;
+
+knnWithNegatedPrefilter
+required_capability: knn_function_v3
+
+from colors metadata _score
+| where knn(rgb_vector, [128,128,0], 10) and not (match(color, "olive") or match(color, "chocolate")) 
+| sort _score desc, color asc
+| keep color, rgb_vector
+| LIMIT 10
+;
+
+color:text | rgb_vector:dense_vector
+sienna     | [160.0, 82.0, 45.0]
+peru       | [205.0, 133.0, 63.0]
+golden rod | [218.0, 165.0, 32.0]
+brown      | [165.0, 42.0, 42.0]
+firebrick  | [178.0, 34.0, 34.0]
+chartreuse | [127.0, 255.0, 0.0]
+gray       | [128.0, 128.0, 128.0]
+green      | [0.0, 128.0, 0.0]
+maroon     | [128.0, 0.0, 0.0]
+orange     | [255.0, 165.0, 0.0]
 ;
 
 knnAfterKeep
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score
 | keep rgb_vector, color, _score 
@@ -101,7 +124,7 @@ rgb_vector:dense_vector
 ;
 
 knnAfterDrop
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score
 | drop primary
@@ -120,7 +143,7 @@ lime           | [0.0, 255.0, 0.0]
 ;
 
 knnAfterEval
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score
 | eval composed_name = locate(color, " ") > 0 
@@ -139,14 +162,12 @@ golden rod | true
 ;
 
 knnWithConjunction
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
-# TODO We need kNN prefiltering here so we get more candidates that pass the filter
 from colors metadata _score 
-| where knn(rgb_vector, [255,255,238], 140) and hex_code like "#FFF*" 
+| where knn(rgb_vector, [255,255,238], 10) and hex_code like "#FFF*" 
 | sort _score desc, color asc
 | keep color, hex_code, rgb_vector
-| limit 10
 ;
 
 color:text    | hex_code:keyword | rgb_vector:dense_vector
@@ -160,11 +181,10 @@ yellow        | #FFFF00          | [255.0, 255.0, 0.0]
 ;
 
 knnWithDisjunctionAndFiltersConjunction
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
-# TODO We need kNN prefiltering here so we get more candidates that pass the filter
 from colors metadata _score 
-| where (knn(rgb_vector, [0,255,255], 140) or knn(rgb_vector, [128, 0, 255], 140)) and primary == true 
+| where (knn(rgb_vector, [0,255,255], 140) or knn(rgb_vector, [128, 0, 255], 10)) and primary == true 
 | keep color, rgb_vector, _score
 | sort _score desc, color asc
 | drop _score
@@ -183,8 +203,31 @@ red        | [255.0, 0.0, 0.0]
 yellow     | [255.0, 255.0, 0.0]
 ;
 
+knnWithNegationsAndFiltersConjunction
+required_capability: knn_function_v3
+
+from colors metadata _score 
+| where (knn(rgb_vector, [0,255,255], 140) and not(primary == true and match(color, "blue"))) 
+| sort _score desc, color asc
+| keep color, rgb_vector
+| limit 10
+;
+
+color:text  | rgb_vector:dense_vector
+cyan        | [0.0, 255.0, 255.0]
+turquoise   | [64.0, 224.0, 208.0]
+aqua marine | [127.0, 255.0, 212.0]
+teal        | [0.0, 128.0, 128.0]
+silver      | [192.0, 192.0, 192.0]
+gray        | [128.0, 128.0, 128.0]
+gainsboro   | [220.0, 220.0, 220.0]
+thistle     | [216.0, 191.0, 216.0]
+lavender    | [230.0, 230.0, 250.0]
+azure       | [240.0, 255.0, 255.0]
+;
+
 knnWithNonPushableConjunction
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score
 | eval composed_name = locate(color, " ") > 0 
@@ -208,7 +251,7 @@ maroon     | false
 ;
 
 testKnnWithNonPushableDisjunctions
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score 
 | where knn(rgb_vector, [128,128,0], 140, {"similarity": 30}) or length(color) > 10 
@@ -224,7 +267,7 @@ papaya whip
 ;
 
 testKnnWithNonPushableDisjunctionsOnComplexExpressions
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors metadata _score 
 | where (knn(rgb_vector, [128,128,0], 140, {"similarity": 70}) and length(color) < 10) or (knn(rgb_vector, [128,0,128], 140, {"similarity": 60}) and primary == false) 
@@ -239,7 +282,7 @@ indigo       | false
 ;
 
 testKnnInStatsNonPushable
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 
 from colors 
 | where length(color) < 10 
@@ -251,7 +294,7 @@ c: long
 ;
 
 testKnnInStatsWithGrouping
-required_capability: knn_function_v2
+required_capability: knn_function_v3
 required_capability: full_text_functions_in_stats_where
 
 from colors 

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KnnFunctionIT.java

@@ -114,6 +114,29 @@ public void testKnnNonPushedDown() {
         }
     }
 
+    public void testKnnWithPrefilters() {
+        float[] queryVector = new float[numDims];
+        Arrays.fill(queryVector, 1.0f);
+
+        // We retrieve 5 from knn, but must be prefiltered with id > 5 or no result will be returned as it would be post-filtered
+        var query = String.format(Locale.ROOT, """
+            FROM test METADATA _score
+            | WHERE knn(vector, %s, 5) AND id > 5
+            | KEEP id, floats, _score, vector
+            | SORT _score DESC
+            | LIMIT 5
+            """, Arrays.toString(queryVector));
+
+        try (var resp = run(query)) {
+            assertColumnNames(resp.columns(), List.of("id", "floats", "_score", "vector"));
+            assertColumnTypes(resp.columns(), List.of("integer", "double", "double", "dense_vector"));
+
+            List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
+            // K = 5, 1 more for every id > 10
+            assertEquals(5, valuesList.size());
+        }
+    }
+
     public void testKnnWithLookupJoin() {
         float[] queryVector = new float[numDims];
         Arrays.fill(queryVector, 1.0f);
@@ -136,7 +159,7 @@ public void testKnnWithLookupJoin() {
 
     @Before
     public void setup() throws IOException {
-        assumeTrue("Needs KNN support", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
+        assumeTrue("Needs KNN support", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
 
         var indexName = "test";
         var client = client().admin().indices();
@@ -163,7 +186,7 @@ public void setup() throws IOException {
         var createRequest = client.prepareCreate(indexName).setMapping(mapping).setSettings(settingsBuilder.build());
         assertAcked(createRequest);
 
-        numDocs = randomIntBetween(10, 20);
+        numDocs = randomIntBetween(15, 25);
         numDims = randomIntBetween(3, 10);
         IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
         float value = 0.0f;
@@ -202,6 +225,5 @@ private void createAndPopulateLookupIndex(IndicesAdminClient client, String look
 
         var createRequest = client.prepareCreate(lookupIndexName).setMapping(mapping).setSettings(settingsBuilder.build());
         assertAcked(createRequest);
-
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1213,7 +1213,7 @@ public enum Cap {
         /**
          * Support knn function
          */
-        KNN_FUNCTION_V2(Build.current().isSnapshot()),
+        KNN_FUNCTION_V3(Build.current().isSnapshot()),
 
         /**
          * Support for the LIKE operator with a list of wildcards.

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java

@@ -481,7 +481,7 @@ private static FunctionDefinition[][] snapshotFunctions() {
                 def(FirstOverTime.class, uni(FirstOverTime::new), "first_over_time"),
                 def(Score.class, uni(Score::new), Score.NAME),
                 def(Term.class, bi(Term::new), "term"),
-                def(Knn.class, Knn::new, "knn"),
+                def(Knn.class, quad(Knn::new), "knn"),
                 def(StGeohash.class, StGeohash::new, "st_geohash"),
                 def(StGeohashToLong.class, StGeohashToLong::new, "st_geohash_to_long"),
                 def(StGeohashToString.class, StGeohashToString::new, "st_geohash_to_string"),
@@ -1208,4 +1208,8 @@ private static <T extends Function> TernaryBuilder<T> tri(TernaryBuilder<T> func
         return function;
     }
 
+    private static <T extends Function> QuaternaryBuilder<T> quad(QuaternaryBuilder<T> function) {
+        return function;
+    }
+
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java

@@ -166,20 +166,19 @@ public boolean equals(Object obj) {
 
     @Override
     public Translatable translatable(LucenePushdownPredicates pushdownPredicates) {
-        // In isolation, full text functions are pushable to source. We check if there are no disjunctions in Or conditions
         return Translatable.YES;
     }
 
     @Override
     public Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
-        return queryBuilder != null ? new TranslationAwareExpressionQuery(source(), queryBuilder) : translate(handler);
+        return queryBuilder != null ? new TranslationAwareExpressionQuery(source(), queryBuilder) : translate(pushdownPredicates, handler);
     }
 
     public QueryBuilder queryBuilder() {
         return queryBuilder;
     }
 
-    protected abstract Query translate(TranslatorHandler handler);
+    protected abstract Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler);
 
     public abstract Expression replaceQueryBuilder(QueryBuilder queryBuilder);
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/Kql.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.esql.querydsl.query.KqlQuery;
 
@@ -93,7 +94,7 @@ protected NodeInfo<? extends Expression> info() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
+    protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         return new KqlQuery(source(), Objects.toString(queryAsObject()));
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/Match.java

@@ -35,6 +35,7 @@
 import org.elasticsearch.xpack.esql.expression.function.OptionalArgument;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.esql.querydsl.query.MatchQuery;
@@ -423,7 +424,7 @@ public Object queryAsObject() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
+    protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         var fieldAttribute = fieldAsFieldAttribute();
         Check.notNull(fieldAttribute, "Match must have a field attribute as the first argument");
         String fieldName = getNameFromFieldAttribute(fieldAttribute);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/MatchPhrase.java

@@ -32,6 +32,7 @@
 import org.elasticsearch.xpack.esql.expression.function.OptionalArgument;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.esql.querydsl.query.MatchPhraseQuery;
@@ -278,7 +279,7 @@ public Object queryAsObject() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
+    protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         var fieldAttribute = fieldAsFieldAttribute();
         Check.notNull(fieldAttribute, "MatchPhrase must have a field attribute as the first argument");
         String fieldName = getNameFromFieldAttribute(fieldAttribute);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/MultiMatch.java

@@ -31,6 +31,7 @@
 import org.elasticsearch.xpack.esql.expression.function.OptionalArgument;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.esql.querydsl.query.MultiMatchQuery;
@@ -335,7 +336,7 @@ protected NodeInfo<? extends Expression> info() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
+    protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         Map<String, Float> fieldsWithBoost = new HashMap<>();
         for (Expression field : fields) {
             var fieldAttribute = Match.fieldAsFieldAttribute(field);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryString.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.xpack.esql.expression.function.OptionalArgument;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
 
 import java.io.IOException;
@@ -345,7 +346,7 @@ protected NodeInfo<? extends Expression> info() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
+    protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         return new QueryStringQuery(source(), Objects.toString(queryAsObject()), Map.of(), queryStringOptions());
     }