Add streaming search with configurable scoring modes

server/src/main/java/org/opensearch/action/ActionModule.java

@@ -289,14 +289,14 @@
 import org.opensearch.action.search.SearchAction;
 import org.opensearch.action.search.SearchScrollAction;
 import org.opensearch.action.search.StreamSearchAction;
-import org.opensearch.action.search.StreamTransportSearchAction;
 import org.opensearch.action.search.TransportClearScrollAction;
 import org.opensearch.action.search.TransportCreatePitAction;
 import org.opensearch.action.search.TransportDeletePitAction;
 import org.opensearch.action.search.TransportGetAllPitsAction;
 import org.opensearch.action.search.TransportMultiSearchAction;
 import org.opensearch.action.search.TransportSearchAction;
 import org.opensearch.action.search.TransportSearchScrollAction;
+import org.opensearch.action.search.TransportStreamSearchAction;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.AutoCreateIndex;
 import org.opensearch.action.support.DestructiveOperations;
@@ -741,9 +741,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         actions.register(MultiGetAction.INSTANCE, TransportMultiGetAction.class, TransportShardMultiGetAction.class);
         actions.register(BulkAction.INSTANCE, TransportBulkAction.class, TransportShardBulkAction.class);
         actions.register(SearchAction.INSTANCE, TransportSearchAction.class);
-        if (FeatureFlags.isEnabled(FeatureFlags.STREAM_TRANSPORT)) {
-            actions.register(StreamSearchAction.INSTANCE, StreamTransportSearchAction.class);
-        }
+        actions.register(StreamSearchAction.INSTANCE, TransportStreamSearchAction.class);
         actions.register(SearchScrollAction.INSTANCE, TransportSearchScrollAction.class);
         actions.register(MultiSearchAction.INSTANCE, TransportMultiSearchAction.class);
         actions.register(ExplainAction.INSTANCE, TransportExplainAction.class);

server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java

@@ -151,6 +151,14 @@ int getBatchReduceSize(int requestBatchedReduceSize, int minBatchReduceSize) {
         return (hasAggs || hasTopDocs) ? Math.min(requestBatchedReduceSize, minBatchReduceSize) : minBatchReduceSize;
     }
 
