adding validation in depth

Signed-off-by: Varun Jain <[email protected]>
vibrantvarun · Oct 16, 2024 · f1c64dc · f1c64dc
1 parent 58cd4c1
commit f1c64dc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java
@@ -79,7 +79,13 @@ public static CollectorManager createHybridCollectorManager(final SearchContext
         final IndexReader reader = searchContext.searcher().getIndexReader();
         final int totalNumDocs = Math.max(0, reader.numDocs());
         HybridQuery hybridQuery = (HybridQuery) searchContext.query();
-        int numDocs = Math.min(hybridQuery.getPaginationDepth(), totalNumDocs);
+        int retrievalSize;
+        if (hybridQuery.getPaginationDepth() == 0) {
+            retrievalSize = searchContext.from() + searchContext.size();
+        } else {
+            retrievalSize = hybridQuery.getPaginationDepth();
+        }
+        int numDocs = Math.min(retrievalSize, totalNumDocs);
         int trackTotalHitsUpTo = searchContext.trackTotalHitsUpTo();
         if (searchContext.sort() != null) {
             validateSortCriteria(searchContext, searchContext.trackScores());

diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java
@@ -214,5 +214,6 @@ protected boolean searchWithCollector(
                 hasTimeout
             );
         }
+
     }
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -214,5 +214,6 @@ protected boolean searchWithCollector( @@
                     hasTimeout
                 );
             }
         }
     }