From 6f6dd566e4871ddb2d0b5829c08aab9e4af90774 Mon Sep 17 00:00:00 2001 From: Vijayan Balasubramanian Date: Fri, 27 Sep 2024 23:15:43 -0700 Subject: [PATCH] KNNIterators should support with and without filters (#2155) * Rename class names to represent both and filter and non filter use cases * Iterator should support with filters Update VectorIterator and NesterVector Iterator to iterate even if there is no filters provided to iterator. Currently this is used by exact search to score either topk docs or all docs when filter is provided by users. However, in future we will be allowing exact search even if there are no filters. Hence, decouple filter and make it option to support both cases. --------- Signed-off-by: Vijayan Balasubramanian --- CHANGELOG.md | 1 + .../knn/index/query/ExactSearcher.java | 47 +++++---- .../ByteVectorIdsKNNIterator.java} | 45 ++++++--- .../{filtered => iterators}/KNNIterator.java | 2 +- .../NestedByteVectorIdsKNNIterator.java} | 32 ++++-- .../NestedVectorIdsKNNIterator.java} | 37 ++++--- .../VectorIdsKNNIterator.java} | 51 ++++++---- .../FilteredIdsKNNByteIteratorTests.java | 50 ---------- .../filtered/FilteredIdsKNNIteratorTests.java | 54 ---------- .../ByteVectorIdsKNNIteratorTests.java | 97 ++++++++++++++++++ .../NestedByteVectorIdsKNNIteratorTests.java} | 36 ++++++- .../NestedVectorIdsKNNIteratorTests.java} | 46 +++++++-- .../iterators/VectorIdsKNNIteratorTests.java | 98 +++++++++++++++++++ 13 files changed, 404 insertions(+), 192 deletions(-) rename src/main/java/org/opensearch/knn/index/query/{filtered/FilteredIdsKNNByteIterator.java => iterators/ByteVectorIdsKNNIterator.java} (57%) rename src/main/java/org/opensearch/knn/index/query/{filtered => iterators}/KNNIterator.java (80%) rename src/main/java/org/opensearch/knn/index/query/{filtered/NestedFilteredIdsKNNByteIterator.java => iterators/NestedByteVectorIdsKNNIterator.java} (54%) rename src/main/java/org/opensearch/knn/index/query/{filtered/NestedFilteredIdsKNNIterator.java => iterators/NestedVectorIdsKNNIterator.java} (59%) rename src/main/java/org/opensearch/knn/index/query/{filtered/FilteredIdsKNNIterator.java => iterators/VectorIdsKNNIterator.java} (65%) delete mode 100644 src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNByteIteratorTests.java delete mode 100644 src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNIteratorTests.java create mode 100644 src/test/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIteratorTests.java rename src/test/java/org/opensearch/knn/index/query/{filtered/NestedFilteredIdsKNNByteIteratorTests.java => iterators/NestedByteVectorIdsKNNIteratorTests.java} (54%) rename src/test/java/org/opensearch/knn/index/query/{filtered/NestedFilteredIdsKNNIteratorTests.java => iterators/NestedVectorIdsKNNIteratorTests.java} (55%) create mode 100644 src/test/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIteratorTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 92f4f0bbe..67879bae7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), * Add short circuit if no live docs are in segments [#2059](https://github.com/opensearch-project/k-NN/pull/2059) * Optimize reduceToTopK in ResultUtil by removing pre-filling and reducing peek calls [#2146](https://github.com/opensearch-project/k-NN/pull/2146) * Update Default Rescore Context based on Dimension [#2149](https://github.com/opensearch-project/k-NN/pull/2149) +* KNNIterators should support with and without filters [#2155](https://github.com/opensearch-project/k-NN/pull/2155) ### Bug Fixes * KNN80DocValues should only be considered for BinaryDocValues fields [#2147](https://github.com/opensearch-project/k-NN/pull/2147) ### Infrastructure diff --git a/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java b/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java index 5b6029766..193cba8c1 100644 --- a/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java +++ b/src/main/java/org/opensearch/knn/index/query/ExactSearcher.java @@ -20,11 +20,11 @@ import org.opensearch.knn.common.FieldInfoExtractor; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.VectorDataType; -import org.opensearch.knn.index.query.filtered.FilteredIdsKNNByteIterator; -import org.opensearch.knn.index.query.filtered.FilteredIdsKNNIterator; -import org.opensearch.knn.index.query.filtered.KNNIterator; -import org.opensearch.knn.index.query.filtered.NestedFilteredIdsKNNByteIterator; -import org.opensearch.knn.index.query.filtered.NestedFilteredIdsKNNIterator; +import org.opensearch.knn.index.query.iterators.ByteVectorIdsKNNIterator; +import org.opensearch.knn.index.query.iterators.VectorIdsKNNIterator; +import org.opensearch.knn.index.query.iterators.KNNIterator; +import org.opensearch.knn.index.query.iterators.NestedByteVectorIdsKNNIterator; +import org.opensearch.knn.index.query.iterators.NestedVectorIdsKNNIterator; import org.opensearch.knn.index.vectorvalues.KNNBinaryVectorValues; import org.opensearch.knn.index.vectorvalues.KNNFloatVectorValues; import org.opensearch.knn.index.vectorvalues.KNNVectorValues; @@ -51,8 +51,9 @@ public class ExactSearcher { */ public Map searchLeaf(final LeafReaderContext leafReaderContext, final ExactSearcherContext exactSearcherContext) throws IOException { - KNNIterator iterator = getMatchedKNNIterator(leafReaderContext, exactSearcherContext); - if (exactSearcherContext.getMatchedDocs().cardinality() <= exactSearcherContext.getK()) { + KNNIterator iterator = getKNNIterator(leafReaderContext, exactSearcherContext); + if (exactSearcherContext.getMatchedDocs() != null + && exactSearcherContext.getMatchedDocs().cardinality() <= exactSearcherContext.getK()) { return scoreAllDocs(iterator); } return searchTopK(iterator, exactSearcherContext.getK()); @@ -98,8 +99,7 @@ private Map searchTopK(KNNIterator iterator, int k) throws IOExc return docToScore; } - private KNNIterator getMatchedKNNIterator(LeafReaderContext leafReaderContext, ExactSearcherContext exactSearcherContext) - throws IOException { + private KNNIterator getKNNIterator(LeafReaderContext leafReaderContext, ExactSearcherContext exactSearcherContext) throws IOException { final KNNQuery knnQuery = exactSearcherContext.getKnnQuery(); final BitSet matchedDocs = exactSearcherContext.getMatchedDocs(); final SegmentReader reader = Lucene.segmentReader(leafReaderContext.reader()); @@ -108,20 +108,18 @@ private KNNIterator getMatchedKNNIterator(LeafReaderContext leafReaderContext, E boolean isNestedRequired = exactSearcherContext.isParentHits() && knnQuery.getParentsFilter() != null; - if (VectorDataType.BINARY == knnQuery.getVectorDataType() && isNestedRequired) { - final KNNVectorValues vectorValues = KNNVectorValuesFactory.getVectorValues(fieldInfo, reader); - return new NestedFilteredIdsKNNByteIterator( - matchedDocs, - knnQuery.getByteQueryVector(), - (KNNBinaryVectorValues) vectorValues, - spaceType, - knnQuery.getParentsFilter().getBitSet(leafReaderContext) - ); - } - if (VectorDataType.BINARY == knnQuery.getVectorDataType()) { final KNNVectorValues vectorValues = KNNVectorValuesFactory.getVectorValues(fieldInfo, reader); - return new FilteredIdsKNNByteIterator( + if (isNestedRequired) { + return new NestedByteVectorIdsKNNIterator( + matchedDocs, + knnQuery.getByteQueryVector(), + (KNNBinaryVectorValues) vectorValues, + spaceType, + knnQuery.getParentsFilter().getBitSet(leafReaderContext) + ); + } + return new ByteVectorIdsKNNIterator( matchedDocs, knnQuery.getByteQueryVector(), (KNNBinaryVectorValues) vectorValues, @@ -142,7 +140,7 @@ private KNNIterator getMatchedKNNIterator(LeafReaderContext leafReaderContext, E final KNNVectorValues vectorValues = KNNVectorValuesFactory.getVectorValues(fieldInfo, reader); if (isNestedRequired) { - return new NestedFilteredIdsKNNIterator( + return new NestedVectorIdsKNNIterator( matchedDocs, knnQuery.getQueryVector(), (KNNFloatVectorValues) vectorValues, @@ -152,8 +150,7 @@ private KNNIterator getMatchedKNNIterator(LeafReaderContext leafReaderContext, E segmentLevelQuantizationInfo ); } - - return new FilteredIdsKNNIterator( + return new VectorIdsKNNIterator( matchedDocs, knnQuery.getQueryVector(), (KNNFloatVectorValues) vectorValues, @@ -180,7 +177,7 @@ public static class ExactSearcherContext { KNNQuery knnQuery; /** * whether the matchedDocs contains parent ids or child ids. This is relevant in the case of - * filtered nested search where the matchedDocs contain the parent ids and {@link NestedFilteredIdsKNNIterator} + * filtered nested search where the matchedDocs contain the parent ids and {@link NestedVectorIdsKNNIterator} * needs to be used. */ boolean isParentHits; diff --git a/src/main/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNByteIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIterator.java similarity index 57% rename from src/main/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNByteIterator.java rename to src/main/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIterator.java index ccfe626a0..b1aea4284 100644 --- a/src/main/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNByteIterator.java +++ b/src/main/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIterator.java @@ -3,11 +3,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.query.filtered; +package org.opensearch.knn.index.query.iterators; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.util.BitSet; import org.apache.lucene.util.BitSetIterator; +import org.opensearch.common.Nullable; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.vectorvalues.KNNBinaryVectorValues; @@ -17,11 +18,9 @@ * Inspired by DiversifyingChildrenFloatKnnVectorQuery in lucene * https://github.com/apache/lucene/blob/7b8aece125aabff2823626d5b939abf4747f63a7/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java#L162 * - * The class is used in KNNWeight to score filtered KNN field by iterating filterIdsArray. + * The class is used in KNNWeight to score all docs, but, it iterates over filterIdsArray if filter is provided */ -public class FilteredIdsKNNByteIterator implements KNNIterator { - // Array of doc ids to iterate - protected final BitSet filterIdsBitSet; +public class ByteVectorIdsKNNIterator implements KNNIterator { protected final BitSetIterator bitSetIterator; protected final byte[] queryVector; protected final KNNBinaryVectorValues binaryVectorValues; @@ -29,18 +28,24 @@ public class FilteredIdsKNNByteIterator implements KNNIterator { protected float currentScore = Float.NEGATIVE_INFINITY; protected int docId; - public FilteredIdsKNNByteIterator( - final BitSet filterIdsBitSet, + public ByteVectorIdsKNNIterator( + @Nullable final BitSet filterIdsBitSet, final byte[] queryVector, final KNNBinaryVectorValues binaryVectorValues, final SpaceType spaceType - ) { - this.filterIdsBitSet = filterIdsBitSet; - this.bitSetIterator = new BitSetIterator(filterIdsBitSet, filterIdsBitSet.length()); + ) throws IOException { + this.bitSetIterator = filterIdsBitSet == null ? null : new BitSetIterator(filterIdsBitSet, filterIdsBitSet.length()); this.queryVector = queryVector; this.binaryVectorValues = binaryVectorValues; this.spaceType = spaceType; - this.docId = bitSetIterator.nextDoc(); + // This cannot be moved inside nextDoc() method since it will break when we have nested field, where + // nextDoc should already be referring to next knnVectorValues + this.docId = getNextDocId(); + } + + public ByteVectorIdsKNNIterator(final byte[] queryVector, final KNNBinaryVectorValues binaryVectorValues, final SpaceType spaceType) + throws IOException { + this(null, queryVector, binaryVectorValues, spaceType); } /** @@ -55,10 +60,10 @@ public int nextDoc() throws IOException { if (docId == DocIdSetIterator.NO_MORE_DOCS) { return DocIdSetIterator.NO_MORE_DOCS; } - int doc = binaryVectorValues.advance(docId); currentScore = computeScore(); - docId = bitSetIterator.nextDoc(); - return doc; + int currentDocId = docId; + docId = getNextDocId(); + return currentDocId; } @Override @@ -72,4 +77,16 @@ protected float computeScore() throws IOException { // scores correspond to closer vectors. return spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector); } + + protected int getNextDocId() throws IOException { + 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 nextDocID; + } } diff --git a/src/main/java/org/opensearch/knn/index/query/filtered/KNNIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/KNNIterator.java similarity index 80% rename from src/main/java/org/opensearch/knn/index/query/filtered/KNNIterator.java rename to src/main/java/org/opensearch/knn/index/query/iterators/KNNIterator.java index 4a105975a..00cbb3aa2 100644 --- a/src/main/java/org/opensearch/knn/index/query/filtered/KNNIterator.java +++ b/src/main/java/org/opensearch/knn/index/query/iterators/KNNIterator.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.query.filtered; +package org.opensearch.knn.index.query.iterators; import java.io.IOException; diff --git a/src/main/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNByteIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIterator.java similarity index 54% rename from src/main/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNByteIterator.java rename to src/main/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIterator.java index b69a90518..3c93ec888 100644 --- a/src/main/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNByteIterator.java +++ b/src/main/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIterator.java @@ -3,33 +3,45 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.query.filtered; +package org.opensearch.knn.index.query.iterators; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.util.BitSet; +import org.opensearch.common.Nullable; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.vectorvalues.KNNBinaryVectorValues; import java.io.IOException; /** - * This iterator iterates filterIdsArray to score. However, it dedupe docs per each parent doc + * This iterator iterates filterIdsArray to score if filter is provided else it iterates over all docs. + * However, it dedupe docs per each parent doc * of which ID is set in parentBitSet and only return best child doc with the highest score. */ -public class NestedFilteredIdsKNNByteIterator extends FilteredIdsKNNByteIterator { +public class NestedByteVectorIdsKNNIterator extends ByteVectorIdsKNNIterator { private final BitSet parentBitSet; - public NestedFilteredIdsKNNByteIterator( - final BitSet filterIdsArray, + public NestedByteVectorIdsKNNIterator( + @Nullable final BitSet filterIdsArray, final byte[] queryVector, final KNNBinaryVectorValues binaryVectorValues, final SpaceType spaceType, final BitSet parentBitSet - ) { + ) throws IOException { super(filterIdsArray, queryVector, binaryVectorValues, spaceType); this.parentBitSet = parentBitSet; } + public NestedByteVectorIdsKNNIterator( + final byte[] queryVector, + final KNNBinaryVectorValues binaryVectorValues, + final SpaceType spaceType, + final BitSet parentBitSet + ) throws IOException { + super(null, queryVector, binaryVectorValues, spaceType); + this.parentBitSet = parentBitSet; + } + /** * Advance to the next best child doc per parent and update score with the best score among child docs from the parent. * DocIdSetIterator.NO_MORE_DOCS is returned when there is no more docs @@ -46,14 +58,18 @@ public int nextDoc() throws IOException { int currentParent = parentBitSet.nextSetBit(docId); int bestChild = -1; + // In order to traverse all children for given parent, we have to use docId < parentId, because, + // kNNVectorValues will not have parent id since DocId is unique per segment. For ex: let's say for doc id 1, there is one child + // and for doc id 5, there are three children. In that case knnVectorValues iterator will have [0, 2, 3, 4] + // 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) { - binaryVectorValues.advance(docId); float score = computeScore(); if (score > currentScore) { bestChild = docId; currentScore = score; } - docId = bitSetIterator.nextDoc(); + docId = getNextDocId(); } return bestChild; diff --git a/src/main/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIterator.java similarity index 59% rename from src/main/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNIterator.java rename to src/main/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIterator.java index 53ac72882..692793b99 100644 --- a/src/main/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNIterator.java +++ b/src/main/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIterator.java @@ -3,10 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.query.filtered; +package org.opensearch.knn.index.query.iterators; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.util.BitSet; +import org.opensearch.common.Nullable; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.query.SegmentLevelQuantizationInfo; import org.opensearch.knn.index.vectorvalues.KNNFloatVectorValues; @@ -14,31 +15,41 @@ import java.io.IOException; /** - * This iterator iterates filterIdsArray to score. However, it dedupe docs per each parent doc + * This iterator iterates filterIdsArray to score if filter is provided else it iterates over all docs. + * However, it dedupe docs per each parent doc * of which ID is set in parentBitSet and only return best child doc with the highest score. */ -public class NestedFilteredIdsKNNIterator extends FilteredIdsKNNIterator { +public class NestedVectorIdsKNNIterator extends VectorIdsKNNIterator { private final BitSet parentBitSet; - NestedFilteredIdsKNNIterator( - final BitSet filterIdsArray, + public NestedVectorIdsKNNIterator( + @Nullable final BitSet filterIdsArray, final float[] queryVector, final KNNFloatVectorValues knnFloatVectorValues, final SpaceType spaceType, final BitSet parentBitSet - ) { + ) throws IOException { this(filterIdsArray, queryVector, knnFloatVectorValues, spaceType, parentBitSet, null, null); } - public NestedFilteredIdsKNNIterator( - final BitSet filterIdsArray, + public NestedVectorIdsKNNIterator( + final float[] queryVector, + final KNNFloatVectorValues knnFloatVectorValues, + final SpaceType spaceType, + final BitSet parentBitSet + ) throws IOException { + this(null, queryVector, knnFloatVectorValues, spaceType, parentBitSet, null, null); + } + + public NestedVectorIdsKNNIterator( + @Nullable final BitSet filterIdsArray, final float[] queryVector, final KNNFloatVectorValues knnFloatVectorValues, final SpaceType spaceType, final BitSet parentBitSet, final byte[] quantizedVector, final SegmentLevelQuantizationInfo segmentLevelQuantizationInfo - ) { + ) throws IOException { super(filterIdsArray, queryVector, knnFloatVectorValues, spaceType, quantizedVector, segmentLevelQuantizationInfo); this.parentBitSet = parentBitSet; } @@ -59,14 +70,18 @@ public int nextDoc() throws IOException { int currentParent = parentBitSet.nextSetBit(docId); int bestChild = -1; + // In order to traverse all children for given parent, we have to use docId < parentId, because, + // kNNVectorValues will not have parent id since DocId is unique per segment. For ex: let's say for doc id 1, there is one child + // and for doc id 5, there are three children. In that case knnVectorValues iterator will have [0, 2, 3, 4] + // 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) { - knnFloatVectorValues.advance(docId); float score = computeScore(); if (score > currentScore) { bestChild = docId; currentScore = score; } - docId = bitSetIterator.nextDoc(); + docId = getNextDocId(); } return bestChild; diff --git a/src/main/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNIterator.java b/src/main/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIterator.java similarity index 65% rename from src/main/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNIterator.java rename to src/main/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIterator.java index 56d291470..9fb354242 100644 --- a/src/main/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNIterator.java +++ b/src/main/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIterator.java @@ -3,11 +3,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.query.filtered; +package org.opensearch.knn.index.query.iterators; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.util.BitSet; import org.apache.lucene.util.BitSetIterator; +import org.opensearch.common.Nullable; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.query.SegmentLevelQuantizationInfo; import org.opensearch.knn.index.query.SegmentLevelQuantizationUtil; @@ -19,11 +20,9 @@ * Inspired by DiversifyingChildrenFloatKnnVectorQuery in lucene * https://github.com/apache/lucene/blob/7b8aece125aabff2823626d5b939abf4747f63a7/lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java#L162 * - * The class is used in KNNWeight to score filtered KNN field by iterating filterIdsArray. + * The class is used in KNNWeight to score all docs, but, it iterates over filterIdsArray if filter is provided */ -public class FilteredIdsKNNIterator implements KNNIterator { - // Array of doc ids to iterate - protected final BitSet filterIdsBitSet; +public class VectorIdsKNNIterator implements KNNIterator { protected final BitSetIterator bitSetIterator; protected final float[] queryVector; private final byte[] quantizedQueryVector; @@ -33,29 +32,35 @@ public class FilteredIdsKNNIterator implements KNNIterator { protected int docId; private final SegmentLevelQuantizationInfo segmentLevelQuantizationInfo; - FilteredIdsKNNIterator( - final BitSet filterIdsBitSet, + public VectorIdsKNNIterator( + @Nullable final BitSet filterIdsBitSet, final float[] queryVector, final KNNFloatVectorValues knnFloatVectorValues, final SpaceType spaceType - ) { + ) throws IOException { this(filterIdsBitSet, queryVector, knnFloatVectorValues, spaceType, null, null); } - public FilteredIdsKNNIterator( - final BitSet filterIdsBitSet, + public VectorIdsKNNIterator(final float[] queryVector, final KNNFloatVectorValues knnFloatVectorValues, final SpaceType spaceType) + throws IOException { + this(null, queryVector, knnFloatVectorValues, spaceType, null, null); + } + + public VectorIdsKNNIterator( + @Nullable final BitSet filterIdsBitSet, final float[] queryVector, final KNNFloatVectorValues knnFloatVectorValues, final SpaceType spaceType, final byte[] quantizedQueryVector, final SegmentLevelQuantizationInfo segmentLevelQuantizationInfo - ) { - this.filterIdsBitSet = filterIdsBitSet; - this.bitSetIterator = new BitSetIterator(filterIdsBitSet, filterIdsBitSet.length()); + ) throws IOException { + this.bitSetIterator = filterIdsBitSet == null ? null : new BitSetIterator(filterIdsBitSet, filterIdsBitSet.length()); this.queryVector = queryVector; this.knnFloatVectorValues = knnFloatVectorValues; this.spaceType = spaceType; - this.docId = bitSetIterator.nextDoc(); + // This cannot be moved inside nextDoc() method since it will break when we have nested field, where + // nextDoc should already be referring to next knnVectorValues + this.docId = getNextDocId(); this.quantizedQueryVector = quantizedQueryVector; this.segmentLevelQuantizationInfo = segmentLevelQuantizationInfo; } @@ -72,10 +77,10 @@ public int nextDoc() throws IOException { if (docId == DocIdSetIterator.NO_MORE_DOCS) { return DocIdSetIterator.NO_MORE_DOCS; } - int doc = knnFloatVectorValues.advance(docId); currentScore = computeScore(); - docId = bitSetIterator.nextDoc(); - return doc; + int currentDocId = docId; + docId = getNextDocId(); + return currentDocId; } @Override @@ -94,4 +99,16 @@ protected float computeScore() throws IOException { return spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector); } } + + protected int getNextDocId() throws IOException { + 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 nextDocID; + } } diff --git a/src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNByteIteratorTests.java b/src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNByteIteratorTests.java deleted file mode 100644 index c52798c05..000000000 --- a/src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNByteIteratorTests.java +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.knn.index.query.filtered; - -import junit.framework.TestCase; -import lombok.SneakyThrows; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.util.FixedBitSet; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.knn.index.vectorvalues.KNNBinaryVectorValues; - -import java.util.Arrays; -import java.util.List; -import java.util.stream.Collectors; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class FilteredIdsKNNByteIteratorTests extends TestCase { - @SneakyThrows - public void testNextDoc_whenCalled_IterateAllDocs() { - final SpaceType spaceType = SpaceType.HAMMING; - final byte[] queryVector = { 1, 2, 3 }; - final int[] filterIds = { 1, 2, 3 }; - final List dataVectors = Arrays.asList(new byte[] { 11, 12, 13 }, new byte[] { 14, 15, 16 }, new byte[] { 17, 18, 19 }); - final List expectedScores = dataVectors.stream() - .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) - .collect(Collectors.toList()); - - KNNBinaryVectorValues values = mock(KNNBinaryVectorValues.class); - when(values.getVector()).thenReturn(dataVectors.get(0), dataVectors.get(1), dataVectors.get(2)); - - FixedBitSet filterBitSet = new FixedBitSet(4); - for (int id : filterIds) { - when(values.advance(id)).thenReturn(id); - filterBitSet.set(id); - } - - // Execute and verify - FilteredIdsKNNByteIterator iterator = new FilteredIdsKNNByteIterator(filterBitSet, queryVector, values, spaceType); - for (int i = 0; i < filterIds.length; i++) { - assertEquals(filterIds[i], iterator.nextDoc()); - assertEquals(expectedScores.get(i), (Float) iterator.score()); - } - assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); - } -} diff --git a/src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNIteratorTests.java b/src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNIteratorTests.java deleted file mode 100644 index 731eed2cc..000000000 --- a/src/test/java/org/opensearch/knn/index/query/filtered/FilteredIdsKNNIteratorTests.java +++ /dev/null @@ -1,54 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.knn.index.query.filtered; - -import lombok.SneakyThrows; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.util.FixedBitSet; -import org.opensearch.knn.KNNTestCase; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.knn.index.vectorvalues.KNNFloatVectorValues; - -import java.util.Arrays; -import java.util.List; -import java.util.stream.Collectors; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class FilteredIdsKNNIteratorTests extends KNNTestCase { - @SneakyThrows - public void testNextDoc_whenCalled_IterateAllDocs() { - final SpaceType spaceType = SpaceType.L2; - final float[] queryVector = { 1.0f, 2.0f, 3.0f }; - final int[] filterIds = { 1, 2, 3 }; - final List dataVectors = Arrays.asList( - new float[] { 11.0f, 12.0f, 13.0f }, - new float[] { 14.0f, 15.0f, 16.0f }, - new float[] { 17.0f, 18.0f, 19.0f } - ); - final List expectedScores = dataVectors.stream() - .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) - .collect(Collectors.toList()); - - KNNFloatVectorValues values = mock(KNNFloatVectorValues.class); - when(values.getVector()).thenReturn(dataVectors.get(0), dataVectors.get(1), dataVectors.get(2)); - - FixedBitSet filterBitSet = new FixedBitSet(4); - for (int id : filterIds) { - when(values.advance(id)).thenReturn(id); - filterBitSet.set(id); - } - - // Execute and verify - FilteredIdsKNNIterator iterator = new FilteredIdsKNNIterator(filterBitSet, queryVector, values, spaceType); - for (int i = 0; i < filterIds.length; i++) { - assertEquals(filterIds[i], iterator.nextDoc()); - assertEquals(expectedScores.get(i), (Float) iterator.score()); - } - assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); - } -} diff --git a/src/test/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIteratorTests.java b/src/test/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIteratorTests.java new file mode 100644 index 000000000..0b1b71286 --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/query/iterators/ByteVectorIdsKNNIteratorTests.java @@ -0,0 +1,97 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.query.iterators; + +import junit.framework.TestCase; +import lombok.SneakyThrows; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.util.FixedBitSet; +import org.mockito.stubbing.OngoingStubbing; +import org.opensearch.knn.index.SpaceType; +import org.opensearch.knn.index.vectorvalues.KNNBinaryVectorValues; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; + +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ByteVectorIdsKNNIteratorTests extends TestCase { + @SneakyThrows + public void testNextDoc_whenCalled_thenIterateAllDocs() { + final SpaceType spaceType = SpaceType.HAMMING; + final byte[] queryVector = { 1, 2, 3 }; + final int[] filterIds = { 1, 2, 3 }; + final List dataVectors = Arrays.asList(new byte[] { 11, 12, 13 }, new byte[] { 14, 15, 16 }, new byte[] { 17, 18, 19 }); + final List expectedScores = dataVectors.stream() + .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) + .collect(Collectors.toList()); + + KNNBinaryVectorValues values = mock(KNNBinaryVectorValues.class); + when(values.getVector()).thenReturn(dataVectors.get(0), dataVectors.get(1), dataVectors.get(2)); + + FixedBitSet filterBitSet = new FixedBitSet(4); + for (int id : filterIds) { + when(values.advance(id)).thenReturn(id); + filterBitSet.set(id); + } + + // Execute and verify + ByteVectorIdsKNNIterator iterator = new ByteVectorIdsKNNIterator(filterBitSet, queryVector, values, spaceType); + for (int i = 0; i < filterIds.length; i++) { + assertEquals(filterIds[i], iterator.nextDoc()); + assertEquals(expectedScores.get(i), iterator.score()); + } + assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); + } + + @SneakyThrows + public void testNextDoc_whenCalled_thenIterateAllDocsWithoutFilter() throws IOException { + final SpaceType spaceType = SpaceType.HAMMING; + final byte[] queryVector = { 1, 2, 3 }; + final List dataVectors = Arrays.asList( + new byte[] { 11, 12, 13 }, + new byte[] { 14, 15, 16 }, + new byte[] { 17, 18, 19 }, + new byte[] { 20, 21, 22 }, + new byte[] { 23, 24, 25 } + ); + final List expectedScores = dataVectors.stream() + .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) + .collect(Collectors.toList()); + + KNNBinaryVectorValues values = mock(KNNBinaryVectorValues.class); + when(values.getVector()).thenReturn( + dataVectors.get(0), + dataVectors.get(1), + dataVectors.get(2), + dataVectors.get(3), + dataVectors.get(4) + ); + + // stub return value when nextDoc is called + OngoingStubbing stubbing = when(values.nextDoc()); + for (int i = 0; i < dataVectors.size(); i++) { + stubbing = stubbing.thenReturn(i); + } + // set last return to be Integer.MAX_VALUE to represent no more docs + stubbing.thenReturn(Integer.MAX_VALUE); + + // Execute and verify + ByteVectorIdsKNNIterator iterator = new ByteVectorIdsKNNIterator(queryVector, values, spaceType); + for (int i = 0; i < dataVectors.size(); i++) { + assertEquals(i, iterator.nextDoc()); + assertEquals(expectedScores.get(i), iterator.score()); + } + assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); + verify(values, never()).advance(anyInt()); + } +} diff --git a/src/test/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNByteIteratorTests.java b/src/test/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIteratorTests.java similarity index 54% rename from src/test/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNByteIteratorTests.java rename to src/test/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIteratorTests.java index 1940ffe12..eff021234 100644 --- a/src/test/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNByteIteratorTests.java +++ b/src/test/java/org/opensearch/knn/index/query/iterators/NestedByteVectorIdsKNNIteratorTests.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.query.filtered; +package org.opensearch.knn.index.query.iterators; import junit.framework.TestCase; import lombok.SneakyThrows; @@ -17,10 +17,13 @@ import java.util.List; import java.util.stream.Collectors; +import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -public class NestedFilteredIdsKNNByteIteratorTests extends TestCase { +public class NestedByteVectorIdsKNNIteratorTests extends TestCase { @SneakyThrows public void testNextDoc_whenIterate_ReturnBestChildDocsPerParent() { final SpaceType spaceType = SpaceType.HAMMING; @@ -45,7 +48,7 @@ public void testNextDoc_whenIterate_ReturnBestChildDocsPerParent() { } // Execute and verify - NestedFilteredIdsKNNByteIterator iterator = new NestedFilteredIdsKNNByteIterator( + NestedByteVectorIdsKNNIterator iterator = new NestedByteVectorIdsKNNIterator( filterBitSet, queryVector, values, @@ -58,4 +61,31 @@ public void testNextDoc_whenIterate_ReturnBestChildDocsPerParent() { assertEquals(expectedScores.get(2), iterator.score()); assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); } + + @SneakyThrows + public void testNextDoc_whenIterateWithoutFilters_thenReturnBestChildDocsPerParent() { + final SpaceType spaceType = SpaceType.HAMMING; + final byte[] queryVector = { 1, 2, 3 }; + // Parent id for 0 -> 1 + // Parent id for 2, 3 -> 4 + // In bit representation, it is 10010. In long, it is 18. + final BitSet parentBitSet = new FixedBitSet(new long[] { 18 }, 5); + final List dataVectors = Arrays.asList(new byte[] { 11, 12, 13 }, new byte[] { 14, 15, 16 }, new byte[] { 17, 18, 19 }); + final List expectedScores = dataVectors.stream() + .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) + .collect(Collectors.toList()); + + KNNBinaryVectorValues values = mock(KNNBinaryVectorValues.class); + when(values.getVector()).thenReturn(dataVectors.get(0), dataVectors.get(1), dataVectors.get(2)); + when(values.nextDoc()).thenReturn(0, 2, 3, Integer.MAX_VALUE); + + // Execute and verify + NestedByteVectorIdsKNNIterator iterator = new NestedByteVectorIdsKNNIterator(queryVector, values, spaceType, parentBitSet); + assertEquals(0, iterator.nextDoc()); + assertEquals(expectedScores.get(0), iterator.score()); + assertEquals(3, iterator.nextDoc()); + assertEquals(expectedScores.get(2), iterator.score()); + assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); + verify(values, never()).advance(anyInt()); + } } diff --git a/src/test/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNIteratorTests.java b/src/test/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIteratorTests.java similarity index 55% rename from src/test/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNIteratorTests.java rename to src/test/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIteratorTests.java index cca789a4d..f94ddb4e1 100644 --- a/src/test/java/org/opensearch/knn/index/query/filtered/NestedFilteredIdsKNNIteratorTests.java +++ b/src/test/java/org/opensearch/knn/index/query/iterators/NestedVectorIdsKNNIteratorTests.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.query.filtered; +package org.opensearch.knn.index.query.iterators; import junit.framework.TestCase; import lombok.SneakyThrows; @@ -17,10 +17,13 @@ import java.util.List; import java.util.stream.Collectors; +import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -public class NestedFilteredIdsKNNIteratorTests extends TestCase { +public class NestedVectorIdsKNNIteratorTests extends TestCase { @SneakyThrows public void testNextDoc_whenIterate_ReturnBestChildDocsPerParent() { final SpaceType spaceType = SpaceType.L2; @@ -53,17 +56,42 @@ public void testNextDoc_whenIterate_ReturnBestChildDocsPerParent() { } // Execute and verify - NestedFilteredIdsKNNIterator iterator = new NestedFilteredIdsKNNIterator( - filterBitSet, - queryVector, - values, - spaceType, - parentBitSet - ); + NestedVectorIdsKNNIterator iterator = new NestedVectorIdsKNNIterator(filterBitSet, queryVector, values, spaceType, parentBitSet); assertEquals(filterIds[0], iterator.nextDoc()); assertEquals(expectedScores.get(0), iterator.score()); assertEquals(filterIds[2], iterator.nextDoc()); assertEquals(expectedScores.get(2), iterator.score()); assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); } + + @SneakyThrows + public void testNextDoc_whenIterateWithoutFilters_thenReturnBestChildDocsPerParent() { + final SpaceType spaceType = SpaceType.L2; + final float[] queryVector = { 1.0f, 2.0f, 3.0f }; + // Parent id for 0 -> 1 + // Parent id for 2, 3 -> 4 + // In bit representation, it is 10010. In long, it is 18. + final BitSet parentBitSet = new FixedBitSet(new long[] { 18 }, 5); + final List dataVectors = Arrays.asList( + new float[] { 11.0f, 12.0f, 13.0f }, + new float[] { 17.0f, 18.0f, 19.0f }, + new float[] { 14.0f, 15.0f, 16.0f } + ); + final List expectedScores = dataVectors.stream() + .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) + .collect(Collectors.toList()); + + KNNFloatVectorValues values = mock(KNNFloatVectorValues.class); + when(values.getVector()).thenReturn(dataVectors.get(0), dataVectors.get(1), dataVectors.get(2)); + when(values.nextDoc()).thenReturn(0, 2, 3, Integer.MAX_VALUE); + + // Execute and verify + NestedVectorIdsKNNIterator iterator = new NestedVectorIdsKNNIterator(queryVector, values, spaceType, parentBitSet); + assertEquals(0, iterator.nextDoc()); + assertEquals(expectedScores.get(0), iterator.score()); + assertEquals(3, iterator.nextDoc()); + assertEquals(expectedScores.get(2), iterator.score()); + assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); + verify(values, never()).advance(anyInt()); + } } diff --git a/src/test/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIteratorTests.java b/src/test/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIteratorTests.java new file mode 100644 index 000000000..96932d0f1 --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/query/iterators/VectorIdsKNNIteratorTests.java @@ -0,0 +1,98 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.query.iterators; + +import lombok.SneakyThrows; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.util.FixedBitSet; +import org.mockito.stubbing.OngoingStubbing; +import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.SpaceType; +import org.opensearch.knn.index.vectorvalues.KNNFloatVectorValues; + +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; + +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class VectorIdsKNNIteratorTests extends KNNTestCase { + @SneakyThrows + public void testNextDoc_whenCalledWithFilters_thenIterateAllDocs() { + final SpaceType spaceType = SpaceType.L2; + final float[] queryVector = { 1.0f, 2.0f, 3.0f }; + final int[] filterIds = { 1, 2, 3 }; + final List dataVectors = Arrays.asList( + new float[] { 11.0f, 12.0f, 13.0f }, + new float[] { 14.0f, 15.0f, 16.0f }, + new float[] { 17.0f, 18.0f, 19.0f } + ); + final List expectedScores = dataVectors.stream() + .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) + .collect(Collectors.toList()); + + KNNFloatVectorValues values = mock(KNNFloatVectorValues.class); + when(values.getVector()).thenReturn(dataVectors.get(0), dataVectors.get(1), dataVectors.get(2)); + + FixedBitSet filterBitSet = new FixedBitSet(4); + for (int id : filterIds) { + when(values.advance(id)).thenReturn(id); + filterBitSet.set(id); + } + + // Execute and verify + VectorIdsKNNIterator iterator = new VectorIdsKNNIterator(filterBitSet, queryVector, values, spaceType); + for (int i = 0; i < filterIds.length; i++) { + assertEquals(filterIds[i], iterator.nextDoc()); + assertEquals(expectedScores.get(i), (Float) iterator.score()); + } + assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); + } + + @SneakyThrows + public void testNextDoc_whenCalledWithoutFilters_thenIterateAllDocs() { + final SpaceType spaceType = SpaceType.L2; + final float[] queryVector = { 1.0f, 2.0f, 3.0f }; + final List dataVectors = Arrays.asList( + new float[] { 11.0f, 12.0f, 13.0f }, + new float[] { 14.0f, 15.0f, 16.0f }, + new float[] { 17.0f, 18.0f, 19.0f }, + new float[] { 20.0f, 21.0f, 22.0f }, + new float[] { 23.0f, 24.0f, 25.0f } + ); + final List expectedScores = dataVectors.stream() + .map(vector -> spaceType.getKnnVectorSimilarityFunction().compare(queryVector, vector)) + .collect(Collectors.toList()); + + KNNFloatVectorValues values = mock(KNNFloatVectorValues.class); + when(values.getVector()).thenReturn( + dataVectors.get(0), + dataVectors.get(1), + dataVectors.get(2), + dataVectors.get(3), + dataVectors.get(4) + ); + // stub return value when nextDoc is called + OngoingStubbing stubbing = when(values.nextDoc()); + for (int i = 0; i < dataVectors.size(); i++) { + stubbing = stubbing.thenReturn(i); + } + // set last return to be Integer.MAX_VALUE to represent no more docs + stubbing.thenReturn(Integer.MAX_VALUE); + // Execute and verify + VectorIdsKNNIterator iterator = new VectorIdsKNNIterator(queryVector, values, spaceType); + for (int i = 0; i < dataVectors.size(); i++) { + assertEquals(i, iterator.nextDoc()); + assertEquals(expectedScores.get(i), (Float) iterator.score()); + } + assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc()); + verify(values, never()).advance(anyInt()); + } +}