Working stream without streamed pages

Author: ivancea

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

@@ -97,7 +97,7 @@ public static Iterator<? extends ToXContent> columnarValues(List<ColumnInfoImpl>
     }
 
     /** Returns a row based representation of the values in the given pages (described by the column infos). */
-    public static Iterator<? extends ToXContent> rowValues(List<ColumnInfoImpl> columns, List<Page> pages, boolean[] nullColumns) {
+    public static Iterator<? extends ToXContent> rowValues(List<ColumnInfoImpl> columns, Iterable<Page> pages, boolean[] nullColumns) {
         final BytesRef scratch = new BytesRef();
         return Iterators.flatMap(pages.iterator(), page -> {
             final int columnCount = columns.size();

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

@@ -17,7 +17,6 @@
 import org.elasticsearch.rest.action.RestCancellableNodeClient;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.esql.action.stream.EsqlQueryResponseStream;
-import org.elasticsearch.xpack.esql.action.stream.NonStreamingEsqlQueryResponseStream;
 
 import java.io.IOException;
 import java.util.List;
@@ -54,16 +53,14 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
 
     protected static RestChannelConsumer restChannelConsumer(EsqlQueryRequest esqlRequest, RestRequest request, NodeClient client) {
         final Boolean partialResults = request.paramAsBoolean("allow_partial_results", null);
-        final boolean stream = request.paramAsBoolean("stream", false);
+        final boolean shouldStream = request.paramAsBoolean("stream", false);
         if (partialResults != null) {
             esqlRequest.allowPartialResults(partialResults);
         }
         LOGGER.debug("Beginning execution of ESQL query.\nQuery string: [{}]", esqlRequest.query());
         // TODO: Create responseStream here, and add to the request object (?). See RestRepositoryVerifyIntegrityAction
         return channel -> {
-            final var responseStream = stream
-                ? EsqlQueryResponseStream.forMediaType(channel, request)
-                : new NonStreamingEsqlQueryResponseStream(channel, request, esqlRequest);
+            final var responseStream = EsqlQueryResponseStream.forMediaType(channel, request, esqlRequest, shouldStream);
             esqlRequest.responseStream(responseStream);
             RestCancellableNodeClient cancellableClient = new RestCancellableNodeClient(client, request.getHttpChannel());
             cancellableClient.execute(EsqlQueryAction.INSTANCE, esqlRequest, responseStream.completionListener());

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

@@ -7,13 +7,19 @@
 
 package org.elasticsearch.xpack.esql.action.stream;
 
+import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.rest.RestChannel;
+import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xpack.esql.action.ColumnInfoImpl;
+import org.elasticsearch.xpack.esql.action.EsqlExecutionInfo;
+import org.elasticsearch.xpack.esql.action.EsqlQueryRequest;
 import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
 import org.elasticsearch.xpack.esql.action.ResponseXContentUtils;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
@@ -24,52 +30,114 @@
  * Default, XContent response stream.
  */
 public class DefaultEsqlQueryResponseStream extends EsqlQueryResponseStream {
-    DefaultEsqlQueryResponseStream(RestChannel restChannel, ToXContent.Params params) {
-        super(restChannel, params);
+    /**
+     * Columns, stored on {@link #doStartResponse}, and used later when sending pages.
+     */
+    @Nullable
+    private List<ColumnInfoImpl> columns;
+
+    private boolean dropNullColumns;
+
+    DefaultEsqlQueryResponseStream(RestChannel restChannel, RestRequest restRequest, EsqlQueryRequest esqlRequest) throws IOException {
+        super(restChannel, restRequest, esqlRequest);
+
+        this.dropNullColumns = restRequest.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false);
     }
 
     @Override
     protected boolean canBeStreamed() {
-        boolean dropNullColumns = params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false);
-        return dropNullColumns == false;
+        return dropNullColumns == false && esqlRequest.columnar() == false;
     }
 
     @Override
     protected void doStartResponse(List<ColumnInfoImpl> columns) {
-        assert params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false) == false : "this method doesn't support dropping null columns";
+        assert dropNullColumns == false : "this method doesn't support dropping null columns";
 
-        var content = new ArrayList<Iterator<? extends ToXContent>>(1);
+        var content = new ArrayList<Iterator<? extends ToXContent>>(3);
 
+        content.add(ChunkedToXContentHelper.startObject());
         content.add(ResponseXContentUtils.allColumns(columns, "columns"));
 
+        // Start the values array, to be filled in
+        content.add(ChunkedToXContentHelper.startArray("values"));
+
         sendChunks(content);
     }
 
     @Override
     protected void doSendPages(Iterable<Page> pages) {
-        // TODO: Implement
+        assert columns != null : "columns must be set before sending pages";
+
+        sendChunks(List.of(ResponseXContentUtils.rowValues(columns, pages, null)));
     }
 
     @Override
     protected void doFinishResponse(EsqlQueryResponse response) {
-        // TODO: Implement
+        var content = new ArrayList<Iterator<? extends ToXContent>>(10);
+
+        // End the values array
+        content.add(ChunkedToXContentHelper.endArray());
+
+        var executionInfo = response.getExecutionInfo();
+        if (executionInfo != null) {
+            if (executionInfo.overallTook() != null) {
+                content.add(
+                    ChunkedToXContentHelper.chunk(
+                        (builder, p) -> builder //
+                            .field("took", executionInfo.overallTook().millis())
+                            .field(EsqlExecutionInfo.IS_PARTIAL_FIELD.getPreferredName(), executionInfo.isPartial())
+                    )
+                );
+            }
+            if (executionInfo.hasMetadataToReport()) {
+                content.add(ChunkedToXContentHelper.field("_clusters", executionInfo, restRequest));
+            }
+        }
+        content.add(
+            ChunkedToXContentHelper.chunk(
+                (builder, p) -> builder //
+                    .field("documents_found", response.documentsFound())
+                    .field("values_loaded", response.valuesLoaded())
+            )
+        );
+
+        var profile = response.profile();
+        if (profile != null) {
+            content.add(ChunkedToXContentHelper.startObject("profile"));
+            content.add(ChunkedToXContentHelper.chunk((b, p) -> {
+                if (executionInfo != null) {
+                    b.field("query", executionInfo.overallTimeSpan());
+                    b.field("planning", executionInfo.planningTimeSpan());
+                }
+                return b;
+            }));
+            content.add(ChunkedToXContentHelper.array("drivers", profile.drivers().iterator(), restRequest));
+            content.add(ChunkedToXContentHelper.array("plans", profile.plans().iterator()));
+            content.add(ChunkedToXContentHelper.endObject());
+        }
+
+        content.add(ChunkedToXContentHelper.endObject());
+
+        sendChunks(content);
     }
 
     @Override
     protected void doSendEverything(EsqlQueryResponse response) {
-        boolean dropNullColumns = params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false);
-
         // TODO: Close the response
         // final Releasable releasable = releasableFromResponse(esqlResponse);
 
         // TODO: Instead of sendChunks, implement a flush() to attach the response to it? Or pass the response to the methods
+        // TODO: Or make "doX" methods return an Iterator<? extends ToXContent> and then concat them all together
 
+        boolean[] nullColumns = null;
         if (dropNullColumns) {
-            sendStartResponseDroppingNullColumns(response.columns(), response.pages());
+            nullColumns = nullColumns(response.columns(), response.pages());
+            sendStartResponseDroppingNullColumns(response.columns(), nullColumns);
         } else {
             doStartResponse(response.columns());
         }
-        doSendPages(response.pages());
+        // doSendPages doesn't work with nullColumns or columnar, so we generate them here directly
+        sendChunks(List.of(ResponseXContentUtils.columnValues(response.columns(), response.pages(), esqlRequest.columnar(), nullColumns)));
         doFinishResponse(response);
     }
 
@@ -78,12 +146,12 @@ protected void doHandleException(Exception e) {
         // TODO: Implement this
     }
 
-    private void sendStartResponseDroppingNullColumns(List<ColumnInfoImpl> columns, List<Page> pages) {
-        assert params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false) : "this method should only be called when dropping null columns";
+    private void sendStartResponseDroppingNullColumns(List<ColumnInfoImpl> columns, boolean[] nullColumns) {
+        assert dropNullColumns : "this method should only be called when dropping null columns";
 
-        var content = new ArrayList<Iterator<? extends ToXContent>>(2);
+        var content = new ArrayList<Iterator<? extends ToXContent>>(3);
 
-        boolean[] nullColumns = nullColumns(columns, pages);
+        content.add(ChunkedToXContentHelper.startObject());
         content.add(ResponseXContentUtils.allColumns(columns, "all_columns"));
         content.add(ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns"));
 
@@ -106,183 +174,4 @@ private boolean allColumnsAreNull(List<Page> pages, int c) {
         }
         return true;
     }
-
-    // TODO:
-    /*
-    On start, check if it's columnar or something else that disable streaming.
-    If streaming is ok, create the StreamingXContentResponse and write the opening fragment.
-    If not, leave it null.
-
-    On dispose, check if the StreamingXContentResponse is null or not.
-    If it's null, send the full response directly.
-    If it's not null, write the final fragment and close the StreamingXContentResponse.
-    */
-
-    /*
-    public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
-        boolean dropNullColumns = params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false);
-        boolean[] nullColumns = dropNullColumns ? nullColumns() : null;
-
-        var content = new ArrayList<Iterator<? extends ToXContent>>(25);
-        content.add(ChunkedToXContentHelper.startObject());
-        if (isAsync) {
-            content.add(ChunkedToXContentHelper.chunk((builder, p) -> {
-                if (asyncExecutionId != null) {
-                    builder.field("id", asyncExecutionId);
-                }
-                builder.field("is_running", isRunning);
-                return builder;
-            }));
-        }
-        if (executionInfo != null && executionInfo.overallTook() != null) {
-            content.add(
-                ChunkedToXContentHelper.chunk(
-                    (builder, p) -> builder //
-                        .field("took", executionInfo.overallTook().millis())
-                        .field(EsqlExecutionInfo.IS_PARTIAL_FIELD.getPreferredName(), executionInfo.isPartial())
-                )
-            );
-        }
-        content.add(
-            ChunkedToXContentHelper.chunk(
-                (builder, p) -> builder //
-                    .field("documents_found", documentsFound)
-                    .field("values_loaded", valuesLoaded)
-            )
-        );
-        if (dropNullColumns) {
-            content.add(ResponseXContentUtils.allColumns(columns, "all_columns"));
-            content.add(ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns"));
-        } else {
-            content.add(ResponseXContentUtils.allColumns(columns, "columns"));
-        }
-        content.add(
-            ChunkedToXContentHelper.array("values", ResponseXContentUtils.columnValues(this.columns, this.pages, columnar, nullColumns))
-        );
-        if (executionInfo != null && executionInfo.hasMetadataToReport()) {
-            content.add(ChunkedToXContentHelper.field("_clusters", executionInfo, params));
-        }
-        if (profile != null) {
-            content.add(ChunkedToXContentHelper.startObject("profile"));
-            content.add(ChunkedToXContentHelper.chunk((b, p) -> {
-                if (executionInfo != null) {
-                    b.field("query", executionInfo.overallTimeSpan());
-                    b.field("planning", executionInfo.planningTimeSpan());
-                }
-                return b;
-            }));
-            content.add(ChunkedToXContentHelper.array("drivers", profile.drivers.iterator(), params));
-            content.add(ChunkedToXContentHelper.array("plans", profile.plans.iterator()));
-            content.add(ChunkedToXContentHelper.endObject());
-        }
-        content.add(ChunkedToXContentHelper.endObject());
-
-        return Iterators.concat(content.toArray(Iterator[]::new));
-    }*/
-
-    /*void startResponse(Releasable releasable) throws IOException {
-        assert hasReferences();
-        assert streamingXContentResponse == null;
-        streamingXContentResponse = new StreamingXContentResponse(restChannel, restChannel.request(), () -> {});
-        streamingXContentResponse.writeFragment(
-            p0 -> ChunkedToXContentHelper.chunk((b, p) -> b.startObject().startArray("log")),
-            releasable
-        );
-    }
-
-    void writePage(RepositoryVerifyIntegrityResponseChunk chunk, Releasable releasable) {
-        assert hasReferences();
-        assert streamingXContentResponse != null;
-
-        if (chunk.type() == RepositoryVerifyIntegrityResponseChunk.Type.ANOMALY) {
-            anomalyCount.incrementAndGet();
-        }
-        streamingXContentResponse.writeFragment(
-            p0 -> ChunkedToXContentHelper.chunk((b, p) -> b.startObject().value(chunk, p).endObject()),
-            releasable
-        );
-    }
-
-    void writeChunk(RepositoryVerifyIntegrityResponseChunk chunk, Releasable releasable) {
-        assert hasReferences();
-        assert streamingXContentResponse != null;
-
-        if (chunk.type() == RepositoryVerifyIntegrityResponseChunk.Type.ANOMALY) {
-            anomalyCount.incrementAndGet();
-        }
-        streamingXContentResponse.writeFragment(
-            p0 -> ChunkedToXContentHelper.chunk((b, p) -> b.startObject().value(chunk, p).endObject()),
-            releasable
-        );
-    }
-
-    @Override
-    protected void closeInternal() {
-        try {
-            assert finalResultListener.isDone();
-            finalResultListener.addListener(new ActionListener<>() {
-                @Override
-                public void onResponse(RepositoryVerifyIntegrityResponse repositoryVerifyIntegrityResponse) {
-                    // success - finish the response with the final results
-                    assert streamingXContentResponse != null;
-                    streamingXContentResponse.writeFragment(
-                        p0 -> ChunkedToXContentHelper.chunk(
-                            (b, p) -> b.endArray()
-                                .startObject("results")
-                                .field("status", repositoryVerifyIntegrityResponse.finalTaskStatus())
-                                .field("final_repository_generation", repositoryVerifyIntegrityResponse.finalRepositoryGeneration())
-                                .field("total_anomalies", anomalyCount.get())
-                                .field(
-                                    "result",
-                                    anomalyCount.get() == 0
-                                        ? repositoryVerifyIntegrityResponse
-                                            .originalRepositoryGeneration() == repositoryVerifyIntegrityResponse.finalRepositoryGeneration()
-                                                ? "pass"
-                                                : "inconclusive due to concurrent writes"
-                                        : "fail"
-                                )
-                                .endObject()
-                                .endObject()
-                        ),
-                        () -> {}
-                    );
-                }
-
-                @Override
-                public void onFailure(Exception e) {
-                    if (streamingXContentResponse != null) {
-                        // failure after starting the response - finish the response with a rendering of the final exception
-                        streamingXContentResponse.writeFragment(
-                            p0 -> ChunkedToXContentHelper.chunk(
-                                (b, p) -> b.endArray()
-                                    .startObject("exception")
-                                    .value((bb, pp) -> ElasticsearchException.generateFailureXContent(bb, pp, e, true))
-                                    .field("status", ExceptionsHelper.status(e))
-                                    .endObject()
-                                    .endObject()
-                            ),
-                            () -> {}
-                        );
-                    } else {
-                        // didn't even get as far as starting to stream the response, must have hit an early exception (e.g. repo not found)
-                        // so we can return this exception directly.
-                        try {
-                            restChannel.sendResponse(new RestResponse(restChannel, e));
-                        } catch (IOException e2) {
-                            e.addSuppressed(e2);
-                            logger.error("error building error response", e);
-                            assert false : e; // shouldn't actually throw anything here
-                            restChannel.request().getHttpChannel().close();
-                        }
-                    }
-                }
-            });
-        } finally {
-            Releasables.closeExpectNoException(streamingXContentResponse);
-        }
-    }
-
-    public ActionListener<RepositoryVerifyIntegrityResponse> getCompletionListener() {
-        return completionListener;
-    }*/
 }