Add streaming search with configurable scoring modes

modules/reindex/src/test/java/org/opensearch/index/reindex/DeleteByQueryBasicTests.java

@@ -54,6 +54,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static org.opensearch.action.search.StreamSearchTransportService.STREAM_SEARCH_ENABLED;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_READ_ONLY;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_READ_ONLY_ALLOW_DELETE;
 import static org.opensearch.index.query.QueryBuilders.matchQuery;
@@ -383,6 +384,63 @@ public void testMissingSources() {
         assertThat(response, matcher().deleted(0).slices(hasSize(0)));
     }
 
+    /**
+     * Regression test to ensure delete-by-query works correctly even when streaming search is enabled.
+     * Since delete-by-query uses scroll searches, it should automatically opt out of the streaming pipeline.
+     */
+    public void testDeleteByQueryWithStreamingEnabled() throws Exception {
+        // Enable streaming search globally
+        Settings streamingSettings = Settings.builder().put(STREAM_SEARCH_ENABLED.getKey(), true).build();
+        assertAcked(client().admin().cluster().prepareUpdateSettings().setTransientSettings(streamingSettings).get());
+
+        try {
+            // Index some test documents
+            indexRandom(
+                true,
+                client().prepareIndex("test").setId("1").setSource("foo", "delete_me"),
+                client().prepareIndex("test").setId("2").setSource("foo", "delete_me"),
+                client().prepareIndex("test").setId("3").setSource("foo", "keep_me")
+            );
+
+            assertHitCount(client().prepareSearch("test").setSize(0).get(), 3);
+
+            // Perform delete-by-query - this should work without "topDocs already consumed" errors
+            BulkByScrollResponse response = deleteByQuery().source("test").filter(termQuery("foo", "delete_me")).refresh(true).get();
+
+            // Verify the delete operation succeeded
+            assertThat(response, matcher().deleted(2));
+            assertHitCount(client().prepareSearch("test").setSize(0).get(), 1);
+
+            // Verify only the expected document remains
+            assertHitCount(client().prepareSearch("test").setQuery(termQuery("foo", "keep_me")).get(), 1);
+        } finally {
+            // Comprehensive cleanup: reset both transient and persistent streaming settings
+            Settings disableStreamingSettings = Settings.builder().putNull(STREAM_SEARCH_ENABLED.getKey()).build();
+            try {
+                // Clear transient settings
+                assertAcked(client().admin().cluster().prepareUpdateSettings().setTransientSettings(disableStreamingSettings).get());
+                // Clear persistent settings (in case they were set)
+                assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(disableStreamingSettings).get());
+
+                // Assert the setting is cleared to prevent test pollution
+                assertFalse(
+                    "Stream search should be disabled after cleanup",
+                    client().admin()
+                        .cluster()
+                        .prepareState()
+                        .get()
+                        .getState()
+                        .getMetadata()
+                        .transientSettings()
+                        .getAsBoolean(STREAM_SEARCH_ENABLED.getKey(), false)
+                );
+            } catch (Exception cleanupException) {
+                // Log cleanup failures but don't fail the test
+                logger.warn("Failed to clean up streaming settings", cleanupException);
+            }
+        }
+    }
+
     /** Enables or disables the cluster disk allocation decider **/
     private void setDiskAllocationDeciderEnabled(boolean value) {
         Settings settings = value

plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/streaming/aggregation/SubAggregationIT.java

@@ -10,6 +10,8 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.action.admin.indices.create.CreateIndexRequest;
 import org.opensearch.action.admin.indices.create.CreateIndexResponse;
 import org.opensearch.action.admin.indices.flush.FlushRequest;
@@ -52,6 +54,7 @@
 
 @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, minNumDataNodes = 3, maxNumDataNodes = 3)
 public class SubAggregationIT extends ParameterizedDynamicSettingsOpenSearchIntegTestCase {
+    private static final Logger logger = LogManager.getLogger(SubAggregationIT.class);
 
     public SubAggregationIT(Settings dynamicSettings) {
         super(dynamicSettings);
@@ -84,6 +87,7 @@ public void setUp() throws Exception {
             .prepareUpdateSettings()
             .setTransientSettings(
                 Settings.builder()
+                    .put("stream.search.enabled", true)
                     .put("search.aggregations.streaming.max_estimated_bucket_count", 1000)
                     .put("search.aggregations.streaming.min_cardinality_ratio", 0.001)
                     .put("search.aggregations.streaming.min_estimated_bucket_count", 1)
@@ -179,6 +183,7 @@ public void tearDown() throws Exception {
             .prepareUpdateSettings()
             .setTransientSettings(
                 Settings.builder()
+                    .putNull("stream.search.enabled")
                     .putNull("search.aggregations.streaming.max_estimated_bucket_count")
                     .putNull("search.aggregations.streaming.min_cardinality_ratio")
                     .putNull("search.aggregations.streaming.min_estimated_bucket_count")
@@ -197,8 +202,14 @@ public void testStreamingAggregationUsed() throws Exception {
             .setSize(0)
             .setRequestCache(false)
             .setProfile(true)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
         assertEquals(90, resp.getHits().getTotalHits().value());
@@ -231,8 +242,14 @@ public void testStreamingAggregationTerm() throws Exception {
             .addAggregation(agg)
             .setSize(0)
             .setRequestCache(false)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
         assertEquals(90, resp.getHits().getTotalHits().value());
@@ -272,8 +289,14 @@ public void testStreamingNumericAggregationUsed() throws Exception {
             .setSize(0)
             .setRequestCache(false)
             .setProfile(true)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
         assertEquals(90, resp.getHits().getTotalHits().value());
@@ -305,8 +328,14 @@ public void testStreamingNumericAggregation() throws Exception {
             .addAggregation(agg)
             .setSize(0)
             .setRequestCache(false)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
 
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
@@ -340,8 +369,14 @@ public void testStreamingAggregationWithoutProfile() throws Exception {
             .addAggregation(agg)
             .setSize(0)
             .setRequestCache(false)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute(); // No profile
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
 
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
@@ -383,8 +418,14 @@ public void testStreamingAggregationNotUsedWithRestrictiveLimits() throws Except
                 .setSize(0)
                 .setRequestCache(false)
                 .setProfile(true)
+                .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
                 .execute();
-            SearchResponse resp = future.actionGet();
+            SearchResponse resp;
+            try {
+                resp = future.actionGet(TimeValue.timeValueSeconds(35));
+            } catch (Exception e) {
+                throw e;
+            }
 
             assertNotNull(resp);
             assertEquals(NUM_SHARDS, resp.getTotalShards());
@@ -446,8 +487,14 @@ public void testStreamingCardinalityAggregationUsed() throws Exception {
             .setSize(0)
             .setRequestCache(false)
             .setProfile(true)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
         assertEquals(90, resp.getHits().getTotalHits().value());
@@ -487,8 +534,14 @@ public void testStreamingCardinalityAggregation() throws Exception {
             .addAggregation(AggregationBuilders.cardinality("cardinality_agg").field("field1").precisionThreshold(1000))
             .setSize(0)
             .setRequestCache(false)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
 
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
@@ -510,8 +563,14 @@ public void testStreamingCardinalityWithPrecisionThreshold() throws Exception {
             .addAggregation(AggregationBuilders.cardinality("cardinality_high").field("field1").precisionThreshold(1000))
             .setSize(0)
             .setRequestCache(false)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
 
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());
@@ -540,8 +599,14 @@ public void testStreamingCardinalityAsSubAggregation() throws Exception {
             .addAggregation(agg)
             .setSize(0)
             .setRequestCache(false)
+            .setCancelAfterTimeInterval(TimeValue.timeValueSeconds(30))
             .execute();
-        SearchResponse resp = future.actionGet();
+        SearchResponse resp;
+        try {
+            resp = future.actionGet(TimeValue.timeValueSeconds(35));
+        } catch (Exception e) {
+            throw e;
+        }
 
         assertNotNull(resp);
         assertEquals(NUM_SHARDS, resp.getTotalShards());

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightStreamPlugin.java

@@ -235,6 +235,21 @@ public Map<String, Supplier<Transport>> getTransports(
         return Collections.emptyMap();
     }
 
+    @Override
+    public Map<String, Supplier<Transport>> getStreamTransports(
+        Settings settings,
+        ThreadPool threadPool,
+        PageCacheRecycler pageCacheRecycler,
+        CircuitBreakerService circuitBreakerService,
+        NamedWriteableRegistry namedWriteableRegistry,
+        NetworkService networkService,
+        Tracer tracer
+    ) {
+        // Return empty map since getTransports() already registers the FLIGHT transport
+        // The NetworkModule will find it via getStreamTransportSupplier() which looks in transportFactories
+        return Collections.emptyMap();
+    }
+
     /**
      * Gets the auxiliary transports for the FlightStream plugin.
      * @param settings The settings for the plugin.

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;
@@ -740,9 +740,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/PreviewFirstPartialReceiver.java

@@ -0,0 +1,26 @@
+/*
+ * 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.action.search;
+
+/**
+ * Marker interface for ActionListeners that can receive the first partial response
+ * in streaming preview-first mode. This is a temporary interface for OSB benchmarking.
+ *
+ * @opensearch.internal
+ */
+public interface PreviewFirstPartialReceiver {
+
+    /**
+     * Called when the first partial response is available.
+     * Implementation should send the response and cancel the search task.
+     *
+     * @param partial the first partial search response
+     */
+    void onPartialResponse(SearchResponse partial);
+}