Avoid infinite loop with multiple knn expressions, add tests

carlosdelest · carlosdelest · commit ef12ee5d7c50 · 2025-07-09T18:27:31.000+02:00
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Knn.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Knn.java
@@ -287,7 +287,7 @@ public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
     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()) {
+        for (Expression filterExpression : filterExpressions()) {
             translatable = translatable.merge(TranslationAware.translatable(filterExpression, pushdownPredicates));
         }
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownConjunctionsToKnnPrefilters.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownConjunctionsToKnnPrefilters.java
@@ -17,6 +17,11 @@
 import java.util.List;
 import java.util.Stack;
 
+/**
+ * Rewrites an expression tree to push down conjunctions in the prefilter of {@link Knn} functions.
+ *  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]))}
+*/
 public class PushDownConjunctionsToKnnPrefilters extends OptimizerRules.OptimizerRule<Filter> {
 
     @Override
@@ -28,25 +33,43 @@ protected LogicalPlan rule(Filter filter) {
         return condition.equals(newCondition) ? filter : filter.with(newCondition);
     }
 
-    private static Expression pushConjunctionsToKnn(Expression expression, List<Expression> filters, Expression addedFilter) {
+    /**
+     * Updates knn function prefilters. This method processes conjunctions so knn functions on one side of the conjunction receive
+     * the other side of the conjunction as a prefilter
+     *
+     * @param expression expression to process recursively
+     * @param filters current filters to apply to the expression. They contain expressions on the other side of the traversed conjunctions
+     * @param addedFilter a new filter to add to the list of filters for the processing
+     * @return the updated expression, or the original expression if it doesn't need to be updated
+     */
+    private static Expression pushConjunctionsToKnn(Expression expression, Stack<Expression> filters, Expression addedFilter) {
         if (addedFilter != null) {
-            filters.add(addedFilter);
+            filters.push(addedFilter);
         }
-        Expression result = switch(expression) {
+        Expression result = switch (expression) {
             case And and:
+                // Traverse both sides of the And, using the other side as the added filter
                 Expression newLeft = pushConjunctionsToKnn(and.left(), filters, and.right());
                 Expression newRight = pushConjunctionsToKnn(and.right(), filters, and.left());
                 if (newLeft.equals(and.left()) && newRight.equals(and.right())) {
                     yield and;
                 }
                 yield and.replaceChildrenSameSize(List.of(newLeft, newRight));
             case Knn knn:
-                yield knn.withFilters(List.copyOf(filters));
+                // 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 recveive the other knn functions as filters and would be constantly
+                // updating
+                List<Expression> newFilters = new ArrayList<>(filters);
+                if (newFilters.size() == knn.filterExpressions().size()) {
+                    yield knn;
+                }
+                yield knn.withFilters(newFilters);
             default:
                 List<Expression> children = expression.children();
                 boolean childrenChanged = false;
 
-                // TODO This copies transformChildren
+                // This copies transformChildren algorithm to avoid unnecessary changes
                 List<Expression> transformedChildren = null;
 
                 for (int i = 0, s = children.size(); i < s; i++) {
@@ -66,7 +89,7 @@ private static Expression pushConjunctionsToKnn(Expression expression, List<Expr
         };
 
         if (addedFilter != null) {
-            filters.removeLast();
+            filters.pop();
         }
 
         return result;
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java
@@ -1628,10 +1628,6 @@ private void testFullTextFunctionWithPushableConjunction(FullTextFunctionTestCas
         assertEquals(expected.toString(), esQuery.query().toString());
     }
 
-    public void testKnn() {
-        testFullTextFunctionWithNonPushableDisjunction(new KnnFunctionTestCase());
-    }
-
     private void testFullTextFunctionWithNonPushableDisjunction(FullTextFunctionTestCase testCase) {
         String query = String.format(Locale.ROOT, """
             from test
@@ -1652,32 +1648,6 @@ private void testFullTextFunctionWithNonPushableDisjunction(FullTextFunctionTest
         assertThat(fieldExtract.child(), instanceOf(EsQueryExec.class));
     }
 
-    public void testKnnPrefilters() {
-        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());
-    }
-
     private void testFullTextFunctionWithPushableDisjunction(FullTextFunctionTestCase testCase) {
         String query = String.format(Locale.ROOT, """
             from test
@@ -1697,12 +1667,11 @@ private void testFullTextFunctionWithPushableDisjunction(FullTextFunctionTestCas
     }
 
     private FullTextFunctionTestCase randomFullTextFunctionTestCase() {
-        return switch (randomIntBetween(0, 4)) {
+        return switch (randomIntBetween(0, 3)) {
             case 0 -> new MatchFunctionTestCase();
             case 1 -> new MatchOperatorTestCase();
             case 2 -> new KqlFunctionTestCase();
             case 3 -> new QueryStringFunctionTestCase();
-            case 4 -> new KnnFunctionTestCase();
             default -> throw new IllegalStateException("Unexpected value");
         };
     }
@@ -1861,6 +1830,150 @@ public void testFullTextFunctionWithStatsBy(FullTextFunctionTestCase testCase) {
         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 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 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 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, 45, "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, 87, "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.addFilterQuery(firstOr));
+        firstKnnQuery.addFilterQuery(secondOr);
+
+        // 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)";
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java
@@ -7954,4 +7954,49 @@ public void testMorePushDownConjunctionsAndNotDisjunctionsToKnnPrefilter() {
         var rightAndPrefilter = as(knn.filterExpressions().get(0), GreaterThan.class);
         assertThat(leftAnd.right(), equalTo(rightAndPrefilter));
     }
+
+    public void testMultipleKnnQueriesInPrefilters() {
+        assumeTrue("knn must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
+
+        /*
+            and
+                or
+                    knn(dense_vector, [0, 1, 2], 10)
+                    integer > 10
+                or
+                    keyword == "test"
+                    knn(dense_vector, [4, 5, 6], 10)
+         */
+        var 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 optimized = planTypes(query);
+
+        var limit = as(optimized, Limit.class);
+        var filter = as(limit.child(), Filter.class);
+        var and = as(filter.condition(), And.class);
+
+        // First OR (knn1 OR integer > 10)
+        var firstOr = as(and.left(), Or.class);
+        var firstKnn = as(firstOr.left(), Knn.class);
+        var integerGt = as(firstOr.right(), GreaterThan.class);
+
+        // Second OR (keyword == "test" OR knn2)
+        var secondOr = as(and.right(), Or.class);
+        var keywordEq = as(secondOr.left(), Equals.class);
+        var secondKnn = as(secondOr.right(), Knn.class);
+
+        // 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));
+
+        // 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()));
+        assertTrue(secondKnnFilters.contains(firstOrWithoutFilters));
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -287,7 +287,7 @@ public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {`
`287`	`287`	`public Translatable translatable(LucenePushdownPredicates pushdownPredicates) {`
`288`	`288`	`Translatable translatable = super.translatable(pushdownPredicates);`
`289`	`289`	`// We need to check whether filter expressions are translatable as well`
`290`		`- for(Expression filterExpression : filterExpressions()) {`
	`290`	`+ for (Expression filterExpression : filterExpressions()) {`
`291`	`291`	`translatable = translatable.merge(TranslationAware.translatable(filterExpression, pushdownPredicates));`
`292`	`292`	`}`
`293`	`293`