Refactor implementations of query phase searcher, add empty QueryColl…

…ectorContext (opensearch-project#13481) * Refactor implementations of query hpase searcher by adding overloaded searchWith method Signed-off-by: Martin Gaievski <[email protected]> * Switched to Empty context add rescoring interface Signed-off-by: Martin Gaievski <[email protected]> * Changed if by simple null check for querycontext argument Signed-off-by: Martin Gaievski <[email protected]> * Added Override annotation for searchWith method Signed-off-by: Martin Gaievski <[email protected]> * Remove old override method from Concurrent Query Phase Searcher Signed-off-by: Martin Gaievski <[email protected]> --------- Signed-off-by: Martin Gaievski <[email protected]>
finnegancarroll · May 10, 2024 · 0e753a9 · 0e753a9
1 parent 2afa058
commit 0e753a9
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -71,6 +71,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [BWC and API enforcement] Reconsider the breaking changes check policy to detect breaking changes against released versions ([#13292](https://github.com/opensearch-project/OpenSearch/pull/13292))
 - Switch to macos-13 runner for precommit and assemble github actions due to macos-latest is now arm64 ([#13412](https://github.com/opensearch-project/OpenSearch/pull/13412))
 - [Revert] Prevent unnecessary fetch sub phase processor initialization during fetch phase execution ([#12503](https://github.com/opensearch-project/OpenSearch/pull/12503))
+- Refactor implementations of query phase searcher, allow QueryCollectorContext to have zero collectors ([#13481](https://github.com/opensearch-project/OpenSearch/pull/13481))
 
 ### Deprecated
 

diff --git a/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java b/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java
@@ -23,10 +23,9 @@
 
 import java.io.IOException;
 import java.util.LinkedList;
+import java.util.Objects;
 import java.util.concurrent.ExecutionException;
 
-import static org.opensearch.search.query.TopDocsCollectorContext.createTopDocsCollectorContext;
-
 /**
  * The implementation of the {@link QueryPhaseSearcher} which attempts to use concurrent
  * search of Apache Lucene segments if it has been enabled.
@@ -46,24 +45,32 @@ protected boolean searchWithCollector(
         ContextIndexSearcher searcher,
         Query query,
         LinkedList<QueryCollectorContext> collectors,
+        QueryCollectorContext queryCollectorContext,
         boolean hasFilterCollector,
         boolean hasTimeout
     ) throws IOException {
-        return searchWithCollectorManager(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
+        return searchWithCollectorManager(
+            searchContext,
+            searcher,
+            query,
+            collectors,
+            queryCollectorContext,
+            hasFilterCollector,
+            hasTimeout
+        );
     }
 
     private static boolean searchWithCollectorManager(
         SearchContext searchContext,
         ContextIndexSearcher searcher,
         Query query,
         LinkedList<QueryCollectorContext> collectorContexts,
+        QueryCollectorContext queryCollectorContext,
         boolean hasFilterCollector,
         boolean timeoutSet
     ) throws IOException {
-        // create the top docs collector last when the other collectors are known
-        final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
-        // add the top docs collector, the first collector context in the chain
-        collectorContexts.addFirst(topDocsFactory);
+        // add the passed collector, the first collector context in the chain
+        collectorContexts.addFirst(Objects.requireNonNull(queryCollectorContext));
 
         final QuerySearchResult queryResult = searchContext.queryResult();
         final CollectorManager<?, ReduceableSearchResult> collectorManager;
@@ -95,7 +102,10 @@ private static boolean searchWithCollectorManager(
             queryResult.terminatedEarly(false);
         }
 
-        return topDocsFactory.shouldRescore();
+        if (queryCollectorContext instanceof RescoringQueryCollectorContext) {
+            return ((RescoringQueryCollectorContext) queryCollectorContext).shouldRescore();
+        }
+        return false;
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java b/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java
@@ -77,6 +77,29 @@ public ScoreMode scoreMode() {
         }
     };
 
+    public static final QueryCollectorContext EMPTY_CONTEXT = new QueryCollectorContext("empty") {
+
+        @Override
+        Collector create(Collector in) throws IOException {
+            return EMPTY_COLLECTOR;
+        }
+
+        @Override
+        CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
+            return new CollectorManager<Collector, ReduceableSearchResult>() {
+                @Override
+                public Collector newCollector() throws IOException {
+                    return EMPTY_COLLECTOR;
+                }
+
+                @Override
+                public ReduceableSearchResult reduce(Collection<Collector> collectors) throws IOException {
+                    return result -> {};
+                }
+            };
+        }
+    };
+
     private String profilerName;
 
     QueryCollectorContext(String profilerName) {

diff --git a/server/src/main/java/org/opensearch/search/query/QueryPhase.java b/server/src/main/java/org/opensearch/search/query/QueryPhase.java
@@ -335,13 +335,12 @@ private static boolean searchWithCollector(
         ContextIndexSearcher searcher,
         Query query,
         LinkedList<QueryCollectorContext> collectors,
+        QueryCollectorContext queryCollectorContext,
         boolean hasFilterCollector,
         boolean timeoutSet
     ) throws IOException {
-        // create the top docs collector last when the other collectors are known
-        final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
-        // add the top docs collector, the first collector context in the chain
-        collectors.addFirst(topDocsFactory);
+        // add passed collector, the first collector context in the chain
+        collectors.addFirst(Objects.requireNonNull(queryCollectorContext));
 
         final Collector queryCollector;
         if (searchContext.getProfilers() != null) {
@@ -370,7 +369,10 @@ private static boolean searchWithCollector(
         for (QueryCollectorContext ctx : collectors) {
             ctx.postProcess(queryResult);
         }
-        return topDocsFactory.shouldRescore();
+        if (queryCollectorContext instanceof RescoringQueryCollectorContext) {
+            return ((RescoringQueryCollectorContext) queryCollectorContext).shouldRescore();
+        }
+        return false;
     }
 
     /**
@@ -440,7 +442,29 @@ protected boolean searchWithCollector(
             boolean hasFilterCollector,
             boolean hasTimeout
         ) throws IOException {
-            return QueryPhase.searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
+            // create the top docs collector last when the other collectors are known
+            final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
+            return searchWithCollector(searchContext, searcher, query, collectors, topDocsFactory, hasFilterCollector, hasTimeout);
+        }
+
+        protected boolean searchWithCollector(
+            SearchContext searchContext,
+            ContextIndexSearcher searcher,
+            Query query,
+            LinkedList<QueryCollectorContext> collectors,
+            QueryCollectorContext queryCollectorContext,
+            boolean hasFilterCollector,
+            boolean hasTimeout
+        ) throws IOException {
+            return QueryPhase.searchWithCollector(
+                searchContext,
+                searcher,
+                query,
+                collectors,
+                queryCollectorContext,
+                hasFilterCollector,
+                hasTimeout
+            );
         }
     }
 }
diff --git a/server/src/main/java/org/opensearch/search/query/RescoringQueryCollectorContext.java b/server/src/main/java/org/opensearch/search/query/RescoringQueryCollectorContext.java
@@ -0,0 +1,27 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.query;
+
+import org.opensearch.common.annotation.PublicApi;
+
+/**
+ * Abstraction that allows indication of whether results should be rescored or not based on
+ * custom logic of exact {@link QueryCollectorContext} implementation.
+ *
+ * @opensearch.api
+ */
+@PublicApi(since = "2.15.0")
+public interface RescoringQueryCollectorContext {
+
+    /**
+     * Indicates if results from the query context should be rescored
+     * @return true if results must be rescored, false otherwise
+     */
+    boolean shouldRescore();
+}
diff --git a/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java b/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java
@@ -95,7 +95,7 @@
  *
  * @opensearch.internal
  */
-public abstract class TopDocsCollectorContext extends QueryCollectorContext {
+public abstract class TopDocsCollectorContext extends QueryCollectorContext implements RescoringQueryCollectorContext {
     protected final int numHits;
 
     TopDocsCollectorContext(String profilerName, int numHits) {