update nextDoc

Signed-off-by: Vijayan Balasubramanian <[email protected]>
VijayanB · Sep 27, 2024 · 92040fd · 92040fd
1 parent 71adc4f
commit 92040fd
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 18 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIterator.java
@@ -60,9 +60,6 @@ public int nextDoc() throws IOException {
         if (docId == DocIdSetIterator.NO_MORE_DOCS) {
             return DocIdSetIterator.NO_MORE_DOCS;
         }
-        if (bitSetIterator != null) {
-            binaryVectorValues.advance(docId);
-        }
         currentScore = computeScore();
         int currentDocId = docId;
         docId = getNextDocId();
@@ -82,9 +79,14 @@ protected float computeScore() throws IOException {
     }
 
     protected int getNextDocId() throws IOException {
-        if (bitSetIterator != null) {
-            return bitSetIterator.nextDoc();
+        if (bitSetIterator == null) {
+            return binaryVectorValues.nextDoc();
+        }
+        int nextDocID = this.bitSetIterator.nextDoc();
+        // For filter case, advance vector values to corresponding doc id from filter bit set
+        if (nextDocID != DocIdSetIterator.NO_MORE_DOCS) {
+            binaryVectorValues.advance(nextDocID);
         }
-        return binaryVectorValues.nextDoc();
+        return nextDocID;
     }
 }
diff --git a/src/main/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIterator.java
@@ -64,9 +64,6 @@ public int nextDoc() throws IOException {
         // and parentBitSet will have [1,5]
         // Hence, we have to iterate till docId from knnVectorValues is less than parentId instead of till equal to parentId
         while (docId != DocIdSetIterator.NO_MORE_DOCS && docId < currentParent) {
-            if (bitSetIterator != null) {
-                binaryVectorValues.advance(docId);
-            }
             float score = computeScore();
             if (score > currentScore) {
                 bestChild = docId;

diff --git a/src/main/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIterator.java
@@ -76,9 +76,6 @@ public int nextDoc() throws IOException {
         // and parentBitSet will have [1,5]
         // Hence, we have to iterate till docId from knnVectorValues is less than parentId instead of till equal to parentId
         while (docId != DocIdSetIterator.NO_MORE_DOCS && docId < currentParent) {
-            if (bitSetIterator != null) {
-                knnFloatVectorValues.advance(docId);
-            }
             float score = computeScore();
             if (score > currentScore) {
                 bestChild = docId;

diff --git a/src/main/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIterator.java
@@ -77,9 +77,6 @@ public int nextDoc() throws IOException {
         if (docId == DocIdSetIterator.NO_MORE_DOCS) {
             return DocIdSetIterator.NO_MORE_DOCS;
         }
-        if (bitSetIterator != null) {
-            knnFloatVectorValues.advance(docId);
-        }
         currentScore = computeScore();
         int currentDocId = docId;
         docId = getNextDocId();
@@ -104,9 +101,14 @@ protected float computeScore() throws IOException {
     }
 
     protected int getNextDocId() throws IOException {
-        if (bitSetIterator != null) {
-            return bitSetIterator.nextDoc();
+        if (bitSetIterator == null) {
+            return knnFloatVectorValues.nextDoc();
+        }
+        int nextDocID = this.bitSetIterator.nextDoc();
+        // For filter case, advance vector values to corresponding doc id from filter bit set
+        if (nextDocID != DocIdSetIterator.NO_MORE_DOCS) {
+            knnFloatVectorValues.advance(nextDocID);
         }
-        return knnFloatVectorValues.nextDoc();
+        return nextDocID;
     }
 }