Specialise the wrapping of vector scorer iterator for KnnVectorValues.DocIndexIterator in ExitableDirectoryReader (#131623)

This commit specialises the wrapping of vector scorer iterator for KnnVectorValues.DocIndexIterator in ExitableDirectoryReader.

This change is necessary now since a new assert in Lucene catches this scenario, since DocIndexIterator is expected in several places which we were previously not exposing. DocIndexIterator can be optimized more than DocIdSetIterator.

closes #131610

Author: ChrisHegarty

server/src/main/java/org/elasticsearch/search/internal/ExitableDirectoryReader.java

@@ -586,8 +586,9 @@ public VectorScorer scorer(byte[] bytes) throws IOException {
             if (scorer == null) {
                 return null;
             }
+            DocIdSetIterator scorerIterator = scorer.iterator();
             return new VectorScorer() {
-                private final DocIdSetIterator iterator = new ExitableDocSetIterator(scorer.iterator(), queryCancellation);
+                private final DocIdSetIterator iterator = exitableIterator(scorerIterator, queryCancellation);
 
                 @Override
                 public float score() throws IOException {
@@ -637,8 +638,9 @@ public VectorScorer scorer(float[] target) throws IOException {
             if (scorer == null) {
                 return null;
             }
+            DocIdSetIterator scorerIterator = scorer.iterator();
             return new VectorScorer() {
-                private final DocIdSetIterator iterator = new ExitableDocSetIterator(scorer.iterator(), queryCancellation);
+                private final DocIdSetIterator iterator = exitableIterator(scorerIterator, queryCancellation);
 
                 @Override
                 public float score() throws IOException {
@@ -663,6 +665,15 @@ public FloatVectorValues copy() throws IOException {
         }
     }
 
+    /** Wraps the iterator in an exitable iterator, specializing for KnnVectorValues.DocIndexIterator. */
+    static DocIdSetIterator exitableIterator(DocIdSetIterator iterator, QueryCancellation queryCancellation) {
+        if (iterator instanceof KnnVectorValues.DocIndexIterator docIndexIterator) {
+            return createExitableIterator(docIndexIterator, queryCancellation);
+        } else {
+            return new ExitableDocSetIterator(iterator, queryCancellation);
+        }
+    }
+
     private static KnnVectorValues.DocIndexIterator createExitableIterator(
         KnnVectorValues.DocIndexIterator delegate,
         QueryCancellation queryCancellation