From e0c3afe5fe2f0048b9422364aeeb620e35a44deb Mon Sep 17 00:00:00 2001 From: Junqiu Lei Date: Fri, 27 Sep 2024 14:47:55 -0700 Subject: [PATCH] Optimize reduceToTopK in ResultUtil by removing pre-filling and reducing peek calls (#2146) Signed-off-by: Junqiu Lei --- CHANGELOG.md | 1 + .../java/org/opensearch/knn/index/query/ResultUtil.java | 7 +++---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a270ec9c0..92f4f0bbe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), * Add AVX512 support to k-NN for FAISS library [#2069](https://github.com/opensearch-project/k-NN/pull/2069) ### Enhancements * 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) ### Bug Fixes * KNN80DocValues should only be considered for BinaryDocValues fields [#2147](https://github.com/opensearch-project/k-NN/pull/2147) diff --git a/src/main/java/org/opensearch/knn/index/query/ResultUtil.java b/src/main/java/org/opensearch/knn/index/query/ResultUtil.java index 2ed70b8b3..f62c09cb0 100644 --- a/src/main/java/org/opensearch/knn/index/query/ResultUtil.java +++ b/src/main/java/org/opensearch/knn/index/query/ResultUtil.java @@ -33,15 +33,14 @@ public final class ResultUtil { public static void reduceToTopK(List> perLeafResults, int k) { // Iterate over all scores to get min competitive score PriorityQueue topKMinQueue = new PriorityQueue<>(k); - for (int i = 0; i < k; i++) { - topKMinQueue.add(-Float.MAX_VALUE); - } int count = 0; for (Map perLeafResult : perLeafResults) { count += perLeafResult.size(); for (Float score : perLeafResult.values()) { - if (topKMinQueue.peek() != null && score > topKMinQueue.peek()) { + if (topKMinQueue.size() < k) { + topKMinQueue.add(score); + } else if (topKMinQueue.peek() != null && score > topKMinQueue.peek()) { topKMinQueue.poll(); topKMinQueue.add(score); }