Add negations tests

Author: carlosdelest

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

@@ -81,6 +81,29 @@ 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_v3
 
@@ -141,12 +164,10 @@ golden rod | true
 knnWithConjunction
 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
@@ -182,6 +203,29 @@ 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_v3
 

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

@@ -1944,6 +1944,44 @@ public void testPushDownConjunctionsToKnnPrefilter() {
         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());
 
@@ -2007,6 +2045,57 @@ public void testNotPushDownKnnWithNonPushablePrefilters() {
         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 notKeywordFilter = wrapWithSingleQuery(
+            query,
+            unscore(boolQuery().mustNot(unscore(termQuery("keyword", "test")))),
+            "keyword",
+            new Source(2, 74, "keyword == \"test\"")
+        );
+
+        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 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));
+
+        // Build the main boolean query structure
+        BoolQueryBuilder expectedQuery = boolQuery()
+            .must(notKeywordFilter)  // 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());