+    /**
+     * Protected accessor for progressListener to allow subclasses to access it.
+     * @return the search progress listener
+     */
+    protected SearchProgressListener progressListener() {
+        return this.progressListener;
+    }
+
     @Override
     public void close() {
         Releasables.close(pendingReduces);
@@ -160,7 +168,7 @@ public void close() {
     public void consumeResult(SearchPhaseResult result, Runnable next) {
         super.consumeResult(result, () -> {});
         QuerySearchResult querySearchResult = result.queryResult();
-        progressListener.notifyQueryResult(querySearchResult.getShardIndex());
+        progressListener.notifyQueryResult(querySearchResult.getShardIndex(), querySearchResult.getSearchShardTarget());
         pendingReduces.consume(querySearchResult, next);
     }
 
@@ -228,7 +236,10 @@ private ReduceResult partialReduce(
         Arrays.sort(toConsume, Comparator.comparingInt(QuerySearchResult::getShardIndex));
 
         for (QuerySearchResult result : toConsume) {
-            topDocsStats.add(result.topDocs(), result.searchTimedOut(), result.terminatedEarly());
+            // Use non-consuming topDocs() for stats aggregation only
+            if (result.hasTopDocs()) {
+                topDocsStats.add(result.topDocs(), result.searchTimedOut(), result.terminatedEarly());
+            }
         }
 
         final TopDocs newTopDocs;
@@ -238,7 +249,9 @@ private ReduceResult partialReduce(
                 topDocsList.add(lastReduceResult.reducedTopDocs);
             }
             for (QuerySearchResult result : toConsume) {
+                // Consume TopDocs exactly once for merge/reduce phase
                 TopDocsAndMaxScore topDocs = result.consumeTopDocs();
+                // For streaming, avoid reassigning shardIndex if already set
                 SearchPhaseController.setShardIndex(topDocs.topDocs, result.getShardIndex());
                 topDocsList.add(topDocs.topDocs);
             }
@@ -462,35 +475,39 @@ private void tryExecuteNext() {
                 runningTask.compareAndSet(null, task);
             }
 
-            executor.execute(new AbstractRunnable() {
-                @Override
-                protected void doRun() {
-                    final ReduceResult thisReduceResult = reduceResult;
-                    long estimatedTotalSize = (thisReduceResult != null ? thisReduceResult.estimatedSize : 0) + task.aggsBufferSize;
-                    final ReduceResult newReduceResult;
-                    try {
-                        final QuerySearchResult[] toConsume = task.consumeBuffer();
-                        if (toConsume == null) {
-                            onAfterReduce(task, null, 0);
+            try {
+                executor.execute(new AbstractRunnable() {
+                    @Override
+                    protected void doRun() {
+                        final ReduceResult thisReduceResult = reduceResult;
+                        long estimatedTotalSize = (thisReduceResult != null ? thisReduceResult.estimatedSize : 0) + task.aggsBufferSize;
+                        final ReduceResult newReduceResult;
+                        try {
+                            final QuerySearchResult[] toConsume = task.consumeBuffer();
+                            if (toConsume == null) {
+                                onAfterReduce(task, null, 0);
+                                return;
+                            }
+                            long estimateRamBytesUsedForReduce = estimateRamBytesUsedForReduce(estimatedTotalSize);
+                            addEstimateAndMaybeBreak(estimateRamBytesUsedForReduce);
+                            estimatedTotalSize += estimateRamBytesUsedForReduce;
+                            ++numReducePhases;
+                            newReduceResult = partialReduce(toConsume, task.emptyResults, topDocsStats, thisReduceResult, numReducePhases);
+                        } catch (Exception t) {
+                            PendingReduces.this.onFailure(t);
                             return;
                         }
-                        long estimateRamBytesUsedForReduce = estimateRamBytesUsedForReduce(estimatedTotalSize);
-                        addEstimateAndMaybeBreak(estimateRamBytesUsedForReduce);
-                        estimatedTotalSize += estimateRamBytesUsedForReduce;
-                        ++numReducePhases;
-                        newReduceResult = partialReduce(toConsume, task.emptyResults, topDocsStats, thisReduceResult, numReducePhases);
-                    } catch (Exception t) {
-                        PendingReduces.this.onFailure(t);
-                        return;
+                        onAfterReduce(task, newReduceResult, estimatedTotalSize);
                     }
-                    onAfterReduce(task, newReduceResult, estimatedTotalSize);
-                }
 
-                @Override
-                public void onFailure(Exception exc) {
-                    PendingReduces.this.onFailure(exc);
-                }
-            });
+                    @Override
+                    public void onFailure(Exception exc) {
+                        PendingReduces.this.onFailure(exc);
+                    }
+                });
+            } catch (Exception e) {
+                onFailure(e);
+            }
         }
 
         private void onAfterReduce(ReduceTask task, ReduceResult newResult, long estimatedSize) {
@@ -516,7 +533,11 @@ private void onAfterReduce(ReduceTask task, ReduceResult newResult, long estimat
                 }
             }
             task.consumeListener();
-            executor.execute(this::tryExecuteNext);
+            try {
+                executor.execute(this::tryExecuteNext);
+            } catch (Exception e) {
+                onFailure(e);
+            }
         }
 
         // Idempotent and thread-safe failure handling
@@ -549,7 +570,10 @@ private synchronized void clearReduceTaskQueue() {
 
         private synchronized SearchPhaseController.TopDocsStats consumeTopDocsStats() {
             for (QuerySearchResult result : buffer) {
-                topDocsStats.add(result.topDocs(), result.searchTimedOut(), result.terminatedEarly());
+                // Use non-consuming topDocs() for stats aggregation only
+                if (result.hasTopDocs()) {
+                    topDocsStats.add(result.topDocs(), result.searchTimedOut(), result.terminatedEarly());
+                }
             }
             return topDocsStats;
         }
@@ -563,6 +587,7 @@ private synchronized List<TopDocs> consumeTopDocs() {
                 topDocsList.add(reduceResult.reducedTopDocs);
             }
             for (QuerySearchResult result : buffer) {
+                // Consume TopDocs exactly once for merge/reduce phase
                 TopDocsAndMaxScore topDocs = result.consumeTopDocs();
                 SearchPhaseController.setShardIndex(topDocs.topDocs, result.getShardIndex());
                 topDocsList.add(topDocs.topDocs);

server/src/main/java/org/opensearch/action/search/SearchPhaseController.java

@@ -246,7 +246,14 @@ static TopDocs mergeTopDocs(Collection<TopDocs> results, int topN, int from) {
     }
 
     static void setShardIndex(TopDocs topDocs, int shardIndex) {
-        assert topDocs.scoreDocs.length == 0 || topDocs.scoreDocs[0].shardIndex == -1 : "shardIndex is already set";
+        // Idempotent assignment: in streaming flows partial reductions may touch the same TopDocs more than once.
+        if (topDocs.scoreDocs.length == 0) {
+            return;
+        }
+        if (topDocs.scoreDocs[0].shardIndex != -1) {
+            // Already set by a previous pass; avoid reassigning to prevent assertion failures
+            return;
+        }
         for (ScoreDoc doc : topDocs.scoreDocs) {
             doc.shardIndex = shardIndex;
         }
@@ -424,6 +431,7 @@ public ReduceContext forFinalReduction() {
         final List<TopDocs> topDocs = new ArrayList<>();
         for (SearchPhaseResult sortedResult : queryResults) {
             QuerySearchResult queryResult = sortedResult.queryResult();
+            // Consume TopDocs exactly once for merge/reduce phase
             final TopDocsAndMaxScore td = queryResult.consumeTopDocs();
             assert td != null;
             topDocsStats.add(td, queryResult.searchTimedOut(), queryResult.terminatedEarly());
@@ -795,40 +803,32 @@ QueryPhaseResultConsumer newSearchPhaseResults(
         Consumer<Exception> onPartialMergeFailure,
         BooleanSupplier isTaskCancelled
     ) {
-        return new QueryPhaseResultConsumer(
-            request,
-            executor,
-            circuitBreaker,
-            this,
-            listener,
-            namedWriteableRegistry,
-            numShards,
-            onPartialMergeFailure,
-            isTaskCancelled
-        );
-    }
-
-    /**
-     * Returns a new {@link StreamQueryPhaseResultConsumer} instance that reduces search responses incrementally.
-     */
-    StreamQueryPhaseResultConsumer newStreamSearchPhaseResults(
-        Executor executor,
-        CircuitBreaker circuitBreaker,
-        SearchProgressListener listener,
-        SearchRequest request,
-        int numShards,
-        Consumer<Exception> onPartialMergeFailure
-    ) {
-        return new StreamQueryPhaseResultConsumer(
-            request,
-            executor,
-            circuitBreaker,
-            this,
-            listener,
-            namedWriteableRegistry,
-            numShards,
-            onPartialMergeFailure
-        );
+        String streamingMode = request.getStreamingSearchMode();
+        if (streamingMode != null) {
+            return new StreamQueryPhaseResultConsumer(
+                request,
+                executor,
+                circuitBreaker,
+                this,
+                listener,
+                namedWriteableRegistry,
+                numShards,
+                onPartialMergeFailure
+            );
+        } else {
+            // Regular QueryPhaseResultConsumer
+            return new QueryPhaseResultConsumer(
+                request,
+                executor,
+                circuitBreaker,
+                this,
+                listener,
+                namedWriteableRegistry,
+                numShards,
+                onPartialMergeFailure,
+                isTaskCancelled
+            );
+        }
     }
 
     /**

server/src/main/java/org/opensearch/action/search/SearchProgressListener.java

@@ -80,6 +80,16 @@ protected void onListShards(List<SearchShard> shards, List<SearchShard> skippedS
      */
     protected void onQueryResult(int shardIndex) {}
 
+    /**
+     * Executed when a shard returns a query result.
+     *
+     * @param shardIndex The index of the shard in the list provided by {@link SearchProgressListener#onListShards})}.
+     * @param shardTarget The shard target that returned the result.
+     */
+    protected void onQueryResult(int shardIndex, SearchShardTarget shardTarget) {
+        onQueryResult(shardIndex);
+    }
+
     /**
      * Executed when a shard reports a query failure.
      *
@@ -100,6 +110,7 @@ protected void onQueryFailure(int shardIndex, SearchShardTarget shardTarget, Exc
      */
     protected void onPartialReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {}
 
+
     /**
      * Executed once when the final reduce is created.
      *
@@ -146,6 +157,17 @@ final void notifyQueryResult(int shardIndex) {
         }
     }
 
+    final void notifyQueryResult(int shardIndex, SearchShardTarget shardTarget) {
+        try {
+            onQueryResult(shardIndex, shardTarget);
+        } catch (Exception e) {
+            logger.warn(
+                () -> new ParameterizedMessage("[{}] Failed to execute progress listener on query result", shards.get(shardIndex)),
+                e
+            );
+        }
+    }
+
     final void notifyQueryFailure(int shardIndex, SearchShardTarget shardTarget, Exception exc) {
         try {
             onQueryFailure(shardIndex, shardTarget, exc);
@@ -165,6 +187,7 @@ final void notifyPartialReduce(List<SearchShard> shards, TotalHits totalHits, In
         }
     }
 
+
     protected final void notifyFinalReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {
         try {
             onFinalReduce(shards, totalHits, aggs, reducePhase);

server/src/main/java/org/opensearch/action/search/SearchQueryThenFetchAsyncAction.java

@@ -36,6 +36,7 @@
 import org.apache.lucene.search.TopFieldDocs;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.routing.GroupShardsIterator;
+import org.opensearch.common.lucene.search.TopDocsAndMaxScore;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.search.SearchPhaseResult;
 import org.opensearch.search.SearchShardTarget;
@@ -149,17 +150,19 @@ protected void onShardResult(SearchPhaseResult result, SearchShardIterator shard
         if (queryResult.isNull() == false
             // disable sort optims for scroll requests because they keep track of the last bottom doc locally (per shard)
             && getRequest().scroll() == null
-            && queryResult.topDocs() != null
-            && queryResult.topDocs().topDocs.getClass() == TopFieldDocs.class) {
-            TopFieldDocs topDocs = (TopFieldDocs) queryResult.topDocs().topDocs;
-            if (bottomSortCollector == null) {
-                synchronized (this) {
-                    if (bottomSortCollector == null) {
-                        bottomSortCollector = new BottomSortValuesCollector(topDocsSize, topDocs.fields);
+            && queryResult.hasTopDocs()) {
+            TopDocsAndMaxScore topDocsAndMaxScore = queryResult.topDocs();
+            if (topDocsAndMaxScore != null && topDocsAndMaxScore.topDocs.getClass() == TopFieldDocs.class) {
+                TopFieldDocs topDocs = (TopFieldDocs) topDocsAndMaxScore.topDocs;
+                if (bottomSortCollector == null) {
+                    synchronized (this) {
+                        if (bottomSortCollector == null) {
+                            bottomSortCollector = new BottomSortValuesCollector(topDocsSize, topDocs.fields);
+                        }
                     }
                 }
+                bottomSortCollector.consumeTopDocs(topDocs, queryResult.sortValueFormats());
             }
-            bottomSortCollector.consumeTopDocs(topDocs, queryResult.sortValueFormats());
         }
         super.onShardResult(result, shardIt);
     }