Patch for Lucene bug 14857

docs/changelog/130254.yaml

@@ -0,0 +1,5 @@
+pr: 130254
+summary: Patch for Lucene bug 14857
+area: Vector Search
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/search/vectors/ESDiversifyingChildrenByteKnnVectorQuery.java

@@ -9,15 +9,20 @@
 
 package org.elasticsearch.search.vectors;
 
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.search.join.DiversifyingChildrenByteKnnVectorQuery;
 import org.elasticsearch.search.profile.query.QueryProfiler;
 
+import java.io.IOException;
+
 public class ESDiversifyingChildrenByteKnnVectorQuery extends DiversifyingChildrenByteKnnVectorQuery implements QueryProfilerProvider {
     private final Integer kParam;
     private long vectorOpsCount;
+    private final int k;
 
     public ESDiversifyingChildrenByteKnnVectorQuery(
         String field,
@@ -29,6 +34,18 @@ public ESDiversifyingChildrenByteKnnVectorQuery(
     ) {
         super(field, query, childFilter, numCands, parentsFilter);
         this.kParam = k;
+        this.k = numCands;
+    }
+
+    @Override
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewrittenQuery = super.rewrite(searcher);
+        if (rewrittenQuery instanceof MatchNoDocsQuery) {
+            // If the rewritten query is a MatchNoDocsQuery, we can return it directly.
+            return rewrittenQuery;
+        }
+        // We don't rewrite this query, so we return it as is.
+        return KnnScoreDocQuery.fromQuery(rewrittenQuery, kParam == null ? k : kParam, searcher);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/vectors/ESDiversifyingChildrenFloatKnnVectorQuery.java

@@ -9,15 +9,20 @@
 
 package org.elasticsearch.search.vectors;
 
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.search.join.DiversifyingChildrenFloatKnnVectorQuery;
 import org.elasticsearch.search.profile.query.QueryProfiler;
 
+import java.io.IOException;
+
 public class ESDiversifyingChildrenFloatKnnVectorQuery extends DiversifyingChildrenFloatKnnVectorQuery implements QueryProfilerProvider {
     private final Integer kParam;
     private long vectorOpsCount;
+    private final int k;
 
     public ESDiversifyingChildrenFloatKnnVectorQuery(
         String field,
@@ -29,6 +34,18 @@ public ESDiversifyingChildrenFloatKnnVectorQuery(
     ) {
         super(field, query, childFilter, numCands, parentsFilter);
         this.kParam = k;
+        this.k = numCands;
+    }
+
+    @Override
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewrittenQuery = super.rewrite(searcher);
+        if (rewrittenQuery instanceof MatchNoDocsQuery) {
+            // If the rewritten query is a MatchNoDocsQuery, we can return it directly.
+            return rewrittenQuery;
+        }
+        // We don't rewrite this query, so we return it as is.
+        return KnnScoreDocQuery.fromQuery(rewrittenQuery, kParam == null ? k : kParam, searcher);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/vectors/ESKnnByteVectorQuery.java

@@ -9,11 +9,15 @@
 
 package org.elasticsearch.search.vectors;
 
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.KnnByteVectorQuery;
+import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TopDocs;
 import org.elasticsearch.search.profile.query.QueryProfiler;
 
+import java.io.IOException;
+
 public class ESKnnByteVectorQuery extends KnnByteVectorQuery implements QueryProfilerProvider {
     private final Integer kParam;
     private long vectorOpsCount;
@@ -23,6 +27,17 @@ public ESKnnByteVectorQuery(String field, byte[] target, Integer k, int numCands
         this.kParam = k;
     }
 
+    @Override
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewrittenQuery = super.rewrite(searcher);
+        if (rewrittenQuery instanceof MatchNoDocsQuery) {
+            // If the rewritten query is a MatchNoDocsQuery, we can return it directly.
+            return rewrittenQuery;
+        }
+        // We don't rewrite this query, so we return it as is.
+        return KnnScoreDocQuery.fromQuery(rewrittenQuery, kParam == null ? k : kParam, searcher);
+    }
+
     @Override
     protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
         // if k param is set, we get only top k results from each shard

server/src/main/java/org/elasticsearch/search/vectors/ESKnnFloatVectorQuery.java

@@ -9,11 +9,15 @@
 
 package org.elasticsearch.search.vectors;
 
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.KnnFloatVectorQuery;
+import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TopDocs;
 import org.elasticsearch.search.profile.query.QueryProfiler;
 
+import java.io.IOException;
+
 public class ESKnnFloatVectorQuery extends KnnFloatVectorQuery implements QueryProfilerProvider {
     private final Integer kParam;
     private long vectorOpsCount;
@@ -23,6 +27,17 @@ public ESKnnFloatVectorQuery(String field, float[] target, Integer k, int numCan
         this.kParam = k;
     }
 
+    @Override
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewrittenQuery = super.rewrite(searcher);
+        if (rewrittenQuery instanceof MatchNoDocsQuery) {
+            // If the rewritten query is a MatchNoDocsQuery, we can return it directly.
+            return rewrittenQuery;
+        }
+        // We don't rewrite this query, so we return it as is.
+        return KnnScoreDocQuery.fromQuery(rewrittenQuery, kParam == null ? k : kParam, searcher);
+    }
+
     @Override
     protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
         // if k param is set, we get only top k results from each shard

server/src/main/java/org/elasticsearch/search/vectors/KnnScoreDocQuery.java

@@ -20,6 +20,7 @@
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.Weight;
 
 import java.io.IOException;
@@ -47,6 +48,19 @@ public class KnnScoreDocQuery extends Query {
     // an object identifying the reader context that was used to build this query
     private final Object contextIdentity;
 
+    static Query fromQuery(Query query, int k, IndexSearcher searcher) throws IOException {
+        if (query instanceof MatchNoDocsQuery) {
+            // If the rewritten query is a MatchNoDocsQuery, we can return it directly.
+            return query;
+        }
+        if (query instanceof KnnScoreDocQuery knnQuery) {
+            return knnQuery;
+        }
+        TopDocs topDocs = searcher.search(query, k);
+        assert topDocs.scoreDocs.length == k;
+        return new KnnScoreDocQuery(topDocs.scoreDocs, searcher.getIndexReader());
+    }
+
     /**
      * Creates a query.
      *
@@ -176,6 +190,9 @@ public float score() {
 
                     @Override
                     public int advanceShallow(int docId) {
+                        if (docId == NO_MORE_DOCS) {
+                            return NO_MORE_DOCS;
+                        }
                         int start = Math.max(upTo, lower);
                         int docIdIndex = Arrays.binarySearch(docs, start, upper, docId + context.docBase);
                         if (docIdIndex < 0) {