Refactoring after review comments

Signed-off-by: Martin Gaievski <[email protected]>
martin-gaievski · Jan 3, 2025 · f3f177c · f3f177c
1 parent 40a8e36
commit f3f177c
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/...ain/java/org/opensearch/neuralsearch/search/collector/HybridTopFieldDocSortCollector.java b/...ain/java/org/opensearch/neuralsearch/search/collector/HybridTopFieldDocSortCollector.java
@@ -47,7 +47,7 @@ public abstract class HybridTopFieldDocSortCollector implements HybridSearchColl
     @Nullable
     private FieldDoc after;
     private FieldComparator<?> firstComparator;
-    // bottom would be set to null per shard.
+    // the array stores bottom elements of the min heap of sorted hits for each sub query
     @Getter(AccessLevel.PACKAGE)
     @VisibleForTesting
     private FieldValueHitQueue.Entry fieldValueLeafTrackers[];
@@ -70,6 +70,7 @@ public abstract class HybridTopFieldDocSortCollector implements HybridSearchColl
     @Getter
     protected float maxScore = 0.0f;
     protected int[] collectedHits;
+    private boolean needsInitialization = true;
 
     // searchSortPartOfIndexSort is used to evaluate whether to perform index sort or not.
     private Boolean searchSortPartOfIndexSort = null;
@@ -250,17 +251,16 @@ protected boolean thresholdCheck(int doc, int subQueryNumber) throws IOException
         The method initializes once per search request.
          */
         protected void initializePriorityQueuesWithComparators(LeafReaderContext context, int numberOfSubQueries) throws IOException {
-            if (compoundScores == null) {
+            if (needsInitialization) {
                 compoundScores = new FieldValueHitQueue[numberOfSubQueries];
                 comparators = new LeafFieldComparator[numberOfSubQueries];
                 queueFull = new boolean[numberOfSubQueries];
                 collectedHits = new int[numberOfSubQueries];
                 for (int i = 0; i < numberOfSubQueries; i++) {
                     initializeLeafFieldComparators(context, i);
                 }
-            }
-            if (Objects.isNull(fieldValueLeafTrackers)) {
                 fieldValueLeafTrackers = new FieldValueHitQueue.Entry[numberOfSubQueries];
+                needsInitialization = false;
             }
             if (initializeLeafComparatorsPerSegmentOnce) {
                 for (int i = 0; i < numberOfSubQueries; i++) {