From e3e54c95c93b7da84c14e2efdd4714cd310a16ba Mon Sep 17 00:00:00 2001
From: Michael Sokolov <sokolov@falutin.net>
Date: Wed, 1 Sep 2021 08:19:11 -0400
Subject: [PATCH] LUCENE-10063: test fixes relating to
 SimpleTextKnnVectorsReader (#273)

---
 .../codecs/simpletext/SimpleTextKnnVectorsReader.java  | 10 ++++------
 .../lucene/index/BaseKnnVectorsFormatTestCase.java     |  5 +++--
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java
index 3a384c53231..a3b9436aae1 100644
--- a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java
+++ b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java
@@ -148,18 +148,16 @@ public class SimpleTextKnnVectorsReader extends KnnVectorsReader {
     VectorValues values = getVectorValues(field);
     if (target.length != values.dimension()) {
       throw new IllegalArgumentException(
-          "incorrect dimension for field "
-              + field
-              + "; expected "
-              + values.dimension()
-              + " but target has "
-              + target.length);
+          "vector dimensions differ: " + target.length + "!=" + values.dimension());
     }
     FieldInfo info = readState.fieldInfos.fieldInfo(field);
     VectorSimilarityFunction vectorSimilarity = info.getVectorSimilarityFunction();
     HitQueue topK = new HitQueue(k, false);
     int doc;
     while ((doc = values.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
+      if (acceptDocs != null && acceptDocs.get(doc) == false) {
+        continue;
+      }
       float[] vector = values.vectorValue();
       float score = vectorSimilarity.compare(vector, target);
       if (vectorSimilarity.reversed) {
diff --git a/lucene/test-framework/src/java/org/apache/lucene/index/BaseKnnVectorsFormatTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/index/BaseKnnVectorsFormatTestCase.java
index 4efc537bef2..22fc2e1ba3e 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/index/BaseKnnVectorsFormatTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/index/BaseKnnVectorsFormatTestCase.java
@@ -859,10 +859,11 @@ public abstract class BaseKnnVectorsFormatTestCase extends BaseIndexFileFormatTe
           }
           // assert that searchNearestVectors returns the expected number of documents, in
           // descending score order
-          int k = random().nextInt(numDoc / 2);
+          int size = ctx.reader().getVectorValues(fieldName).size();
+          int k = random().nextInt(size / 2 + 1);
           TopDocs results =
               ctx.reader().searchNearestVectors(fieldName, randomVector(dimension), k, liveDocs);
-          assertEquals(k, results.scoreDocs.length);
+          assertEquals(Math.min(k, size), results.scoreDocs.length);
           for (int i = 0; i < k - 1; i++) {
             assertTrue(results.scoreDocs[i].score >= results.scoreDocs[i + 1].score);
           }