knn functions won't appear as prefilters for other knn functions

Author: carlosdelest

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

@@ -10,15 +10,18 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 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.BinaryLogic;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Stack;
 
 /**
  * Rewrites an expression tree to push down conjunctions in the prefilter of {@link Knn} functions.
+ * knn functions won't contain other knn functions as a prefilter, to avoid circular dependencies.
  *  Given an expression tree like {@code (A OR B) AND (C AND knn())} this rule will rewrite it to
  *     {@code (A OR B) AND (C AND knn(filterExpressions = [(A OR B), C]))}
 */
@@ -56,12 +59,12 @@ private static Expression pushConjunctionsToKnn(Expression expression, Stack<Exp
                 }
                 yield and.replaceChildrenSameSize(List.of(newLeft, newRight));
             case Knn knn:
-                // Create a copy of the filters, and check the number of existing filters. We don't check for equality
-                // as having two knn functions on opposite sides of the And is valid, but could lead to infinite number
-                // of changes as the knn functions would receive the other knn functions as filters and would be constantly
-                // updating
-                List<Expression> newFilters = new ArrayList<>(filters);
-                if (newFilters.size() == knn.filterExpressions().size()) {
+                // We don't want knn expressions to have other knn expressions as a prefilter to avoid circular dependencies
+                List<Expression> newFilters = filters.stream()
+                    .map(PushDownConjunctionsToKnnPrefilters::removeKnn)
+                    .filter(Objects::nonNull)
+                    .toList();
+                if (newFilters.equals(knn.filterExpressions())) {
                     yield knn;
                 }
                 yield knn.withFilters(newFilters);
@@ -94,4 +97,34 @@ private static Expression pushConjunctionsToKnn(Expression expression, Stack<Exp
 
         return result;
     }
+
+    /**
+     * Removes knn functions from the expression tree
+     * @param expression expression to process
+     * @return expression without knn functions, or null if the expression is a knn function
+     */
+    private static Expression removeKnn(Expression expression) {
+        if (expression.children().isEmpty()) {
+            return expression;
+        }
+        if (expression instanceof Knn) {
+            return null;
+        }
+
+        List<Expression> filteredChildren = expression.children()
+            .stream()
+            .map(PushDownConjunctionsToKnnPrefilters::removeKnn)
+            .filter(Objects::nonNull)
+            .toList();
+        if (filteredChildren.equals(expression.children())) {
+            return expression;
+        } else if (filteredChildren.isEmpty()) {
+            return null;
+        } else if (expression instanceof BinaryLogic && filteredChildren.size() == 1) {
+            // Simplify an AND / OR expression to a single child
+            return filteredChildren.getFirst();
+        } else {
+            return expression.replaceChildrenSameSize(filteredChildren);
+        }
+    }
 }

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

@@ -1944,7 +1944,6 @@ public void testPushDownConjunctionsToKnnPrefilter() {
         assertEquals(expectedQuery.toString(), queryExec.query().toString());
     }
 
-
     public void testPushDownNegatedConjunctionsToKnnPrefilter() {
         assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
 
@@ -2050,7 +2049,8 @@ public void testPushDownComplexNegationsToKnnPrefilter() {
 
         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)))
+            | 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"));
 
@@ -2060,11 +2060,17 @@ public void testPushDownComplexNegationsToKnnPrefilter() {
         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(2, 74, "keyword == \"test\"")
+            new Source(3, 6, "NOT ((keyword == \"test\") or knn(dense_vector, [4, 5, 6], 10))")
         );
 
         QueryBuilder notIntegerGt10 = wrapWithSingleQuery(
@@ -2075,21 +2081,13 @@ public void testPushDownComplexNegationsToKnnPrefilter() {
         );
 
         KnnVectorQueryBuilder firstKnn = new KnnVectorQueryBuilder("dense_vector", new float[] { 0, 1, 2 }, 10, null, null, null);
-        KnnVectorQueryBuilder firstKnnFilter = 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);
-        KnnVectorQueryBuilder secondKnnFilter = new KnnVectorQueryBuilder("dense_vector", new float[] { 4, 5, 6 }, 10, null, null, null);
 
-        firstKnn.addFilterQuery(boolQuery()
-            .must(notKeywordFilter)
-            .must(unscore(boolQuery().mustNot(secondKnnFilter))));
-
-        secondKnn.addFilterQuery(boolQuery()
-            .should(firstKnnFilter)
-            .should(notIntegerGt10));
+        firstKnn.addFilterQuery(notKeywordFilter);
+        secondKnn.addFilterQuery(notIntegerGt10);
 
         // Build the main boolean query structure
-        BoolQueryBuilder expectedQuery = boolQuery()
-            .must(notKeywordFilter)  // NOT (keyword == "test")
+        BoolQueryBuilder expectedQuery = boolQuery().must(notKeywordQuery)  // NOT (keyword == "test")
             .must(unscore(boolQuery().mustNot(secondKnn)))
             .must(boolQuery().should(firstKnn).should(notIntegerGt10));
 
@@ -2112,14 +2110,6 @@ 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,
@@ -2130,14 +2120,6 @@ public void testMultipleKnnQueriesInPrefilters() {
 
         // 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(
@@ -2149,13 +2131,10 @@ public void testMultipleKnnQueriesInPrefilters() {
 
         // 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);
-        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);
+        firstKnnQuery.addFilterQuery(keywordQuery);
+        secondKnnQuery.addFilterQuery(integerRangeQuery);
 
         // Top-level AND combining both ORs
         var expectedQuery = boolQuery().must(firstOr).must(secondOr);

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

@@ -8019,13 +8019,13 @@ public void testMultipleKnnQueriesInPrefilters() {
         // First KNN should have the second OR as its filter
         List<Expression> firstKnnFilters = firstKnn.filterExpressions();
         assertThat(firstKnnFilters.size(), equalTo(1));
-        var secondOrWithoutFilters = secondOr.replaceChildren(List.of(secondOr.left(), secondKnn.withFilters(List.of())));
-        assertTrue(firstKnnFilters.contains(secondOrWithoutFilters));
+        var secondOrWithoutKnn = secondOr.replaceChildren(List.of(secondOr.left(), Literal.TRUE));
+        assertTrue(firstKnnFilters.contains(secondOrWithoutKnn));
 
         // Second KNN should have the first OR as its filter
         List<Expression> secondKnnFilters = secondKnn.filterExpressions();
         assertThat(secondKnnFilters.size(), equalTo(1));
-        var firstOrWithoutFilters = firstOr.replaceChildren(List.of(firstKnn.withFilters(List.of()), firstOr.right()));
+        var firstOrWithoutFilters = firstOr.replaceChildren(List.of(Literal.TRUE, firstOr.right()));
         assertTrue(secondKnnFilters.contains(firstOrWithoutFilters));
     }
 }