Skip to content

Commit

Permalink
More conversions to ChunkedXContentBuilder
Browse files Browse the repository at this point in the history
  • Loading branch information
@thecoop
thecoop committed Oct 7, 2024
1 parent 1028c48 commit 5b0d0d0
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 105 deletions.
29 changes: 11 additions & 18 deletions .../src/main/java/org/elasticsearch/action/admin/cluster/reroute/ClusterRerouteResponse.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,17 @@
import org.elasticsearch.action.support.master.IsAcknowledgedSupplier;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.routing.allocation.RoutingExplanations;
import org.elasticsearch.common.collect.Iterators;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.logging.DeprecationCategory;
import org.elasticsearch.common.logging.DeprecationLogger;
import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
import org.elasticsearch.common.xcontent.ChunkedToXContent;
import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
import org.elasticsearch.core.RestApiVersion;
import org.elasticsearch.rest.action.search.RestSearchAction;
import org.elasticsearch.xcontent.ToXContent;

import java.io.IOException;
import java.util.Collections;
import java.util.Iterator;
import java.util.Objects;

Expand Down Expand Up @@ -96,20 +94,15 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerP
}

@Override
public Iterator<? extends ToXContent> toXContentChunkedV7(ToXContent.Params outerParams) {
return Iterators.concat(
Iterators.single((builder, params) -> builder.startObject().field(ACKNOWLEDGED_KEY, isAcknowledged())),
emitState(outerParams)
? ChunkedToXContentHelper.wrapWithObject("state", state.toXContentChunked(outerParams))
: Collections.emptyIterator(),
Iterators.single((builder, params) -> {
if (params.paramAsBoolean("explain", false)) {
explanations.toXContent(builder, params);
}

builder.endObject();
return builder;
})
);
public Iterator<? extends ToXContent> toXContentChunkedV7(ToXContent.Params params) {
return ChunkedToXContent.builder(params).object(b -> {
b.field(ACKNOWLEDGED_KEY, isAcknowledged());
if (emitState(params)) {
b.xContentObject("state", state);
}
if (params.paramAsBoolean("explain", false)) {
b.append(explanations);
}
});
}
}
10 changes: 10 additions & 0 deletions server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,10 @@ private void addChunk(ToXContent content) {
builder.add(Objects.requireNonNull(content));
}

public ToXContent.Params params() {
return params;
}

private void startObject() {
addChunk((b, p) -> b.startObject());
}
Expand Down Expand Up @@ -351,6 +355,12 @@ public ChunkedToXContentBuilder field(String name, ToXContent value) {
return this;
}

public ChunkedToXContentBuilder field(String name, ChunkedToXContent value) {
addChunk((b, p) -> b.field(name));
append(value);
return this;
}

public ChunkedToXContentBuilder field(String name, Object value) {
addChunk((b, p) -> b.field(name, value));
return this;
Expand Down
57 changes: 21 additions & 36 deletions server/src/main/java/org/elasticsearch/script/ScriptCacheStats.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,15 +21,8 @@
import java.util.Iterator;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;

import static org.elasticsearch.common.collect.Iterators.concat;
import static org.elasticsearch.common.collect.Iterators.flatMap;
import static org.elasticsearch.common.collect.Iterators.single;
import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.endArray;
import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.endObject;
import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.field;
import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.startArray;
import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.startObject;
import static org.elasticsearch.script.ScriptCacheStats.Fields.SCRIPT_CACHE_STATS;

// This class is deprecated in favor of ScriptStats and ScriptContextStats
Expand Down Expand Up @@ -76,35 +69,27 @@ public void writeTo(StreamOutput out) throws IOException {

@Override
public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
return concat(
startObject(SCRIPT_CACHE_STATS),
startObject(Fields.SUM),
general != null
? concat(
field(ScriptStats.Fields.COMPILATIONS, general.getCompilations()),
field(ScriptStats.Fields.CACHE_EVICTIONS, general.getCacheEvictions()),
field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, general.getCompilationLimitTriggered()),
endObject(),
endObject()
)
: concat(single((builder, params) -> {
var sum = sum();
return builder.field(ScriptStats.Fields.COMPILATIONS, sum.getCompilations())
.field(ScriptStats.Fields.CACHE_EVICTIONS, sum.getCacheEvictions())
.field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, sum.getCompilationLimitTriggered())
.endObject();
}), startArray(Fields.CONTEXTS), flatMap(context.keySet().stream().sorted().iterator(), ctx -> {
var stats = context.get(ctx);
return concat(
startObject(),
field(Fields.CONTEXT, ctx),
field(ScriptStats.Fields.COMPILATIONS, stats.getCompilations()),
field(ScriptStats.Fields.CACHE_EVICTIONS, stats.getCacheEvictions()),
field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, stats.getCompilationLimitTriggered()),
endObject()
Function<ScriptStats, ToXContent> statsFields = s -> (b, p) -> b.field(ScriptStats.Fields.COMPILATIONS, s.getCompilations())
.field(ScriptStats.Fields.CACHE_EVICTIONS, s.getCacheEvictions())
.field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, s.getCompilationLimitTriggered());

return ChunkedToXContent.builder(outerParams).object(SCRIPT_CACHE_STATS, sb -> {
if (general != null) {
sb.xContentObject(Fields.SUM, statsFields.apply(general));
} else {
sb.object(Fields.SUM, sum -> {
sum.append(statsFields.apply(sum()));
sum.array(
Fields.CONTEXTS,
context.entrySet().stream().sorted(Map.Entry.comparingByKey()).iterator(),
(eb, e) -> eb.object(ebo -> {
ebo.field(Fields.CONTEXT, e.getKey());
ebo.append(statsFields.apply(e.getValue()));
})
);
}), endArray(), endObject())
);
});
}
});
}

