Cleanup

Signed-off-by: Atri Sharma <atrisharma@Atris-Mac-Studio.local>

Author: Atri Sharma

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

@@ -91,6 +91,11 @@ public class SearchResponse extends ActionResponse implements StatusToXContentOb
     private static final ParseField TIMED_OUT = new ParseField("timed_out");
     private static final ParseField TERMINATED_EARLY = new ParseField("terminated_early");
     private static final ParseField NUM_REDUCE_PHASES = new ParseField("num_reduce_phases");
+    
+    // Streaming REST fields
+    private static final ParseField IS_PARTIAL = new ParseField("is_partial");
+    private static final ParseField SEQUENCE_NUMBER = new ParseField("sequence_number");
+    private static final ParseField TOTAL_PARTIALS = new ParseField("total_partials");
 
     private final SearchResponseSections internalResponse;
     private final String scrollId;
@@ -132,6 +137,12 @@ public SearchResponse(StreamInput in) throws IOException {
         }
         skippedShards = in.readVInt();
         pointInTimeId = in.readOptionalString();
+        // Read streaming fields
+        if (in.getVersion().onOrAfter(Version.V_2_15_0)) {
+            isPartial = in.readBoolean();
+            sequenceNumber = in.readVInt();
+            totalPartials = in.readVInt();
+        }
     }
 
     public SearchResponse(
@@ -385,6 +396,17 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
         if (getNumReducePhases() != 1) {
             builder.field(NUM_REDUCE_PHASES.getPreferredName(), getNumReducePhases());
         }
+        
+        // Write streaming fields if it's a streaming partial
+        // or if we have sequenced streaming responses
+        if (isPartial) {
+            builder.field(IS_PARTIAL.getPreferredName(), isPartial);
+            builder.field(SEQUENCE_NUMBER.getPreferredName(), sequenceNumber);
+            if (totalPartials > 0) {
+                builder.field(TOTAL_PARTIALS.getPreferredName(), totalPartials);
+            }
+        }
+        
         RestActions.buildBroadcastShardsHeader(
             builder,
             params,
@@ -423,6 +445,9 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
         int skippedShards = 0; // 0 for BWC
         String scrollId = null;
         String searchContextId = null;
+        boolean isPartial = false;
+        int sequenceNumber = 0;
+        int totalPartials = 0;
         List<ShardSearchFailure> failures = new ArrayList<>();
         Clusters clusters = Clusters.EMPTY;
         List<SearchExtBuilder> extBuilders = new ArrayList<>();
@@ -443,6 +468,12 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
                     terminatedEarly = parser.booleanValue();
                 } else if (NUM_REDUCE_PHASES.match(currentFieldName, parser.getDeprecationHandler())) {
                     numReducePhases = parser.intValue();
+                } else if (IS_PARTIAL.match(currentFieldName, parser.getDeprecationHandler())) {
+                    isPartial = parser.booleanValue();
+                } else if (SEQUENCE_NUMBER.match(currentFieldName, parser.getDeprecationHandler())) {
+                    sequenceNumber = parser.intValue();
+                } else if (TOTAL_PARTIALS.match(currentFieldName, parser.getDeprecationHandler())) {
+                    totalPartials = parser.intValue();
                 } else {
                     parser.skipChildren();
                 }
@@ -583,6 +614,22 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
             clusters,
             searchContextId
         );
+        SearchResponse response = new SearchResponse(
+            internalResponse,
+            scrollId,
+            totalShards,
+            successfulShards,
+            skippedShards,
+            tookInMillis,
+            phaseTook,
+            failures.toArray(ShardSearchFailure.EMPTY_ARRAY),
+            clusters,
+            searchContextId
+        );
+        response.setPartial(isPartial);
+        response.setSequenceNumber(sequenceNumber);
+        response.setTotalPartials(totalPartials);
+        return response;
     }
 
     @Override
@@ -603,6 +650,13 @@ public void writeTo(StreamOutput out) throws IOException {
         }
         out.writeVInt(skippedShards);
         out.writeOptionalString(pointInTimeId);
+        
+        // Write streaming fields
+        if (out.getVersion().onOrAfter(Version.V_2_15_0)) {
+            out.writeBoolean(isPartial);
+            out.writeVInt(sequenceNumber);
+            out.writeVInt(totalPartials);
+        }
     }
 
     @Override

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

@@ -32,7 +32,6 @@
 import org.opensearch.search.aggregations.bucket.LocalBucketCountThresholds;
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.streaming.Streamable;
 import org.opensearch.search.streaming.StreamingCostMetrics;
 
 import java.io.IOException;
@@ -46,7 +45,7 @@
 import static java.util.Collections.emptyList;
 import static org.opensearch.search.aggregations.InternalOrder.isKeyOrder;
 
-public class StreamNumericTermsAggregator extends TermsAggregator implements Streamable {
+public class StreamNumericTermsAggregator extends TermsAggregator {
     private static final Logger logger = LogManager.getLogger(StreamNumericTermsAggregator.class);
     private final ResultStrategy<?, ?> resultStrategy;
     private final ValuesSource.Numeric valuesSource;

server/src/main/java/org/opensearch/rest/action/search/StreamingPreviewFirstListener.java

@@ -16,7 +16,6 @@
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.streaming.Streamable;
 import org.opensearch.search.streaming.StreamingCostMetrics;
 
 import java.io.IOException;
@@ -29,7 +28,7 @@
  *
  * @opensearch.internal
  */
-public class StreamCardinalityAggregator extends CardinalityAggregator implements Streamable {
+public class StreamCardinalityAggregator extends CardinalityAggregator {
 
     private Collector streamCollector;