#17593 Add Extension Points For Pre and Post Collection of Scores in QueryPhase

CHANGELOG.md

@@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Workload Management] Modify logging message in WorkloadGroupService ([#18712](https://github.com/opensearch-project/OpenSearch/pull/18712))
 - Add BooleanQuery rewrite moving constant-scoring must clauses to filter clauses ([#18510](https://github.com/opensearch-project/OpenSearch/issues/18510))
 - Add functionality for plugins to inject QueryCollectorContext during QueryPhase ([#18637](https://github.com/opensearch-project/OpenSearch/pull/18637))
+- Add QueryPhaseExtension interface for pre/post score collection hooks ([#17593](https://github.com/opensearch-project/OpenSearch/issues/17593))
 - Add support for non-timing info in profiler ([#18460](https://github.com/opensearch-project/OpenSearch/issues/18460))
 - [Rule-based auto tagging] Bug fix and improvements ([#18726](https://github.com/opensearch-project/OpenSearch/pull/18726))
 - Extend Approximation Framework to other numeric types ([#18530](https://github.com/opensearch-project/OpenSearch/issues/18530))

server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java

@@ -39,6 +39,29 @@ public class ConcurrentQueryPhaseSearcher extends DefaultQueryPhaseSearcher {
      */
     public ConcurrentQueryPhaseSearcher() {}
 
+    @Override
+    public boolean searchWith(
+        SearchContext searchContext,
+        ContextIndexSearcher searcher,
+        Query query,
+        LinkedList<QueryCollectorContext> collectors,
+        boolean hasFilterCollector,
+        boolean hasTimeout
+    ) throws IOException {
+        return executeWithExtensions(searchContext, () -> {
+            QueryCollectorContext queryCollectorContext = getQueryCollectorContext(searchContext, hasFilterCollector);
+            return searchWithCollectorManager(
+                searchContext,
+                searcher,
+                query,
+                collectors,
+                queryCollectorContext,
+                hasFilterCollector,
+                hasTimeout
+            );
+        });
+    }
+
     @Override
     protected boolean searchWithCollector(
         SearchContext searchContext,

server/src/main/java/org/opensearch/search/query/QueryPhase.java

@@ -34,6 +34,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.BooleanClause;
@@ -430,7 +431,57 @@ public boolean searchWith(
             boolean hasFilterCollector,
             boolean hasTimeout
         ) throws IOException {
-            return searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
+            return executeWithExtensions(
+                searchContext,
+                () -> searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout)
+            );
+        }
+
+        /**
+         * Executes the provided search operation with before/after extension hooks.
+         * This method handles extension execution and error handling consistently.
+         */
+        protected boolean executeWithExtensions(SearchContext searchContext, SearchOperation searchOperation) throws IOException {
+            List<QueryPhaseExtension> extensions = queryPhaseExtensions();
+
+            // Execute beforeScoreCollection extensions
+            for (QueryPhaseExtension extension : extensions) {
+                try {
+                    extension.beforeScoreCollection(searchContext);
+                } catch (Exception e) {
+                    LOGGER.warn(
+                        new ParameterizedMessage("Failed to execute beforeScoreCollection extension [{}]", extension.getClass().getName()),
+                        e
+                    );
+                }
+            }
+
+            try {
+                return searchOperation.execute();
+            } finally {
+                // Execute afterScoreCollection extensions
+                for (QueryPhaseExtension extension : extensions) {
+                    try {
+                        extension.afterScoreCollection(searchContext);
+                    } catch (Exception e) {
+                        LOGGER.warn(
+                            new ParameterizedMessage(
+                                "Failed to execute afterScoreCollection extension [{}]",
+                                extension.getClass().getName()
+                            ),
+                            e
+                        );
+                    }
+                }
+            }
+        }
+
+        /**
+         * Functional interface for search operations that can throw IOException
+         */
+        @FunctionalInterface
+        protected interface SearchOperation {
+            boolean execute() throws IOException;
         }
 
         @Override
@@ -447,10 +498,19 @@ protected boolean searchWithCollector(
             boolean hasTimeout
         ) throws IOException {
             QueryCollectorContext queryCollectorContext = getQueryCollectorContext(searchContext, hasFilterCollector);
-            return searchWithCollector(searchContext, searcher, query, collectors, queryCollectorContext, hasFilterCollector, hasTimeout);
+            return QueryPhase.searchWithCollector(
+                searchContext,
+                searcher,
+                query,
+                collectors,
+                queryCollectorContext,
+                hasFilterCollector,
+                hasTimeout
+            );
         }
 
-        private QueryCollectorContext getQueryCollectorContext(SearchContext searchContext, boolean hasFilterCollector) throws IOException {
+        protected QueryCollectorContext getQueryCollectorContext(SearchContext searchContext, boolean hasFilterCollector)
+            throws IOException {
             // create the top docs collector last when the other collectors are known
             final Optional<QueryCollectorContext> queryCollectorContextOpt = QueryCollectorContextSpecRegistry.getQueryCollectorContextSpec(
                 searchContext,

server/src/main/java/org/opensearch/search/query/QueryPhaseExtension.java

@@ -0,0 +1,46 @@
+/*
+ * 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.ExperimentalApi;
+import org.opensearch.search.internal.SearchContext;
+
+/**
+ * Extension point interface that allows plugins to hook into the query phase
+ * before and after score collection. This enables custom CollectorManager
+ * implementations and score processing for advanced search features like
+ * hybrid queries and neural search.
+ *
+ * <p>This API is experimental and may change in future versions based on
+ * feedback from plugin implementations.</p>
+ *
+ * @opensearch.api
+ */
+@ExperimentalApi
+public interface QueryPhaseExtension {
+
+    /**
+     * Called before score collection begins in the query phase.
+     * This allows extensions to set up custom state or modify the search context
+     * before the main query execution.
+     *
+     * @param searchContext the current search context
+     */
+    void beforeScoreCollection(SearchContext searchContext);
+
+    /**
+     * Called after score collection completes in the query phase.
+     * This allows extensions to process collected scores or perform
+     * post-collection operations.
+     *
+     * @param searchContext the current search context
+     */
+    void afterScoreCollection(SearchContext searchContext);
+
+}

server/src/main/java/org/opensearch/search/query/QueryPhaseSearcher.java

@@ -17,7 +17,9 @@
 import org.opensearch.search.internal.SearchContext;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.LinkedList;
+import java.util.List;
 
 /**
  * The extension point which allows to plug in custom search implementation to be
@@ -53,4 +55,12 @@ boolean searchWith(
     default AggregationProcessor aggregationProcessor(SearchContext searchContext) {
         return new DefaultAggregationProcessor();
     }
+
+    /**
+     * Get the list of query phase extensions that should be executed before and after score collection.
+     * @return list of query phase extensions, empty list if none
+     */
+    default List<QueryPhaseExtension> queryPhaseExtensions() {
+        return Collections.emptyList();
+    }
 }