From 92fdb44c019ea5f6c83e87b94c764f3ffe620c5b Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Tue, 7 Feb 2023 12:25:32 -0800 Subject: [PATCH] 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 (cherry picked from commit 586958e3ee2f99da85a6ec28ea4565c9ec572208) Co-authored-by: John Mazanec --- src/main/java/org/opensearch/knn/index/KNNWeight.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/main/java/org/opensearch/knn/index/KNNWeight.java b/src/main/java/org/opensearch/knn/index/KNNWeight.java index 909e7222f2..a6e0cd3506 100644 --- 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);