Addressed review comments, mainly refactoring

Signed-off-by: Martin Gaievski <[email protected]>

Author: martin-gaievski

src/main/java/org/opensearch/neuralsearch/query/HybridBulkScorer.java

@@ -36,6 +36,7 @@ public class HybridBulkScorer extends BulkScorer {
     private final float[][] windowScores;
     private final HybridQueryDocIdStream hybridQueryDocIdStream;
     private final int maxDoc;
+    private int[] docIds;
 
     /**
      * Constructor for HybridBulkScorer
@@ -45,8 +46,9 @@ public class HybridBulkScorer extends BulkScorer {
      */
     public HybridBulkScorer(List<Scorer> scorers, boolean needsScores, int maxDoc) {
         long cost = 0;
-        this.scorers = new Scorer[scorers.size()];
-        for (int subQueryIndex = 0; subQueryIndex < scorers.size(); subQueryIndex++) {
+        int numOfQueries = scorers.size();
+        this.scorers = new Scorer[numOfQueries];
+        for (int subQueryIndex = 0; subQueryIndex < numOfQueries; subQueryIndex++) {
             Scorer scorer = scorers.get(subQueryIndex);
             if (Objects.isNull(scorer)) {
                 continue;
@@ -55,12 +57,14 @@ public HybridBulkScorer(List<Scorer> scorers, boolean needsScores, int maxDoc) {
             this.scorers[subQueryIndex] = scorer;
         }
         this.cost = cost;
-        this.hybridSubQueryScorer = new HybridSubQueryScorer(scorers.size());
+        this.hybridSubQueryScorer = new HybridSubQueryScorer(numOfQueries);
         this.needsScores = needsScores;
         this.matching = new FixedBitSet(WINDOW_SIZE);
         this.windowScores = new float[this.scorers.length][WINDOW_SIZE];
         this.maxDoc = maxDoc;
         this.hybridQueryDocIdStream = new HybridQueryDocIdStream(this);
+        this.docIds = new int[numOfQueries];
+        Arrays.fill(docIds, DocIdSetIterator.NO_MORE_DOCS);
     }
 
     @Override
@@ -69,15 +73,15 @@ public int score(LeafCollector collector, Bits acceptDocs, int min, int max) thr
         // making sure we are not going over the global limit defined by maxDoc
         max = Math.min(max, maxDoc);
         // advance all scorers to the segment's minimum doc id
-        int[] docsIds = advance(min, scorers);
-        while (allDocIdsUsed(docsIds, max) == false) {
-            scoreWindow(collector, acceptDocs, min, max, docsIds);
+        advance(min, scorers);
+        while (allDocIdsUsed(docIds, max) == false) {
+            scoreWindow(collector, acceptDocs, min, max, docIds);
         }
-        return getNextDocIdCandidate(docsIds);
+        return getNextDocIdCandidate(docIds);
     }
 
     private void scoreWindow(LeafCollector collector, Bits acceptDocs, int min, int max, int[] docIds) throws IOException {
-        // pick the lowest out of all not yet used doc ids
+        // find the first document ID below the maximum threshold to establish the next scoring window boundary
         int topDoc = -1;
         for (int docId : docIds) {
             if (docId < max) {
@@ -150,11 +154,9 @@ private void scoreWindowIntoBitSetWithSubqueryScorers(
     /**
      * Advance all scorers to the next document that is >= min
      */
-    private int[] advance(int min, Scorer[] scorers) throws IOException {
-        int[] docIds = new int[scorers.length];
+    private void advance(int min, Scorer[] scorers) throws IOException {
         for (int subQueryIndex = 0; subQueryIndex < scorers.length; subQueryIndex++) {
             if (Objects.isNull(scorers[subQueryIndex])) {
-                docIds[subQueryIndex] = DocIdSetIterator.NO_MORE_DOCS;
                 continue;
             }
             DocIdSetIterator it = scorers[subQueryIndex].iterator();
@@ -164,7 +166,6 @@ private int[] advance(int min, Scorer[] scorers) throws IOException {
             }
             docIds[subQueryIndex] = doc;
         }
-        return docIds;
     }
 
     private boolean allDocIdsUsed(int[] docsIds, int max) {

src/main/java/org/opensearch/neuralsearch/query/HybridQueryDocIdStream.java

@@ -18,6 +18,7 @@
  */
 @RequiredArgsConstructor
 public class HybridQueryDocIdStream extends DocIdStream {
+    private static final int BLOCK_SHIFT = 6;
     private final HybridBulkScorer hybridBulkScorer;
     @Setter
     private int base;
@@ -41,7 +42,7 @@ public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException
                 // calculate actual document ID within the window
                 // idx << 6 is equivalent to idx * 64 (block offset)
                 // numberOfTrailingZeros gives position within the block
-                final int docIndexInWindow = (idx << 6) | numberOfTrailingZeros;
+                final int docIndexInWindow = (idx << BLOCK_SHIFT) | numberOfTrailingZeros;
                 float[][] windowScores = hybridBulkScorer.getWindowScores();
                 for (int subQueryIndex = 0; subQueryIndex < windowScores.length; subQueryIndex++) {
                     if (Objects.isNull(windowScores[subQueryIndex])) {

src/main/java/org/opensearch/neuralsearch/search/collector/HybridTopScoreDocCollector.java

@@ -23,6 +23,7 @@
 import org.apache.lucene.util.PriorityQueue;
 
 import lombok.extern.log4j.Log4j2;
+import org.opensearch.neuralsearch.query.HybridSubQueryScorer;
 import org.opensearch.neuralsearch.search.HitsThresholdChecker;
 
 /**
@@ -132,22 +133,18 @@ public void setScorer(Scorable scorer) throws IOException {
 
         @Override
         public void collect(int doc) throws IOException {
-            if (Objects.isNull(getCompoundQueryScorer())) {
+            HybridSubQueryScorer compoundQueryScorer = getCompoundQueryScorer();
+            if (Objects.isNull(compoundQueryScorer)) {
                 return;
             }
             ensureSubQueryScoreQueues();
             // Increment total hit count which represents unique doc found on the shard
             totalHits++;
-            float[] scores = getCompoundQueryScorer().getSubQueryScores();
+            float[] scores = compoundQueryScorer.getSubQueryScores();
             int docWithBase = doc + docBase;
             for (int subQueryIndex = 0; subQueryIndex < scores.length; subQueryIndex++) {
                 float score = scores[subQueryIndex];
-                // if score is 0.0 there is no hits for that sub-query
-                if (score <= 0) {
-                    continue;
-                }
-
-                if (score < minScoreThresholds[subQueryIndex]) {
+                if (isNonCompetitiveScore(score, subQueryIndex)) {
                     continue;
                 }
 
@@ -172,6 +169,10 @@ public void collect(int doc) throws IOException {
             }
         }
 
+        private boolean isNonCompetitiveScore(float score, int subQueryIndex) {
+            return score <= 0 && score < minScoreThresholds[subQueryIndex];
+        }
+
         /**
          * Initialize compound score queues for each sub query if it's not initialized already
          */