ESQL - KNN functions with non-pushed down filters

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluator.java

@@ -61,10 +61,16 @@ protected LuceneQueryEvaluator(BlockFactory blockFactory, ShardConfig[] shards)
     }
 
     public Block executeQuery(Page page) {
-        // Lucene based operators retrieve DocVectors as first block
-        Block block = page.getBlock(0);
-        assert block instanceof DocBlock : "LuceneQueryExpressionEvaluator expects DocBlock as input";
-        DocVector docs = (DocVector) block.asVector();
+        // Search for DocVector block
+        Block docBlock = null;
+        for (int i = 0; i < page.getBlockCount(); i++) {
+            if (page.getBlock(i) instanceof DocBlock) {
+                docBlock = page.getBlock(i);
+                break;
+            }
+        }
+        assert docBlock != null : "LuceneQueryExpressionEvaluator expects a DocBlock";
+        DocVector docs = (DocVector) docBlock.asVector();
         try {
             if (docs.singleSegmentNonDecreasing()) {
                 return evalSingleSegmentNonDecreasing(docs).asBlock();

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ScoreOperator.java

@@ -9,7 +9,6 @@
 
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
-import org.elasticsearch.compute.data.DocVector;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.DoubleVector;
 import org.elasticsearch.compute.data.Page;
@@ -46,9 +45,9 @@ public ScoreOperator(BlockFactory blockFactory, ExpressionScorer scorer, int sco
 
     @Override
     protected Page process(Page page) {
-        assert page.getBlockCount() >= 2 : "Expected at least 2 blocks, got " + page.getBlockCount();
-        assert page.getBlock(0).asVector() instanceof DocVector : "Expected a DocVector, got " + page.getBlock(0).asVector();
-        assert page.getBlock(1).asVector() instanceof DoubleVector : "Expected a DoubleVector, got " + page.getBlock(1).asVector();
+        assert page.getBlockCount() > scoreBlockPosition : "Expected to get a score block in position " + scoreBlockPosition;
+        assert page.getBlock(scoreBlockPosition).asVector() instanceof DoubleVector
+            : "Expected a DoubleVector as a score block, got " + page.getBlock(scoreBlockPosition).asVector();
 
         Block[] blocks = new Block[page.getBlockCount()];
         for (int i = 0; i < page.getBlockCount(); i++) {

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_v3
+required_capability: knn_function_v4
 
 // tag::knn-function[]
 from colors metadata _score 
@@ -30,7 +30,7 @@ chartreuse | [127.0, 255.0, 0.0]
 ;
 
 knnSearchWithSimilarityOption
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score 
 | where knn(rgb_vector, [255,192,203], 140, {"similarity": 40})
@@ -46,7 +46,7 @@ wheat      | [245.0, 222.0, 179.0]
 ;
 
 knnHybridSearch
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score 
 | where match(color, "blue") or knn(rgb_vector, [65,105,225], 10)
@@ -68,7 +68,7 @@ yellow     | [255.0, 255.0, 0.0]
 ;
 
 knnWithPrefilter
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score
 | where knn(rgb_vector, [128,128,0], 10) and (match(color, "olive") or match(color, "green")) 
@@ -82,7 +82,7 @@ green            | [0.0, 128.0, 0.0]
 ;
 
 knnWithNegatedPrefilter
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score
 | where knn(rgb_vector, [128,128,0], 10) and not (match(color, "olive") or match(color, "chocolate")) 
@@ -105,7 +105,7 @@ orange     | [255.0, 165.0, 0.0]
 ;
 
 knnAfterKeep
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score
 | keep rgb_vector, color, _score 
@@ -124,7 +124,7 @@ rgb_vector:dense_vector
 ;
 
 knnAfterDrop
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score
 | drop primary
@@ -143,7 +143,7 @@ lime           | [0.0, 255.0, 0.0]
 ;
 
 knnAfterEval
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score
 | eval composed_name = locate(color, " ") > 0 
@@ -162,7 +162,7 @@ golden rod | true
 ;
 
 knnWithConjunction
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score 
 | where knn(rgb_vector, [255,255,238], 10) and hex_code like "#FFF*" 
@@ -181,7 +181,7 @@ yellow        | #FFFF00          | [255.0, 255.0, 0.0]
 ;
 
 knnWithDisjunctionAndFiltersConjunction
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score 
 | where (knn(rgb_vector, [0,255,255], 140) or knn(rgb_vector, [128, 0, 255], 10)) and primary == true 
@@ -204,7 +204,7 @@ yellow     | [255.0, 255.0, 0.0]
 ;
 
 knnWithNegationsAndFiltersConjunction
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score 
 | where (knn(rgb_vector, [0,255,255], 140) and not(primary == true and match(color, "blue"))) 
@@ -227,62 +227,76 @@ azure       | [240.0, 255.0, 255.0]
 ;
 
 knnWithNonPushableConjunction
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score
 | eval composed_name = locate(color, " ") > 0 
-| where knn(rgb_vector, [128,128,0], 140) and composed_name == false
+| where knn(rgb_vector, [128,128,0], 10) and composed_name == false
 | sort _score desc, color asc
-| keep color, composed_name
-| limit 10
+| keep color
 ;
 
-color:text | composed_name:boolean
-olive      | false
-sienna     | false
-chocolate  | false
-peru       | false
-brown      | false
-firebrick  | false
-chartreuse | false
-gray       | false
-green      | false
-maroon     | false
+color:text
+olive
+sienna
+chocolate
+peru
+brown
+firebrick
+chartreuse
+green
+maroon
 ;
 
 testKnnWithNonPushableDisjunctions
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors metadata _score 
-| where knn(rgb_vector, [128,128,0], 140, {"similarity": 30}) or length(color) > 10 
+| where knn(rgb_vector, [128,128,0], 10) or length(color) > 10 
 | sort _score desc, color asc
 | keep color 
 ;
 
 color:text
 olive
+sienna
+chocolate
+peru
+golden rod
+brown
+firebrick
+chartreuse
+green
+maroon
 aqua marine
 lemon chiffon
 papaya whip
 ;
 
-testKnnWithNonPushableDisjunctionsOnComplexExpressions
-required_capability: knn_function_v3
+testKnnWithNonPushableConjunctionsOnComplexExpressions
+required_capability: knn_function_v4
 
 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) 
+| where knn(rgb_vector, [128,128,0], 10) and length(color) < 7 and knn(rgb_vector, [128,0,128], 10) and primary == false 
 | sort _score desc, color asc
-| keep color, primary
+| keep color
 ;
 
-color:text   | primary:boolean
-olive        | false
-purple       | false
-indigo       | false
+color:text
+olive
+purple
+indigo
+sienna
+brown
+peru
+maroon
+navy
+tomato
+orange
 ;
 
 testKnnInStatsNonPushable
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 
 from colors 
 | where length(color) < 10 
@@ -294,7 +308,7 @@ c: long
 ;
 
 testKnnInStatsWithGrouping
-required_capability: knn_function_v3
+required_capability: knn_function_v4
 required_capability: full_text_functions_in_stats_where
 
 from colors 

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
@@ -31,7 +32,9 @@
 import static org.elasticsearch.index.IndexMode.LOOKUP;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
 
+@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE", reason = "debug")
 public class KnnFunctionIT extends AbstractEsqlIntegTestCase {
 
     private final Map<Integer, List<Float>> indexedVectors = new HashMap<>();
@@ -157,9 +160,71 @@ public void testKnnWithLookupJoin() {
         );
     }
 
+    public void testKnnNotPushedDown() {
+        float[] queryVector = new float[numDims];
+        Arrays.fill(queryVector, 1.0f);
+
+        // We retrieve 5 from knn, and 5 from the non-pushed down disjunction. They are disjoint so we get 10 as a result
+        var query = String.format(Locale.ROOT, """
+            FROM test
+            | WHERE knn(vector, %s, 5) OR (length(keyword) > 5 AND length(keyword) <= 10)
+            | KEEP id, vector, keyword
+            | SORT id ASC
+            | LIMIT 20
+            """, Arrays.toString(queryVector));
+
+        try (var resp = run(query)) {
+            assertColumnNames(resp.columns(), List.of("id", "vector", "keyword"));
+            assertColumnTypes(resp.columns(), List.of("integer", "dense_vector", "keyword"));
+
+            List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
+            assertEquals(10, valuesList.size());
+        }
+    }
+
+    public void testKnnPrefiltersNotPushedDown() {
+        float[] queryVector = new float[numDims];
+        Arrays.fill(queryVector, 1.0f);
+
+        // We retrieve 5 from knn, but must be prefiltered with the non-pushed down conjunction
+        var query = String.format(Locale.ROOT, """
+            FROM test
+            | WHERE knn(vector, %s, 5) AND length(keyword) > 5 AND length(keyword) <= 10
+            | SORT id ASC
+            | LIMIT 20
+            """, Arrays.toString(queryVector));
+
+        try (var resp = run(query)) {
+            // No added columns
+            assertThat(resp.columns().size(), equalTo(4));
+            List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
+            assertEquals(5, valuesList.size());
+        }
+    }
+
+    public void testKnnPrefiltersNotPushedDownWithScoring() {
+        float[] queryVector = new float[numDims];
+        Arrays.fill(queryVector, 1.0f);
+
+        // We retrieve 5 from knn, but must be prefiltered with the non-pushed down conjunction
+        var query = String.format(Locale.ROOT, """
+            FROM test METADATA _score
+            | WHERE knn(vector, %s, 5) AND length(keyword) > 5 AND length(keyword) <= 10
+            | SORT id ASC
+            | LIMIT 20
+            """, Arrays.toString(queryVector));
+
+        try (var resp = run(query)) {
+            // No added columns
+            assertThat(resp.columns().size(), equalTo(5));
+            List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
+            assertEquals(5, valuesList.size());
+        }
+    }
+
     @Before
     public void setup() throws IOException {
-        assumeTrue("Needs KNN support", EsqlCapabilities.Cap.KNN_FUNCTION_V3.isEnabled());
+        assumeTrue("Needs KNN support", EsqlCapabilities.Cap.KNN_FUNCTION_V4.isEnabled());
 
         var indexName = "test";
         var client = client().admin().indices();
@@ -176,6 +241,9 @@ public void setup() throws IOException {
             .startObject("floats")
             .field("type", "float")
             .endObject()
+            .startObject("keyword")
+            .field("type", "keyword")
+            .endObject()
             .endObject()
             .endObject();
 
@@ -195,7 +263,8 @@ public void setup() throws IOException {
             for (int j = 0; j < numDims; j++) {
                 vector.add(value++);
             }
-            docs[i] = prepareIndex("test").setId("" + i).setSource("id", String.valueOf(i), "floats", vector, "vector", vector);
+            docs[i] = prepareIndex("test").setId("" + i)
+                .setSource("id", String.valueOf(i), "floats", vector, "vector", vector, "keyword", randomAlphaOfLength(i));
             indexedVectors.put(i, vector);
         }
 

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

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