Skip two phase rescore for sort query (#1898)

* skip two phase rescore for sort query

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

* ad changelog

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

---------

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

Author: chishui

CHANGELOG.md

@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Add `previous_score_field` to `by_field` rerank processor to avoid overwriting existing document fields ([#1880](https://github.com/opensearch-project/neural-search/pull/1880)) (OpenSearch [#21440](https://github.com/opensearch-project/OpenSearch/issues/21440))
 * [Hybrid Query] Block hybrid query execution with `search_type=dfs_query_then_fetch` ([#1873](https://github.com/opensearch-project/neural-search/pull/1873))
 * [Hybrid Query] Fix `hybrid_score_explanation` returning a single normalization block for hybrid queries on indices that contain a nested field ([#1876](https://github.com/opensearch-project/neural-search/pull/1876))
+* [Two-phase] Skip two phase rescore for sort query ([#1898](https://github.com/opensearch-project/neural-search/pull/1898))
 
 ### Infrastructure
 * [Hybrid Query] Fix flaky `HybridQueryExplainIT` explanation test by asserting per-document explanation trees keyed by `_id` instead of hit position ([#1899](https://github.com/opensearch-project/neural-search/pull/1899))

src/main/java/org/opensearch/neuralsearch/processor/NeuralSparseTwoPhaseProcessor.java

@@ -27,7 +27,11 @@
 import org.opensearch.search.pipeline.SearchRequestProcessor;
 import org.opensearch.search.rescore.QueryRescorerBuilder;
 import org.opensearch.search.rescore.RescorerBuilder;
+import org.opensearch.search.sort.ScoreSortBuilder;
+import org.opensearch.search.sort.SortBuilder;
+import org.opensearch.search.sort.SortOrder;
 
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
@@ -103,6 +107,13 @@ public SearchRequest processRequest(final SearchRequest request) {
         if (!enabled || pruneRatio == 0f) {
             return request;
         }
+        // Two-phase rescore is incompatible with explicit sort (other than _score DESC).
+        // OpenSearch rejects sort + rescore in DefaultSearchContext.preProcess(), so when the
+        // request already has a non-_score-DESC sort, skip the optimization and let the full
+        // neural_sparse query run as a single phase. Correctness over latency.
+        if (hasIncompatibleSort(request.source())) {
+            return request;
+        }
         QueryBuilder queryBuilder = request.source().query();
         // Collect the nested NeuralSparseQueryBuilder in the whole query.
         Multimap<AbstractNeuralQueryBuilder<?>, Float> queryBuilderMap = collectNeuralQueryBuilderWithSparseEmbedding(
@@ -139,6 +150,28 @@ private QueryBuilder getNestedQueryBuilderFromNeuralSparseQueryBuilderMap(
         return boolQueryBuilder;
     }
 
+    /**
+     * Returns true when the request specifies a sort that OpenSearch will treat as an explicit
+     * sort at execution time, in which case adding a rescore would cause
+     * {@code DefaultSearchContext.preProcess()} to throw
+     * "Cannot use [sort] option in conjunction with [rescore]".
+     *
+     * Mirrors the optimization in {@code SortBuilder.buildSort}: a single {@code _score DESC}
+     * is collapsed to "no sort" and is therefore compatible. Anything else (a field sort, a
+     * non-default _score order, or a multi-element sort list) is incompatible.
+     */
+    static boolean hasIncompatibleSort(final SearchSourceBuilder searchSourceBuilder) {
+        List<SortBuilder<?>> sorts = searchSourceBuilder.sorts();
+        if (sorts == null || sorts.isEmpty()) {
+            return false;
+        }
+        if (sorts.size() == 1) {
+            SortBuilder<?> only = sorts.get(0);
+            return !(only instanceof ScoreSortBuilder) || only.order() != SortOrder.DESC;
+        }
+        return true;
+    }
+
     private float getOriginQueryWeightAfterRescore(final SearchSourceBuilder searchSourceBuilder) {
         if (Objects.isNull(searchSourceBuilder.rescores())) {
             return 1.0f;

src/test/java/org/opensearch/neuralsearch/processor/NeuralSparseTwoPhaseProcessorTests.java

@@ -18,6 +18,9 @@
 import org.opensearch.neuralsearch.util.prune.PruneUtils;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.rescore.QueryRescorerBuilder;
+import org.opensearch.search.sort.ScoreSortBuilder;
+import org.opensearch.search.sort.SortBuilders;
+import org.opensearch.search.sort.SortOrder;
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.util.Collections;
@@ -256,6 +259,64 @@ public void testProcessRequest_whenTwoPhaseEnabledWithNeuralQueryNonSparseEmbedd
         assertNull(searchRequest.source().rescores());
     }
 
+    public void testProcessRequest_whenSortByField_thenSkipRescore() throws Exception {
+        NeuralSparseTwoPhaseProcessor.Factory factory = new NeuralSparseTwoPhaseProcessor.Factory();
+        NeuralSparseQueryBuilder neuralQueryBuilder = new NeuralSparseQueryBuilder();
+        SearchRequest searchRequest = new SearchRequest();
+        searchRequest.source(
+            new SearchSourceBuilder().query(neuralQueryBuilder).sort(SortBuilders.fieldSort("entityId").order(SortOrder.ASC))
+        );
+        NeuralSparseTwoPhaseProcessor processor = createTestProcessor(factory, 0.5f, true, 4.0f, 10000);
+        processor.processRequest(searchRequest);
+        assertNull(searchRequest.source().rescores());
+    }
+
+    public void testProcessRequest_whenSortByScoreDescAndField_thenSkipRescore() throws Exception {
+        NeuralSparseTwoPhaseProcessor.Factory factory = new NeuralSparseTwoPhaseProcessor.Factory();
+        NeuralSparseQueryBuilder neuralQueryBuilder = new NeuralSparseQueryBuilder();
+        SearchRequest searchRequest = new SearchRequest();
+        searchRequest.source(
+            new SearchSourceBuilder().query(neuralQueryBuilder)
+                .sort(new ScoreSortBuilder())
+                .sort(SortBuilders.fieldSort("entityId").order(SortOrder.ASC))
+        );
+        NeuralSparseTwoPhaseProcessor processor = createTestProcessor(factory, 0.5f, true, 4.0f, 10000);
+        processor.processRequest(searchRequest);
+        assertNull(searchRequest.source().rescores());
+    }
+
+    public void testProcessRequest_whenSortByScoreDescOnly_thenAddRescore() throws Exception {
+        NeuralSparseTwoPhaseProcessor.Factory factory = new NeuralSparseTwoPhaseProcessor.Factory();
+        NeuralSparseQueryBuilder neuralQueryBuilder = new NeuralSparseQueryBuilder();
+        SearchRequest searchRequest = new SearchRequest();
+        searchRequest.source(new SearchSourceBuilder().query(neuralQueryBuilder).sort(new ScoreSortBuilder()));
+        NeuralSparseTwoPhaseProcessor processor = createTestProcessor(factory, 0.5f, true, 4.0f, 10000);
+        processor.processRequest(searchRequest);
+        assertNotNull(searchRequest.source().rescores());
+    }
+
+    public void testProcessRequest_whenSortByScoreAsc_thenSkipRescore() throws Exception {
+        NeuralSparseTwoPhaseProcessor.Factory factory = new NeuralSparseTwoPhaseProcessor.Factory();
+        NeuralSparseQueryBuilder neuralQueryBuilder = new NeuralSparseQueryBuilder();
+        SearchRequest searchRequest = new SearchRequest();
+        searchRequest.source(new SearchSourceBuilder().query(neuralQueryBuilder).sort(new ScoreSortBuilder().order(SortOrder.ASC)));
+        NeuralSparseTwoPhaseProcessor processor = createTestProcessor(factory, 0.5f, true, 4.0f, 10000);
+        processor.processRequest(searchRequest);
+        assertNull(searchRequest.source().rescores());
+    }
+
+    public void testProcessRequest_whenSortByScoreDescWithTrackScores_thenAddRescore() throws Exception {
+        NeuralSparseTwoPhaseProcessor.Factory factory = new NeuralSparseTwoPhaseProcessor.Factory();
+        NeuralSparseQueryBuilder neuralQueryBuilder = new NeuralSparseQueryBuilder();
+        SearchRequest searchRequest = new SearchRequest();
+        searchRequest.source(
+            new SearchSourceBuilder().query(neuralQueryBuilder).sort(new ScoreSortBuilder()).trackScores(true)
+        );
+        NeuralSparseTwoPhaseProcessor processor = createTestProcessor(factory, 0.5f, true, 4.0f, 10000);
+        processor.processRequest(searchRequest);
+        assertNotNull(searchRequest.source().rescores());
+    }
+
     public void testType() throws Exception {
         NeuralSparseTwoPhaseProcessor.Factory factory = new NeuralSparseTwoPhaseProcessor.Factory();
         NeuralSparseTwoPhaseProcessor processor = createTestProcessor(factory);