Some refactoring after code review

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

CHANGELOG.md

@@ -16,7 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Features
 ### Enhancements
 ### Bug Fixes
-- Fixed merge logic for multiple collector result case ([#800](https://github.com/opensearch-project/neural-search/pull/800)) 
+- Fix for missing HybridQuery results when concurrent segment search is enabled ([#800](https://github.com/opensearch-project/neural-search/pull/800))
 ### Infrastructure
 ### Documentation
 ### Maintenance

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java

@@ -4,7 +4,6 @@
  */
 package org.opensearch.neuralsearch.search.query;
 
-import com.google.common.annotations.VisibleForTesting;
 import lombok.RequiredArgsConstructor;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.Collector;
@@ -20,7 +19,7 @@
 import org.opensearch.common.lucene.search.TopDocsAndMaxScore;
 import org.opensearch.neuralsearch.search.HitsThresholdChecker;
 import org.opensearch.neuralsearch.search.HybridTopScoreDocCollector;
-import org.opensearch.neuralsearch.search.util.ScoreDocsMerger;
+import org.opensearch.neuralsearch.search.util.HybridQueryScoreDocsMerger;
 import org.opensearch.neuralsearch.search.util.TopDocsMerger;
 import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.internal.ContextIndexSearcher;
@@ -34,7 +33,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 
@@ -56,7 +54,7 @@ public abstract class HybridCollectorManager implements CollectorManager<Collect
     @Nullable
     private final Weight filterWeight;
     private static final float boost_factor = 1f;
-    private final ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+    private final HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
     private final TopDocsMerger topDocsMerger = new TopDocsMerger(scoreDocsMerger);
 
     /**

src/main/java/org/opensearch/neuralsearch/search/util/ScoreDocsMerger.java → src/main/java/org/opensearch/neuralsearch/search/util/HybridQueryScoreDocsMerger.java

@@ -18,12 +18,14 @@
  * Merges two ScoreDoc arrays into one
  */
 @NoArgsConstructor
-public class ScoreDocsMerger<T extends ScoreDoc> {
+public class HybridQueryScoreDocsMerger<T extends ScoreDoc> {
 
     private static final int MIN_NUMBER_OF_ELEMENTS_IN_SCORE_DOC = 3;
 
     /**
      * Merge two score docs objects, result ScoreDocs[] object will have all hits per sub-query from both original objects.
+     * Input and output ScoreDocs are in format that is specific to Hybrid Query. This method should not be used for ScoreDocs from
+     * other query types.
      * Logic is based on assumption that hits of every sub-query are sorted by score.
      * Method returns new object and doesn't mutate original ScoreDocs arrays.
      * @param sourceScoreDocs original score docs from query result

src/main/java/org/opensearch/neuralsearch/search/util/TopDocsMerger.java

@@ -20,7 +20,7 @@
 @RequiredArgsConstructor
 public class TopDocsMerger {
 
-    private final ScoreDocsMerger<ScoreDoc> scoreDocsMerger;
+    private final HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger;
     @VisibleForTesting
     protected static final Comparator<ScoreDoc> SCORE_DOC_BY_SCORE_COMPARATOR = Comparator.comparing((scoreDoc) -> scoreDoc.score);
 

src/test/java/org/opensearch/neuralsearch/search/util/ScoreDocsMergerTests.java → src/test/java/org/opensearch/neuralsearch/search/util/HybridQueryScoreDocsMergerTests.java

@@ -13,12 +13,12 @@
 import static org.opensearch.neuralsearch.search.util.HybridSearchResultFormatUtil.MAGIC_NUMBER_START_STOP;
 import static org.opensearch.neuralsearch.search.util.HybridSearchResultFormatUtil.MAGIC_NUMBER_DELIMITER;
 
-public class ScoreDocsMergerTests extends OpenSearchQueryTestCase {
+public class HybridQueryScoreDocsMergerTests extends OpenSearchQueryTestCase {
 
     private static final float DELTA_FOR_ASSERTION = 0.001f;
 
     public void testIncorrectInput_whenScoreDocsAreNullOrNotEnoughElements_thenFail() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
 
         ScoreDoc[] scores = new ScoreDoc[] {
             createStartStopElementForHybridSearchResults(2),
@@ -51,7 +51,7 @@ public void testIncorrectInput_whenScoreDocsAreNullOrNotEnoughElements_thenFail(
     }
 
     public void testMergeScoreDocs_whenBothTopDocsHasHits_thenSuccessful() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
 
         ScoreDoc[] scoreDocsOriginal = new ScoreDoc[] {
             createStartStopElementForHybridSearchResults(0),
@@ -90,7 +90,7 @@ public void testMergeScoreDocs_whenBothTopDocsHasHits_thenSuccessful() {
     }
 
     public void testMergeScoreDocs_whenOneTopDocsHasHitsAndOtherIsEmpty_thenSuccessful() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
 
         ScoreDoc[] scoreDocsOriginal = new ScoreDoc[] {
             createStartStopElementForHybridSearchResults(0),
@@ -123,7 +123,7 @@ public void testMergeScoreDocs_whenOneTopDocsHasHitsAndOtherIsEmpty_thenSuccessf
     }
 
     public void testMergeScoreDocs_whenBothTopDocsHasNoHits_thenSuccessful() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
 
         ScoreDoc[] scoreDocsOriginal = new ScoreDoc[] {
             createStartStopElementForHybridSearchResults(0),

src/test/java/org/opensearch/neuralsearch/search/util/TopDocsMergerTests.java

@@ -22,7 +22,7 @@ public class TopDocsMergerTests extends OpenSearchQueryTestCase {
 
     @SneakyThrows
     public void testMergeScoreDocs_whenBothTopDocsHasHits_thenSuccessful() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
         TopDocsMerger topDocsMerger = new TopDocsMerger(scoreDocsMerger);
 
         TopDocs topDocsOriginal = new TopDocs(
@@ -78,7 +78,7 @@ public void testMergeScoreDocs_whenBothTopDocsHasHits_thenSuccessful() {
 
     @SneakyThrows
     public void testMergeScoreDocs_whenOneTopDocsHasHitsAndOtherIsEmpty_thenSuccessful() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
         TopDocsMerger topDocsMerger = new TopDocsMerger(scoreDocsMerger);
 
         TopDocs topDocsOriginal = new TopDocs(
@@ -130,7 +130,7 @@ public void testMergeScoreDocs_whenOneTopDocsHasHitsAndOtherIsEmpty_thenSuccessf
 
     @SneakyThrows
     public void testMergeScoreDocs_whenBothTopDocsHasNoHits_thenSuccessful() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
         TopDocsMerger topDocsMerger = new TopDocsMerger(scoreDocsMerger);
 
         TopDocs topDocsOriginal = new TopDocs(
@@ -172,7 +172,7 @@ public void testMergeScoreDocs_whenBothTopDocsHasNoHits_thenSuccessful() {
 
     @SneakyThrows
     public void testThreeSequentialMerges_whenAllTopDocsHasHits_thenSuccessful() {
-        ScoreDocsMerger<ScoreDoc> scoreDocsMerger = new ScoreDocsMerger<>();
+        HybridQueryScoreDocsMerger<ScoreDoc> scoreDocsMerger = new HybridQueryScoreDocsMerger<>();
         TopDocsMerger topDocsMerger = new TopDocsMerger(scoreDocsMerger);
 
         TopDocs topDocsOriginal = new TopDocs(