Introduce Pre Post Score Collection Hooks to QueryPhase

Signed-off-by: Atri Sharma <[email protected]>

Author: atris

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/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,56 @@ public boolean searchWith(
             boolean hasFilterCollector,
             boolean hasTimeout
         ) throws IOException {
-            return searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
+            // Fast path - skip extension logic entirely if no extensions are registered
+            List<QueryPhaseExtension> extensions = queryPhaseExtensions();
+            if (extensions == null || extensions.isEmpty()) {
+                return searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
+            }
+
+            // Execute beforeScoreCollection extensions
+            for (QueryPhaseExtension extension : extensions) {
+                try {
+                    extension.beforeScoreCollection(searchContext);
+                } catch (Exception e) {
+                    if (extension.failOnError()) {
+                        throw new QueryPhaseExecutionException(
+                            searchContext.shardTarget(),
+                            "Failed to execute beforeScoreCollection extension [" + extension.getClass().getName() + "]",
+                            e
+                        );
+                    }
+                    LOGGER.warn(
+                        new ParameterizedMessage("Failed to execute beforeScoreCollection extension [{}]", extension.getClass().getName()),
+                        e
+                    );
+                }
+            }
+
+            try {
+                return searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
+            } finally {
+                // Execute afterScoreCollection extensions
+                for (QueryPhaseExtension extension : extensions) {
+                    try {
+                        extension.afterScoreCollection(searchContext);
+                    } catch (Exception e) {
+                        if (extension.failOnError()) {
+                            throw new QueryPhaseExecutionException(
+                                searchContext.shardTarget(),
+                                "Failed to execute afterScoreCollection extension [" + extension.getClass().getName() + "]",
+                                e
+                            );
+                        }
+                        LOGGER.warn(
+                            new ParameterizedMessage(
+                                "Failed to execute afterScoreCollection extension [{}]",
+                                extension.getClass().getName()
+                            ),
+                            e
+                        );
+                    }
+                }
+            }
         }
 
         @Override
@@ -447,7 +497,15 @@ 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 {

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

@@ -0,0 +1,53 @@
+/*
+ * 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;
+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.
+ *
+ * @opensearch.api
+ */
+@PublicApi(since = "3.2.0")
+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);
+
+    /**
+     * Determines whether failures in this extension should fail the entire query.
+     * When true, exceptions thrown by this extension will propagate and fail the search.
+     * When false (default), exceptions are logged and the search continues.
+     *
+     * @return true if extension failures should fail the query, false otherwise
+     */
+    default boolean failOnError() {
+        return false;
+    }
+}

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();
+    }
 }