diff --git a/src/main/java/org/opensearch/knn/index/query/KNNWeight.java b/src/main/java/org/opensearch/knn/index/query/KNNWeight.java index c7f2ddf47..716aed412 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNWeight.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNWeight.java @@ -182,7 +182,10 @@ public Scorer scorer(LeafReaderContext context) throws IOException { .collect(Collectors.toMap(KNNQueryResult::getId, result -> knnEngine.score(result.getScore(), spaceType))); int maxDoc = Collections.max(scores.keySet()) + 1; DocIdSetBuilder docIdSetBuilder = new DocIdSetBuilder(maxDoc); - DocIdSetBuilder.BulkAdder setAdder = docIdSetBuilder.grow(maxDoc); + + // The docIdSetIterator will contain the docids of the returned results. So, before adding results to + // the builder, we can grow to results.length + DocIdSetBuilder.BulkAdder setAdder = docIdSetBuilder.grow(results.length); Arrays.stream(results).forEach(result -> setAdder.add(result.getId())); DocIdSetIterator docIdSetIter = docIdSetBuilder.build().iterator(); return new KNNScorer(this, docIdSetIter, scores, boost);