Fix test failure TestKnnByteVectorQuery#testTimeLimitingKnnCollectorM…

…anager (#13312)
apache · Apr 18, 2024 · 821fc42 · 821fc42
1 parent 94ced43
commit 821fc42
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 50 deletions.
diff --git a/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java
@@ -42,54 +42,62 @@ public KnnCollector newCollector(int visitedLimit, LeafReaderContext context) th
     if (queryTimeout == null) {
       return collector;
     }
-    return new KnnCollector() {
-      @Override
-      public boolean earlyTerminated() {
-        return queryTimeout.shouldExit() || collector.earlyTerminated();
-      }
-
-      @Override
-      public void incVisitedCount(int count) {
-        collector.incVisitedCount(count);
-      }
-
-      @Override
-      public long visitedCount() {
-        return collector.visitedCount();
-      }
-
-      @Override
-      public long visitLimit() {
-        return collector.visitLimit();
-      }
-
-      @Override
-      public int k() {
-        return collector.k();
-      }
-
-      @Override
-      public boolean collect(int docId, float similarity) {
-        return collector.collect(docId, similarity);
-      }
-
-      @Override
-      public float minCompetitiveSimilarity() {
-        return collector.minCompetitiveSimilarity();
-      }
-
-      @Override
-      public TopDocs topDocs() {
-        TopDocs docs = collector.topDocs();
-
-        // Mark results as partial if timeout is met
-        TotalHits.Relation relation =
-            queryTimeout.shouldExit()
-                ? TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO
-                : docs.totalHits.relation;
-
-        return new TopDocs(new TotalHits(docs.totalHits.value, relation), docs.scoreDocs);
-      }
-    };
+    return new TimeLimitingKnnCollector(collector);
+  }
+
+  class TimeLimitingKnnCollector implements KnnCollector {
+    private final KnnCollector collector;
+
+    TimeLimitingKnnCollector(KnnCollector collector) {
+      this.collector = collector;
+    }
+
+    @Override
+    public boolean earlyTerminated() {
+      return queryTimeout.shouldExit() || collector.earlyTerminated();
+    }
+
+    @Override
+    public void incVisitedCount(int count) {
+      collector.incVisitedCount(count);
+    }
+
+    @Override
+    public long visitedCount() {
+      return collector.visitedCount();
+    }
+
+    @Override
+    public long visitLimit() {
+      return collector.visitLimit();
+    }
+
+    @Override
+    public int k() {
+      return collector.k();
+    }
+
+    @Override
+    public boolean collect(int docId, float similarity) {
+      return collector.collect(docId, similarity);
+    }
+
+    @Override
+    public float minCompetitiveSimilarity() {
+      return collector.minCompetitiveSimilarity();
+    }
+
+    @Override
+    public TopDocs topDocs() {
+      TopDocs docs = collector.topDocs();
+
+      // Mark results as partial if timeout is met
+      TotalHits.Relation relation =
+          queryTimeout.shouldExit()
+              ? TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO
+              : docs.totalHits.relation;
+
+      return new TopDocs(new TotalHits(docs.totalHits.value, relation), docs.scoreDocs);
+    }
   }
 }
diff --git a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
@@ -816,7 +816,8 @@ public void testTimeLimitingKnnCollectorManager() throws IOException {
           noTimeoutManager.newCollector(Integer.MAX_VALUE, searcher.leafContexts.get(0));
 
       // Check that a normal collector is created without timeout
-      assertTrue(noTimeoutCollector instanceof TopKnnCollector);
+      assertFalse(
+          noTimeoutCollector instanceof TimeLimitingKnnCollectorManager.TimeLimitingKnnCollector);
       noTimeoutCollector.collect(0, 0);
       assertFalse(noTimeoutCollector.earlyTerminated());