Fix generating QueryBuilder prefilters for knn when they cannot be pushed down

Author: carlosdelest

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

@@ -266,3 +266,20 @@ c: long       | primary: boolean
 41            | false          
 9             | true           
 ;
+
+testKnnUsesPrefiltering
+required_capability: knn_function_v3
+
+from colors metadata _score
+| where knn(rgb_vector, [255, 0, 0], 5) and primary == true
+| sort _score desc, color asc
+| keep color
+;
+
+color:text
+red            
+gray           
+black          
+magenta        
+yellow
+;

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

@@ -253,6 +253,22 @@ private Map<String, Object> knnQueryOptions() throws InvalidArgumentException {
         return matchOptions;
     }
 
+    @Override
+    public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
+        return new Knn(source(), field(), query(), k(), options(), queryBuilder, filterExpressions());
+    }
+
+    @Override
+    public Translatable translatable(LucenePushdownPredicates pushdownPredicates) {
+        Translatable translatable = super.translatable(pushdownPredicates);
+        // We need to check whether filter expressions are translatable as well
+        for (Expression filterExpression : filterExpressions()) {
+            translatable = translatable.merge(TranslationAware.translatable(filterExpression, pushdownPredicates));
+        }
+
+        return translatable;
+    }
+
     @Override
     protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
         var fieldAttribute = Match.fieldAsFieldAttribute(field());
@@ -272,28 +288,19 @@ protected Query translate(LucenePushdownPredicates pushdownPredicates, Translato
 
         List<QueryBuilder> filterQueries = new ArrayList<>();
         for (Expression filterExpression : filterExpressions()) {
-            filterQueries.add(handler.asQuery(pushdownPredicates, filterExpression).toQueryBuilder());
+            if (filterExpression instanceof TranslationAware translationAware) {
+                // We can only translate filter expressions that are translatable. In case any is not translatable,
+                // Knn won't be pushed down as it will not be translatable so it's safe not to translate all filters and check them
+                // when creating an evaluator for the non-pushed down query
+                if (translationAware.translatable(pushdownPredicates) == Translatable.YES) {
+                    filterQueries.add(handler.asQuery(pushdownPredicates, filterExpression).toQueryBuilder());
+                }
+            }
         }
 
         return new KnnQuery(source(), fieldName, queryAsFloats, opts, filterQueries);
     }
 
-    @Override
-    public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
-        return new Knn(source(), field(), query(), k(), options(), queryBuilder, filterExpressions());
-    }
-
-    @Override
-    public Translatable translatable(LucenePushdownPredicates pushdownPredicates) {
-        Translatable translatable = super.translatable(pushdownPredicates);
-        // We need to check whether filter expressions are translatable as well
-        for (Expression filterExpression : filterExpressions()) {
-            translatable = translatable.merge(TranslationAware.translatable(filterExpression, pushdownPredicates));
-        }
-
-        return translatable;
-    }
-
     public Expression withFilters(List<Expression> filterExpressions) {
         return new Knn(source(), field(), query(), k(), options(), queryBuilder(), filterExpressions);
     }

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

@@ -1979,30 +1979,50 @@ public void testMultipleKnnQueriesInPrefilters() {
         var queryExec = as(field.child(), EsQueryExec.class);
 
         KnnVectorQueryBuilder firstKnnQuery = new KnnVectorQueryBuilder("dense_vector", new float[] { 0, 1, 2 }, 10, null, null, null);
+        KnnVectorQueryBuilder firstKnnQueryAsFilter = 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, 45, "integer > 10")
+            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);
+        KnnVectorQueryBuilder secondKnnQueryAsFilter = 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, 87, "keyword == \"test\"")
+            new Source(2, 66, "keyword == \"test\"")
         );
 
         // First OR (knn1 OR integer > 10)
         var firstOr = boolQuery().should(firstKnnQuery).should(integerRangeQuery);
+        var firstOrAsFilter = boolQuery().should(firstKnnQueryAsFilter).should(integerRangeQuery);
         // Second OR (keyword == "test" OR knn2)
-        var secondOr = boolQuery().should(keywordQuery).should(secondKnnQuery.addFilterQuery(firstOr));
-        firstKnnQuery.addFilterQuery(secondOr);
+        var secondOr = boolQuery().should(keywordQuery).should(secondKnnQuery);
+        var secondOrAsFilter = boolQuery().should(keywordQuery).should(secondKnnQueryAsFilter);
+        // Add prefilters to the knn queries. knn queries in prefilters don't have prefilters so we use copies of the queries
+        firstKnnQuery.addFilterQuery(secondOrAsFilter);
+        secondKnnQuery.addFilterQuery(firstOrAsFilter);
 
         // Top-level AND combining both ORs
         var expectedQuery = boolQuery().must(firstOr).must(secondOr);