Refactoring to optimize extractHybridQuery method calls

Signed-off-by: Varun Jain <[email protected]>
opensearch-project · Jan 8, 2025 · 62498ad · 62498ad
1 parent d39a1b8
commit 62498ad
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 17 deletions.
diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java
@@ -59,7 +59,7 @@ public final class HybridQueryBuilder extends AbstractQueryBuilder<HybridQueryBu
     private int paginationDepth;
     static final int MAX_NUMBER_OF_SUB_QUERIES = 5;
     private final static int DEFAULT_PAGINATION_DEPTH = 10;
-    private static final int LOWER_BOUND_OF_PAGINATION_DEPTH = 1;
+    private static final int LOWER_BOUND_OF_PAGINATION_DEPTH = 0;
 
     public HybridQueryBuilder(StreamInput in) throws IOException {
         super(in);
@@ -327,8 +327,10 @@ private Collection<Query> toQueries(Collection<QueryBuilder> queryBuilders, Quer
     }
 
     private static void validatePaginationDepth(final int paginationDepth, final QueryShardContext queryShardContext) {
-        if (paginationDepth < LOWER_BOUND_OF_PAGINATION_DEPTH) {
-            throw new IllegalArgumentException(String.format(Locale.ROOT, "pagination_depth should be greater than 0"));
+        if (paginationDepth <= LOWER_BOUND_OF_PAGINATION_DEPTH) {
+            throw new IllegalArgumentException(
+                String.format(Locale.ROOT, "pagination_depth should be greater than %s", LOWER_BOUND_OF_PAGINATION_DEPTH)
+            );
         }
         // compare pagination depth with OpenSearch setting index.max_result_window
         // see https://opensearch.org/docs/latest/install-and-configure/configuring-opensearch/index-settings/

diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java
@@ -54,6 +54,7 @@
 import static org.opensearch.neuralsearch.search.util.HybridSearchResultFormatUtil.createFieldDocStartStopElementForHybridSearchResults;
 import static org.opensearch.neuralsearch.search.util.HybridSearchResultFormatUtil.createFieldDocDelimiterElementForHybridSearchResults;
 import static org.opensearch.neuralsearch.search.util.HybridSearchResultFormatUtil.createSortFieldsForDelimiterResults;
+import static org.opensearch.neuralsearch.util.HybridQueryUtil.isHybridQueryWrappedInBooleanQuery;
 
 /**
  * Collector manager based on HybridTopScoreDocCollector that allows users to parallelize counting the number of hits.
@@ -483,7 +484,7 @@ private ReduceableSearchResult reduceSearchResults(final List<ReduceableSearchRe
      * @return results size to collected
      */
     private static int getSubqueryResultsRetrievalSize(final SearchContext searchContext) {
-        HybridQuery hybridQuery = unwrapHybridQuery(searchContext.query());
+        HybridQuery hybridQuery = unwrapHybridQuery(searchContext);
         int paginationDepth = hybridQuery.getQueryContext().getPaginationDepth();
 
         // Switch to from+size retrieval size during standard hybrid query execution.
@@ -496,10 +497,11 @@ private static int getSubqueryResultsRetrievalSize(final SearchContext searchCon
     /**
      * Unwraps a HybridQuery from either a direct query or a nested BooleanQuery
      */
-    private static HybridQuery unwrapHybridQuery(Query query) {
+    private static HybridQuery unwrapHybridQuery(final SearchContext searchContext) {
         HybridQuery hybridQuery;
+        Query query = searchContext.query();
         // In case of nested fields and alias filter, hybrid query is wrapped under bool query and lies in the first clause.
-        if (query instanceof BooleanQuery) {
+        if (isHybridQueryWrappedInBooleanQuery(searchContext, searchContext.query())) {
             BooleanQuery booleanQuery = (BooleanQuery) query;
             hybridQuery = (HybridQuery) booleanQuery.clauses().get(0).getQuery();
         } else {

diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java
@@ -29,9 +29,8 @@
 
 import lombok.extern.log4j.Log4j2;
 
-import static org.opensearch.neuralsearch.util.HybridQueryUtil.hasAliasFilter;
-import static org.opensearch.neuralsearch.util.HybridQueryUtil.hasNestedFieldOrNestedDocs;
 import static org.opensearch.neuralsearch.util.HybridQueryUtil.isHybridQuery;
+import static org.opensearch.neuralsearch.util.HybridQueryUtil.isHybridQueryWrappedInBooleanQuery;
 
 /**
  * Custom search implementation to be used at {@link QueryPhase} for Hybrid Query search. For queries other than Hybrid the
@@ -74,16 +73,9 @@ private QueryPhaseSearcher getQueryPhaseSearcher(final SearchContext searchConte
             : defaultQueryPhaseSearcherWithEmptyCollectorContext;
     }
 
-    private static boolean isWrappedHybridQuery(final Query query) {
-        return query instanceof BooleanQuery
-            && ((BooleanQuery) query).clauses().stream().anyMatch(clauseQuery -> clauseQuery.getQuery() instanceof HybridQuery);
-    }
-
     @VisibleForTesting
     protected Query extractHybridQuery(final SearchContext searchContext, final Query query) {
-        if ((hasAliasFilter(query, searchContext) || hasNestedFieldOrNestedDocs(query, searchContext))
-            && isWrappedHybridQuery(query)
-            && !((BooleanQuery) query).clauses().isEmpty()) {
+        if (isHybridQueryWrappedInBooleanQuery(searchContext, query)) {
             List<BooleanClause> booleanClauses = ((BooleanQuery) query).clauses();
             if (!(booleanClauses.get(0).getQuery() instanceof HybridQuery)) {
                 throw new IllegalStateException("cannot process hybrid query due to incorrect structure of top level query");

diff --git a/src/main/java/org/opensearch/neuralsearch/util/HybridQueryUtil.java b/src/main/java/org/opensearch/neuralsearch/util/HybridQueryUtil.java
@@ -56,12 +56,18 @@ public static boolean hasNestedFieldOrNestedDocs(final Query query, final Search
         return searchContext.mapperService().hasNested() && new NestedHelper(searchContext.mapperService()).mightMatchNestedDocs(query);
     }
 
-    private static boolean isWrappedHybridQuery(final Query query) {
+    public static boolean isWrappedHybridQuery(final Query query) {
         return query instanceof BooleanQuery
             && ((BooleanQuery) query).clauses().stream().anyMatch(clauseQuery -> clauseQuery.getQuery() instanceof HybridQuery);
     }
 
     public static boolean hasAliasFilter(final Query query, final SearchContext searchContext) {
         return Objects.nonNull(searchContext.aliasFilter());
     }
+
+    public static boolean isHybridQueryWrappedInBooleanQuery(final SearchContext searchContext, final Query query) {
+        return ((hasAliasFilter(query, searchContext) || hasNestedFieldOrNestedDocs(query, searchContext))
+            && isWrappedHybridQuery(query)
+            && !((BooleanQuery) query).clauses().isEmpty());
+    }
 }