Integrated with the compute service and stream working

Author: ivancea

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/stream/DefaultEsqlQueryResponseStream.java

@@ -53,6 +53,8 @@ protected boolean canBeStreamed() {
     protected void doStartResponse(List<ColumnInfoImpl> columns) {
         assert dropNullColumns == false : "this method doesn't support dropping null columns";
 
+        this.columns = columns;
+
         var content = new ArrayList<Iterator<? extends ToXContent>>(3);
 
         content.add(ChunkedToXContentHelper.startObject());

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/stream/EsqlQueryResponseStream.java

@@ -23,10 +23,14 @@
 import org.elasticsearch.xpack.esql.action.EsqlQueryRequest;
 import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
 import org.elasticsearch.xpack.esql.arrow.ArrowFormat;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.expression.function.UnsupportedAttribute;
 import org.elasticsearch.xpack.esql.formatter.TextFormat;
 import org.elasticsearch.xpack.esql.plugin.EsqlMediaTypeParser;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 
@@ -54,7 +58,7 @@ public static EsqlQueryResponseStream forMediaType(
             return new NonStreamingEsqlQueryResponseStream(restChannel, restRequest, esqlRequest);
         }
 
-        MediaType mediaType = EsqlMediaTypeParser.getResponseMediaType(restRequest, XContentType.JSON);
+        MediaType mediaType = EsqlMediaTypeParser.getResponseMediaType(restRequest, esqlRequest);
 
         if (mediaType instanceof TextFormat) {
             // TODO: Add support
@@ -71,6 +75,7 @@ public static EsqlQueryResponseStream forMediaType(
     protected final RestRequest restRequest;
     protected final EsqlQueryRequest esqlRequest;
 
+    // TODO: Maybe create this on startResponse()? Does creating this do something with the response? Can we still safely set headers?
     private final StreamingXContentResponse streamingXContentResponse;
 
     /**
@@ -95,14 +100,27 @@ protected EsqlQueryResponseStream(RestChannel restChannel, RestRequest restReque
     /**
      * Starts the response stream. This is the first method to be called
      */
-    public final void startResponse(List<ColumnInfoImpl> columns) {
+    public final void startResponse(List<Attribute> schema) {
         assert initialStreamChunkSent : "startResponse() called more than once";
         assert finished == false : "sendPages() called on a finished stream";
 
         if (canBeStreamed() == false) {
             return;
         }
 
+        // TODO: Copied from TransportEsqlQueryAction#toResponse. Deduplicate this code
+        List<ColumnInfoImpl> columns = schema.stream().map(c -> {
+            List<String> originalTypes;
+            if (c instanceof UnsupportedAttribute ua) {
+                // Sort the original types so they are easier to test against and prettier.
+                originalTypes = new ArrayList<>(ua.originalTypes());
+                Collections.sort(originalTypes);
+            } else {
+                originalTypes = null;
+            }
+            return new ColumnInfoImpl(c.name(), c.dataType().outputType(), originalTypes);
+        }).toList();
+
         doStartResponse(columns);
         initialStreamChunkSent = true;
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.compute.operator.exchange.ExchangeSink;
 import org.elasticsearch.compute.operator.exchange.ExchangeSinkHandler;
 import org.elasticsearch.compute.operator.exchange.ExchangeSourceHandler;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -49,6 +50,7 @@
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.esql.action.EsqlExecutionInfo;
 import org.elasticsearch.xpack.esql.action.EsqlQueryAction;
+import org.elasticsearch.xpack.esql.action.stream.EsqlQueryResponseStream;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.enrich.EnrichLookupService;
@@ -189,6 +191,7 @@ public void execute(
         Configuration configuration,
         FoldContext foldContext,
         EsqlExecutionInfo execInfo,
+        EsqlQueryResponseStream responseStream,
         ActionListener<Result> listener
     ) {
         assert ThreadPool.assertCurrentThreadPool(
@@ -204,12 +207,31 @@ public void execute(
 
         // we have no sub plans, so we can just execute the given plan
         if (subplans == null || subplans.isEmpty()) {
-            executePlan(sessionId, rootTask, flags, physicalPlan, configuration, foldContext, execInfo, null, listener, null);
+            responseStream.startResponse(physicalPlan.output());
+
+            executePlan(
+                sessionId,
+                rootTask,
+                flags,
+                physicalPlan,
+                configuration,
+                foldContext,
+                execInfo,
+                null,
+                listener,
+                responseStream,
+                null
+            );
             return;
         }
 
         final List<Page> collectedPages = Collections.synchronizedList(new ArrayList<>());
-        PhysicalPlan mainPlan = new OutputExec(subplansAndMainPlan.v2(), collectedPages::add);
+        PhysicalPlan mainPlan = new OutputExec(subplansAndMainPlan.v2(), page -> {
+            collectedPages.add(page);
+            responseStream.sendPages(List.of(page));
+        });
+
+        responseStream.startResponse(mainPlan.output());
 
         listener = listener.delegateResponse((l, e) -> {
             collectedPages.forEach(p -> Releasables.closeExpectNoException(p::releaseBlocks));
@@ -279,6 +301,7 @@ public void execute(
                             exchangeService.finishSinkHandler(childSessionId, e);
                             subPlanListener.onFailure(e);
                         }),
+                        null,
                         () -> exchangeSink.createExchangeSink(() -> {})
                     );
                 }
@@ -294,9 +317,10 @@ public void executePlan(
         Configuration configuration,
         FoldContext foldContext,
         EsqlExecutionInfo execInfo,
-        String profileQualifier,
+        @Nullable String profileQualifier,
         ActionListener<Result> listener,
-        Supplier<ExchangeSink> exchangeSinkSupplier
+        @Nullable EsqlQueryResponseStream responseStream,
+        @Nullable Supplier<ExchangeSink> exchangeSinkSupplier
     ) {
         Tuple<PhysicalPlan, PhysicalPlan> coordinatorAndDataNodePlan = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(
             physicalPlan,
@@ -310,7 +334,11 @@ public void executePlan(
         PhysicalPlan coordinatorPlan = coordinatorAndDataNodePlan.v1();
 
         if (exchangeSinkSupplier == null) {
-            coordinatorPlan = new OutputExec(coordinatorAndDataNodePlan.v1(), collectedPages::add);
+            assert responseStream != null : "responseStream must not be null when exchangeSinkSupplier is null";
+            coordinatorPlan = new OutputExec(coordinatorAndDataNodePlan.v1(), page -> {
+                collectedPages.add(page);
+                responseStream.sendPages(List.of(page));
+            });
         }
 
         PhysicalPlan dataNodePlan = coordinatorAndDataNodePlan.v2();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java

@@ -250,6 +250,7 @@ private void innerExecute(Task task, EsqlQueryRequest request, ActionListener<Es
             configuration,
             foldCtx,
             executionInfo,
+            request.responseStream(),
             resultListener
         );
         planExecutor.esql(