Change initial size of DocIdSetBuilder (#502) (#753)

Changes initial size of the docidsetbuilder used for iterating over results for k-NN queries. Originally, it was set to the maximum docid. This changes it to be the number of docs returned. Signed-off-by: John Mazanec <[email protected]> (cherry picked from commit 586958e) Co-authored-by: John Mazanec <[email protected]>
opensearch-project · Feb 7, 2023 · 92fdb44 · 92fdb44
1 parent 417e988
commit 92fdb44
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/main/java/org/opensearch/knn/index/KNNWeight.java b/src/main/java/org/opensearch/knn/index/KNNWeight.java
@@ -186,7 +186,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);