diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java index c005414704d0f..d28652c2a0162 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java @@ -29,7 +29,7 @@ import org.elasticsearch.xpack.esql.core.type.DataType; import java.io.IOException; -import java.util.Collections; +import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.Objects; @@ -122,7 +122,7 @@ static EsqlQueryResponse deserialize(BlockStreamInput in) throws IOException { long documentsFound = in.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED) ? in.readVLong() : 0; long valuesLoaded = in.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED) ? in.readVLong() : 0; if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) { - profile = in.readOptionalWriteable(Profile::new); + profile = in.readOptionalWriteable(Profile::readFrom); } boolean columnar = in.readBoolean(); EsqlExecutionInfo executionInfo = null; @@ -224,75 +224,66 @@ public EsqlExecutionInfo getExecutionInfo() { return executionInfo; } - private Iterator asyncPropertiesOrEmpty() { + @Override + @SuppressWarnings("unchecked") + public Iterator toXContentChunked(ToXContent.Params params) { + boolean dropNullColumns = params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false); + boolean[] nullColumns = dropNullColumns ? nullColumns() : null; + + var content = new ArrayList>(25); + content.add(ChunkedToXContentHelper.startObject()); if (isAsync) { - return ChunkedToXContentHelper.chunk((builder, params) -> { + content.add(ChunkedToXContentHelper.chunk((builder, p) -> { if (asyncExecutionId != null) { builder.field("id", asyncExecutionId); } builder.field("is_running", isRunning); return builder; - }); - } else { - return Collections.emptyIterator(); + })); } - } - - @Override - public Iterator toXContentChunked(ToXContent.Params params) { - boolean dropNullColumns = params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false); - boolean[] nullColumns = dropNullColumns ? nullColumns() : null; - - Iterator tookTime; if (executionInfo != null && executionInfo.overallTook() != null) { - tookTime = 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("took", executionInfo.overallTook().millis()) + .field(EsqlExecutionInfo.IS_PARTIAL_FIELD.getPreferredName(), executionInfo.isPartial()) + ) ); - } else { - tookTime = Collections.emptyIterator(); } - - Iterator meta = ChunkedToXContentHelper.chunk((builder, p) -> { - builder.field("documents_found", documentsFound); - builder.field("values_loaded", valuesLoaded); - return builder; - }); - - Iterator columnHeadings = dropNullColumns - ? Iterators.concat( - ResponseXContentUtils.allColumns(columns, "all_columns"), - ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns") + content.add( + ChunkedToXContentHelper.chunk( + (builder, p) -> builder // + .field("documents_found", documentsFound) + .field("values_loaded", valuesLoaded) ) - : ResponseXContentUtils.allColumns(columns, "columns"); - Iterator valuesIt = ResponseXContentUtils.columnValues(this.columns, this.pages, columnar, nullColumns); - Iterator executionInfoRender = executionInfo != null && executionInfo.hasMetadataToReport() - ? ChunkedToXContentHelper.field("_clusters", executionInfo, params) - : Collections.emptyIterator(); - return Iterators.concat( - ChunkedToXContentHelper.startObject(), - asyncPropertiesOrEmpty(), - tookTime, - meta, - columnHeadings, - ChunkedToXContentHelper.array("values", valuesIt), - executionInfoRender, - profileRenderer(params), - ChunkedToXContentHelper.endObject() ); - } - - private Iterator profileRenderer(ToXContent.Params params) { - if (profile == null) { - return Collections.emptyIterator(); + if (dropNullColumns) { + content.add(ResponseXContentUtils.allColumns(columns, "all_columns")); + content.add(ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns")); + } else { + content.add(ResponseXContentUtils.allColumns(columns, "columns")); } - return Iterators.concat(ChunkedToXContentHelper.startObject("profile"), ChunkedToXContentHelper.chunk((b, p) -> { - if (executionInfo != null) { - b.field("query", executionInfo.overallTimeSpan()); - b.field("planning", executionInfo.planningTimeSpan()); - } - return b; - }), ChunkedToXContentHelper.array("drivers", profile.drivers.iterator(), params), ChunkedToXContentHelper.endObject()); + 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.endObject()); + } + content.add(ChunkedToXContentHelper.endObject()); + + return Iterators.concat(content.toArray(Iterator[]::new)); } public boolean[] nullColumns() { @@ -396,41 +387,15 @@ public EsqlResponse responseInternal() { return esqlResponse; } - public static class Profile implements Writeable { - private final List drivers; - - public Profile(List drivers) { - this.drivers = drivers; - } + public record Profile(List drivers) implements Writeable { - public Profile(StreamInput in) throws IOException { - this.drivers = in.readCollectionAsImmutableList(DriverProfile::readFrom); + public static Profile readFrom(StreamInput in) throws IOException { + return new Profile(in.readCollectionAsImmutableList(DriverProfile::readFrom)); } @Override public void writeTo(StreamOutput out) throws IOException { out.writeCollection(drivers); } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - Profile profile = (Profile) o; - return Objects.equals(drivers, profile.drivers); - } - - @Override - public int hashCode() { - return Objects.hash(drivers); - } - - List drivers() { - return drivers; - } } } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponseProfileTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponseProfileTests.java index a0357c2393971..97407768b9347 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponseProfileTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponseProfileTests.java @@ -21,7 +21,7 @@ public class EsqlQueryResponseProfileTests extends AbstractWireSerializingTestCase { @Override protected Writeable.Reader instanceReader() { - return EsqlQueryResponse.Profile::new; + return EsqlQueryResponse.Profile::readFrom; } @Override