Knn dont filter match all (#14936)

If a given filter is a MatchAllQuery, we shouldn't bother going through the filtered search path. It adds unnecessary overhead through:

building a bitset for filter checking
checking a bit set to see if a vector matches a filter

Author: benwtrent

lucene/CHANGES.txt

@@ -210,6 +210,8 @@ Optimizations
 
 * GITHUB#14935: Speed up PostingsEnum#nextPostings when block is encoded as bitset. (Guo Feng)
 
+* GITHUB#14936: Don't do the filtered knn path when the provided filter is a MatchAllDocsQuery. (Ben Trent)
+
 Changes in Runtime Behavior
 ---------------------
 * GITHUB#14823: Decrease TieredMergePolicy's default number of segments per

lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java

@@ -80,16 +80,28 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
 
     final Weight filterWeight;
     if (filter != null) {
-      BooleanQuery booleanQuery =
-          new BooleanQuery.Builder()
-              .add(filter, BooleanClause.Occur.FILTER)
-              .add(new FieldExistsQuery(field), BooleanClause.Occur.FILTER)
-              .build();
-      Query rewritten = indexSearcher.rewrite(booleanQuery);
-      if (rewritten.getClass() == MatchNoDocsQuery.class) {
-        return rewritten;
+      // rewrite inner filter query first to determine if its a match all
+      // or match no docs query, so we can skip the knn search
+      Query rewrittenFilter = filter.rewrite(indexSearcher);
+      if (rewrittenFilter.getClass() == MatchNoDocsQuery.class) {
+        // If the filter is a match no docs query, we can also skip it
+        return rewrittenFilter;
+      }
+      if (rewrittenFilter.getClass() != MatchAllDocsQuery.class) {
+        BooleanQuery booleanQuery =
+            new BooleanQuery.Builder()
+                .add(filter, BooleanClause.Occur.FILTER)
+                .add(new FieldExistsQuery(field), BooleanClause.Occur.FILTER)
+                .build();
+        Query rewritten = indexSearcher.rewrite(booleanQuery);
+        if (rewritten.getClass() == MatchNoDocsQuery.class) {
+          return rewritten;
+        }
+        filterWeight = rewritten.createWeight(indexSearcher, ScoreMode.COMPLETE_NO_SCORES, 1f);
+      } else {
+        // If the filter is a match all docs query, we can skip it
+        filterWeight = null;
       }
-      filterWeight = indexSearcher.createWeight(rewritten, ScoreMode.COMPLETE_NO_SCORES, 1f);
     } else {
       filterWeight = null;
     }

lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java

@@ -250,6 +250,20 @@ public void testFilterWithNoVectorMatches() throws IOException {
     }
   }
 
+  public void testMatchAllFilter() throws IOException {
+    try (Directory indexStore =
+            getIndexStore("field", new float[] {0, 1}, new float[] {1, 2}, new float[] {0, 0});
+        IndexReader reader = DirectoryReader.open(indexStore)) {
+      IndexSearcher searcher = newSearcher(reader);
+
+      // make sure we don't drop to exact search, even though the filter matches fewer than k docs
+      Query kvq =
+          getThrowingKnnVectorQuery("field", new float[] {0, 0}, 10, new MatchAllDocsQuery());
+      TopDocs topDocs = searcher.search(kvq, 3);
+      assertEquals(3, topDocs.totalHits.value());
+    }
+  }
+
   /** testDimensionMismatch */
   public void testDimensionMismatch() throws IOException {
     try (Directory indexStore =

lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java

@@ -81,13 +81,14 @@
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.KnnFloatVectorQuery;
-import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.search.VectorScorer;
+import org.apache.lucene.search.knn.KnnCollectorManager;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.FSDirectory;
 import org.apache.lucene.tests.codecs.asserting.AssertingKnnVectorsFormat;
@@ -1980,9 +1981,8 @@ protected void assertRecall(VectorSimilarityFunction similarity, double min, dou
       for (String queryString : testQueries) {
         computeLineEmbedding(queryString, queryEmbedding);
 
-        // pass match-all "filter" to force full traversal, bypassing graph
-        KnnFloatVectorQuery exactQuery =
-            new KnnFloatVectorQuery("field", queryEmbedding, 1000, new MatchAllDocsQuery());
+        // gather exact results first
+        Query exactQuery = buildExactKnnQuery("field", queryEmbedding, 10000);
         assertEquals(numQueries, searcher.count(exactQuery)); // Same for exact search
 
         KnnFloatVectorQuery query = new KnnFloatVectorQuery("field", queryEmbedding, efSearch);
@@ -2202,4 +2202,18 @@ public void testMergeOffHeapByteSizeMaps() {
     assertEquals(101L, (long) r.get("d"));
     assertEquals(6L, (long) r.get("e"));
   }
+
+  private static Query buildExactKnnQuery(String fieldName, float[] queryVector, int totalDocs) {
+    return new KnnFloatVectorQuery(fieldName, queryVector, totalDocs) {
+      @Override
+      protected TopDocs approximateSearch(
+          LeafReaderContext context,
+          Bits acceptDocs,
+          int visitedLimit,
+          KnnCollectorManager knnCollectorManager)
+          throws IOException {
+        return exactSearch(context, DocIdSetIterator.all(totalDocs), null);
+      }
+    };
+  }
 }