From 54ed9c550b5831bb4a2f1968b4ba8b49319c43fb Mon Sep 17 00:00:00 2001 From: "ievgen.degtiarenko" Date: Fri, 6 Jun 2025 10:36:04 +0200 Subject: [PATCH 1/3] Simplify EsqlQueryResponse --- .../xpack/esql/action/EsqlQueryResponse.java | 144 +++++++----------- 1 file changed, 56 insertions(+), 88 deletions(-) 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..37c6ae7621b01 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 @@ -34,6 +34,7 @@ import java.util.List; import java.util.Objects; import java.util.Optional; +import java.util.function.Supplier; import static org.elasticsearch.TransportVersions.ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED; @@ -122,7 +123,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 +225,68 @@ public EsqlExecutionInfo getExecutionInfo() { return executionInfo; } - private Iterator asyncPropertiesOrEmpty() { - if (isAsync) { - return ChunkedToXContentHelper.chunk((builder, params) -> { - 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()) - ); - } 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") - ) - : 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), + conditionalChunkedXContent(isAsync, () -> ChunkedToXContentHelper.chunk((builder, p) -> { + if (asyncExecutionId != null) { + builder.field("id", asyncExecutionId); + } + builder.field("is_running", isRunning); + return builder; + })), + conditionalChunkedXContent( + executionInfo != null && executionInfo.overallTook() != null, + () -> ChunkedToXContentHelper.chunk( + (builder, p) -> builder // + .field("took", executionInfo.overallTook().millis()) + .field(EsqlExecutionInfo.IS_PARTIAL_FIELD.getPreferredName(), executionInfo.isPartial()) + ) + ), + ChunkedToXContentHelper.chunk( + (builder, p) -> builder // + .field("documents_found", documentsFound) + .field("values_loaded", valuesLoaded) + ), + dropNullColumns + ? Iterators.concat( + ResponseXContentUtils.allColumns(columns, "all_columns"), + ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns") + ) + : ResponseXContentUtils.allColumns(columns, "columns"), + ChunkedToXContentHelper.array("values", ResponseXContentUtils.columnValues(this.columns, this.pages, columnar, nullColumns)), + conditionalChunkedXContent( + executionInfo != null && executionInfo.hasMetadataToReport(), + () -> ChunkedToXContentHelper.field("_clusters", executionInfo, params) + ), + conditionalChunkedXContent( + profile != null, + () -> 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() + ) + ), ChunkedToXContentHelper.endObject() ); } - private Iterator profileRenderer(ToXContent.Params params) { - if (profile == null) { - return Collections.emptyIterator(); - } - 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()); + private Iterator conditionalChunkedXContent( + boolean condition, + Supplier> chunkedXContent + ) { + return condition ? chunkedXContent.get() : Collections.emptyIterator(); } public boolean[] nullColumns() { @@ -396,41 +390,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; - } } } From 1aae75ab0004bcd22d40a5702ff5d846ed88a603 Mon Sep 17 00:00:00 2001 From: "ievgen.degtiarenko" Date: Fri, 6 Jun 2025 11:24:14 +0200 Subject: [PATCH 2/3] fix compilation --- .../xpack/esql/action/EsqlQueryResponseProfileTests.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From ad0a3d0eb4b689d5231083101cee1ff4f19e4ce9 Mon Sep 17 00:00:00 2001 From: "ievgen.degtiarenko" Date: Thu, 12 Jun 2025 08:44:40 +0200 Subject: [PATCH 3/3] use array list instead --- .../xpack/esql/action/EsqlQueryResponse.java | 85 +++++++++---------- 1 file changed, 41 insertions(+), 44 deletions(-) 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 37c6ae7621b01..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,12 +29,11 @@ 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; import java.util.Optional; -import java.util.function.Supplier; import static org.elasticsearch.TransportVersions.ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED; @@ -226,67 +225,65 @@ public EsqlExecutionInfo getExecutionInfo() { } @Override + @SuppressWarnings("unchecked") public Iterator toXContentChunked(ToXContent.Params params) { boolean dropNullColumns = params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false); boolean[] nullColumns = dropNullColumns ? nullColumns() : null; - return Iterators.concat( - ChunkedToXContentHelper.startObject(), - conditionalChunkedXContent(isAsync, () -> ChunkedToXContentHelper.chunk((builder, p) -> { + var content = new ArrayList>(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; - })), - conditionalChunkedXContent( - executionInfo != null && executionInfo.overallTook() != null, - () -> ChunkedToXContentHelper.chunk( + })); + } + 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) - ), - dropNullColumns - ? Iterators.concat( - ResponseXContentUtils.allColumns(columns, "all_columns"), - ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns") - ) - : ResponseXContentUtils.allColumns(columns, "columns"), - ChunkedToXContentHelper.array("values", ResponseXContentUtils.columnValues(this.columns, this.pages, columnar, nullColumns)), - conditionalChunkedXContent( - executionInfo != null && executionInfo.hasMetadataToReport(), - () -> ChunkedToXContentHelper.field("_clusters", executionInfo, params) - ), - conditionalChunkedXContent( - profile != null, - () -> 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() - ) - ), - ChunkedToXContentHelper.endObject() + ) ); - } + 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.endObject()); + } + content.add(ChunkedToXContentHelper.endObject()); - private Iterator conditionalChunkedXContent( - boolean condition, - Supplier> chunkedXContent - ) { - return condition ? chunkedXContent.get() : Collections.emptyIterator(); + return Iterators.concat(content.toArray(Iterator[]::new)); } public boolean[] nullColumns() {