/**
Expand Down
77 changes: 26 additions & 51 deletions x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.xcontent.ChunkedToXContent;
import org.elasticsearch.common.xcontent.ChunkedToXContentBuilder;
import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
import org.elasticsearch.compute.data.BlockFactory;
import org.elasticsearch.compute.data.BlockStreamInput;
Expand All @@ -30,7 +29,6 @@
import org.elasticsearch.xpack.esql.core.type.DataType;

import java.io.IOException;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
Expand Down Expand Up @@ -186,58 +184,35 @@ public EsqlExecutionInfo getExecutionInfo() {
return executionInfo;
}

private Iterator<? extends ToXContent> asyncPropertiesOrEmpty() {
if (isAsync) {
return ChunkedToXContentHelper.singleChunk((builder, params) -> {
if (asyncExecutionId != null) {
builder.field("id", asyncExecutionId);
}
builder.field("is_running", isRunning);
return builder;
});
} else {
return Collections.emptyIterator();
}
}

@Override
public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
boolean dropNullColumns = params.paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false);
boolean[] nullColumns = dropNullColumns ? nullColumns() : null;

Iterator<ToXContent> tookTime;
if (executionInfo != null && executionInfo.overallTook() != null) {
tookTime = ChunkedToXContentHelper.singleChunk((builder, p) -> {
builder.field("took", executionInfo.overallTook().millis());
return builder;
});
} else {
tookTime = Collections.emptyIterator();
}
return ChunkedToXContent.builder(params).object(b -> {
boolean dropNullColumns = b.params().paramAsBoolean(DROP_NULL_COLUMNS_OPTION, false);
boolean[] nullColumns = dropNullColumns ? nullColumns() : null;

Iterator<? extends ToXContent> columnHeadings = dropNullColumns
? Iterators.concat(
ResponseXContentUtils.allColumns(columns, "all_columns"),
ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns")
)
: ResponseXContentUtils.allColumns(columns, "columns");
Iterator<? extends ToXContent> valuesIt = ResponseXContentUtils.columnValues(this.columns, this.pages, columnar, nullColumns);
Iterator<ToXContent> profileRender = profile == null
? List.<ToXContent>of().iterator()
: ChunkedToXContentHelper.field("profile", profile, params);
Iterator<ToXContent> executionInfoRender = executionInfo == null || executionInfo.isCrossClusterSearch() == false
? List.<ToXContent>of().iterator()
: ChunkedToXContentHelper.field("_clusters", executionInfo, params);
return Iterators.concat(
ChunkedToXContentHelper.startObject(),
asyncPropertiesOrEmpty(),
tookTime,
columnHeadings,
ChunkedToXContentHelper.array("values", valuesIt),
executionInfoRender,
profileRender,
ChunkedToXContentHelper.endObject()
);
if (isAsync) {
if (asyncExecutionId != null) {
b.field("id", asyncExecutionId);
}
b.field("is_running", isRunning);
}
if (executionInfo != null && executionInfo.overallTook() != null) {
b.field("took", executionInfo.overallTook().millis());
}
if (dropNullColumns) {
b.append(ResponseXContentUtils.allColumns(columns, "all_columns"))
.append(ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns"));
} else {
b.append(ResponseXContentUtils.allColumns(columns, "columns"));
}
b.array("values", ResponseXContentUtils.columnValues(this.columns, this.pages, columnar, nullColumns));
if (executionInfo != null && executionInfo.isCrossClusterSearch()) {
b.field("_clusters", executionInfo);
}
if (profile != null) {
b.field("profile", profile);
}
});
}

private boolean[] nullColumns() {
Expand Down

0 comments on commit 5b0d0d0

Please sign in to comment.