Skip to content

Commit 5f43af5

Browse files
idegtiarenkoidegtiarenko
authored
[8.19] Add query plans to profile output (#130181)
1 parent d11943a commit 5f43af5

File tree

14 files changed

+210
-121
lines changed

14 files changed

+210
-121
lines changed

‎server/src/main/java/org/elasticsearch/TransportVersions.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,7 @@ static TransportVersion def(int id) {
254254
public static final TransportVersion ML_INFERENCE_ELASTIC_DENSE_TEXT_EMBEDDINGS_ADDED_8_19 = def(8_841_0_59);
255255
public static final TransportVersion ML_INFERENCE_COHERE_API_VERSION_8_19 = def(8_841_0_60);
256256
public static final TransportVersion ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED_8_19 = def(8_841_0_61);
257+
public static final TransportVersion ESQL_PROFILE_INCLUDE_PLAN_8_19 = def(8_841_0_62);
257258

258259
/*
259260
* STOP! READ THIS FIRST! No, really,

‎test/framework/src/main/java/org/elasticsearch/test/rest/ESRestTestCase.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2699,9 +2699,11 @@ public static Request newXContentRequest(HttpMethod method, String endpoint, ToX
26992699
}
27002700

27012701
protected static MapMatcher getProfileMatcher() {
2702-
return matchesMap().entry("query", instanceOf(Map.class))
2702+
return matchesMap() //
2703+
.entry("query", instanceOf(Map.class))
27032704
.entry("planning", instanceOf(Map.class))
2704-
.entry("drivers", instanceOf(List.class));
2705+
.entry("drivers", instanceOf(List.class))
2706+
.entry("plans", instanceOf(List.class));
27052707
}
27062708

27072709
protected static MapMatcher getResultMatcher(boolean includeMetadata, boolean includePartial, boolean includeDocumentsFound) {

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/DriverCompletionInfo.java‎

Lines changed: 45 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
package org.elasticsearch.compute.operator;
99

10+
import org.elasticsearch.TransportVersions;
1011
import org.elasticsearch.common.io.stream.StreamInput;
1112
import org.elasticsearch.common.io.stream.StreamOutput;
1213
import org.elasticsearch.common.io.stream.Writeable;
@@ -24,23 +25,34 @@
2425
* <strong>roughly</strong> the number of documents times the number of
2526
* fields per document. Except {@code null} values don't count.
2627
* And multivalued fields count as many times as there are values.
27-
* @param collectedProfiles {@link DriverProfile}s from each driver. These are fairly cheap to build but
28+
* @param driverProfiles {@link DriverProfile}s from each driver. These are fairly cheap to build but
2829
* not free so this will be empty if the {@code profile} option was not set in
2930
* the request.
3031
*/
31-
public record DriverCompletionInfo(long documentsFound, long valuesLoaded, List<DriverProfile> collectedProfiles) implements Writeable {
32+
public record DriverCompletionInfo(
33+
long documentsFound,
34+
long valuesLoaded,
35+
List<DriverProfile> driverProfiles,
36+
List<PlanProfile> planProfiles
37+
) implements Writeable {
3238

3339
/**
3440
* Completion info we use when we didn't properly complete any drivers.
3541
* Usually this is returned with an error, but it's also used when receiving
3642
* responses from very old nodes.
3743
*/
38-
public static final DriverCompletionInfo EMPTY = new DriverCompletionInfo(0, 0, List.of());
44+
public static final DriverCompletionInfo EMPTY = new DriverCompletionInfo(0, 0, List.of(), List.of());
3945

4046
/**
4147
* Build a {@link DriverCompletionInfo} for many drivers including their profile output.
4248
*/
43-
public static DriverCompletionInfo includingProfiles(List<Driver> drivers) {
49+
public static DriverCompletionInfo includingProfiles(
50+
List<Driver> drivers,
51+
String description,
52+
String clusterName,
53+
String nodeName,
54+
String planTree
55+
) {
4456
long documentsFound = 0;
4557
long valuesLoaded = 0;
4658
List<DriverProfile> collectedProfiles = new ArrayList<>(drivers.size());
@@ -52,7 +64,12 @@ public static DriverCompletionInfo includingProfiles(List<Driver> drivers) {
5264
}
5365
collectedProfiles.add(p);
5466
}
55-
return new DriverCompletionInfo(documentsFound, valuesLoaded, collectedProfiles);
67+
return new DriverCompletionInfo(
68+
documentsFound,
69+
valuesLoaded,
70+
collectedProfiles,
71+
List.of(new PlanProfile(description, clusterName, nodeName, planTree))
72+
);
5673
}
5774

5875
/**
@@ -69,49 +86,63 @@ public static DriverCompletionInfo excludingProfiles(List<Driver> drivers) {
6986
valuesLoaded += o.valuesLoaded();
7087
}
7188
}
72-
return new DriverCompletionInfo(documentsFound, valuesLoaded, List.of());
89+
return new DriverCompletionInfo(documentsFound, valuesLoaded, List.of(), List.of());
7390
}
7491

75-
public DriverCompletionInfo(StreamInput in) throws IOException {
76-
this(in.readVLong(), in.readVLong(), in.readCollectionAsImmutableList(DriverProfile::new));
92+
public static DriverCompletionInfo readFrom(StreamInput in) throws IOException {
93+
return new DriverCompletionInfo(
94+
in.readVLong(),
95+
in.readVLong(),
96+
in.readCollectionAsImmutableList(DriverProfile::new),
97+
in.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_INCLUDE_PLAN_8_19)
98+
? in.readCollectionAsImmutableList(PlanProfile::readFrom)
99+
: List.of()
100+
);
77101
}
78102

79103
@Override
80104
public void writeTo(StreamOutput out) throws IOException {
81105
out.writeVLong(documentsFound);
82106
out.writeVLong(valuesLoaded);
83-
out.writeCollection(collectedProfiles, (o, v) -> v.writeTo(o));
107+
out.writeCollection(driverProfiles);
108+
if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_INCLUDE_PLAN_8_19)) {
109+
out.writeCollection(planProfiles);
110+
}
84111
}
85112

86113
public static class Accumulator {
87114
private long documentsFound;
88115
private long valuesLoaded;
89-
private final List<DriverProfile> collectedProfiles = new ArrayList<>();
116+
private final List<DriverProfile> driverProfiles = new ArrayList<>();
117+
private final List<PlanProfile> planProfiles = new ArrayList<>();
90118

91119
public void accumulate(DriverCompletionInfo info) {
92120
this.documentsFound += info.documentsFound;
93121
this.valuesLoaded += info.valuesLoaded;
94-
this.collectedProfiles.addAll(info.collectedProfiles);
122+
this.driverProfiles.addAll(info.driverProfiles);
123+
this.planProfiles.addAll(info.planProfiles);
95124
}
96125

97126
public DriverCompletionInfo finish() {
98-
return new DriverCompletionInfo(documentsFound, valuesLoaded, collectedProfiles);
127+
return new DriverCompletionInfo(documentsFound, valuesLoaded, driverProfiles, planProfiles);
99128
}
100129
}
101130

102131
public static class AtomicAccumulator {
103132
private final AtomicLong documentsFound = new AtomicLong();
104133
private final AtomicLong valuesLoaded = new AtomicLong();
105134
private final List<DriverProfile> collectedProfiles = Collections.synchronizedList(new ArrayList<>());
135+
private final List<PlanProfile> planProfiles = Collections.synchronizedList(new ArrayList<>());
106136

107137
public void accumulate(DriverCompletionInfo info) {
108138
this.documentsFound.addAndGet(info.documentsFound);
109139
this.valuesLoaded.addAndGet(info.valuesLoaded);
110-
this.collectedProfiles.addAll(info.collectedProfiles);
140+
this.collectedProfiles.addAll(info.driverProfiles);
141+
this.planProfiles.addAll(info.planProfiles);
111142
}
112143

113144
public DriverCompletionInfo finish() {
114-
return new DriverCompletionInfo(documentsFound.get(), valuesLoaded.get(), collectedProfiles);
145+
return new DriverCompletionInfo(documentsFound.get(), valuesLoaded.get(), collectedProfiles, planProfiles);
115146
}
116147
}
117148
}

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/PlanProfile.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.compute.operator;
9+
10+
import org.elasticsearch.common.io.stream.StreamInput;
11+
import org.elasticsearch.common.io.stream.StreamOutput;
12+
import org.elasticsearch.common.io.stream.Writeable;
13+
import org.elasticsearch.xcontent.ToXContentObject;
14+
import org.elasticsearch.xcontent.XContentBuilder;
15+
16+
import java.io.IOException;
17+
18+
public record PlanProfile(String description, String clusterName, String nodeName, String planTree) implements Writeable, ToXContentObject {
19+
20+
public static PlanProfile readFrom(StreamInput in) throws IOException {
21+
return new PlanProfile(in.readString(), in.readString(), in.readString(), in.readString());
22+
}
23+
24+
@Override
25+
public void writeTo(StreamOutput out) throws IOException {
26+
out.writeString(description);
27+
out.writeString(clusterName);
28+
out.writeString(nodeName);
29+
out.writeString(planTree);
30+
}
31+
32+
@Override
33+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
34+
return builder.startObject()
35+
.field("description", description)
36+
.field("cluster_name", clusterName)
37+
.field("node_name", nodeName)
38+
.field("plan", planTree)
39+
.endObject();
40+
}
41+
}

‎x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/PushQueriesIT.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -317,7 +317,9 @@ private void testPushQuery(
317317
result,
318318
getResultMatcher(result).entry(
319319
"profile",
320-
matchesMap().entry("drivers", instanceOf(List.class))
320+
matchesMap() //
321+
.entry("drivers", instanceOf(List.class))
322+
.entry("plans", instanceOf(List.class))
321323
.entry("planning", matchesMap().extraOk())
322324
.entry("query", matchesMap().extraOk())
323325
),

‎x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/StoredFieldsSequentialIT.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,9 @@ private void testQuery(Double percent, String query, int documentsFound, boolean
111111
matchesMap().entry("documents_found", documentsFound)
112112
.entry(
113113
"profile",
114-
matchesMap().entry("drivers", instanceOf(List.class))
114+
matchesMap() //
115+
.entry("drivers", instanceOf(List.class))
116+
.entry("plans", instanceOf(List.class))
115117
.entry("planning", matchesMap().extraOk())
116118
.entry("query", matchesMap().extraOk())
117119
)

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

Lines changed: 13 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import org.elasticsearch.compute.data.BlockStreamInput;
2121
import org.elasticsearch.compute.data.Page;
2222
import org.elasticsearch.compute.operator.DriverProfile;
23+
import org.elasticsearch.compute.operator.PlanProfile;
2324
import org.elasticsearch.core.AbstractRefCounted;
2425
import org.elasticsearch.core.Nullable;
2526
import org.elasticsearch.core.Releasable;
@@ -121,7 +122,7 @@ static EsqlQueryResponse deserialize(BlockStreamInput in) throws IOException {
121122
long documentsFound = in.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED_8_19) ? in.readVLong() : 0;
122123
long valuesLoaded = in.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED_8_19) ? in.readVLong() : 0;
123124
if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
124-
profile = in.readOptionalWriteable(Profile::new);
125+
profile = in.readOptionalWriteable(Profile::readFrom);
125126
}
126127
boolean columnar = in.readBoolean();
127128
EsqlExecutionInfo executionInfo = null;
@@ -261,6 +262,7 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
261262
ob.field("planning", executionInfo.planningTimeSpan());
262263
}
263264
ob.array("drivers", profile.drivers.iterator(), ChunkedToXContentBuilder::append);
265+
ob.array("plans", profile.plans.iterator());
264266
}));
265267
}
266268
});
@@ -372,41 +374,23 @@ public EsqlResponse responseInternal() {
372374
return esqlResponse;
373375
}
374376

375-
public static class Profile implements Writeable {
376-
private final List<DriverProfile> drivers;
377+
public record Profile(List<DriverProfile> drivers, List<PlanProfile> plans) implements Writeable {
377378

378-
public Profile(List<DriverProfile> drivers) {
379-
this.drivers = drivers;
380-
}
381-
382-
public Profile(StreamInput in) throws IOException {
383-
this.drivers = in.readCollectionAsImmutableList(DriverProfile::new);
379+
public static Profile readFrom(StreamInput in) throws IOException {
380+
return new Profile(
381+
in.readCollectionAsImmutableList(DriverProfile::new),
382+
in.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_INCLUDE_PLAN_8_19)
383+
? in.readCollectionAsImmutableList(PlanProfile::readFrom)
384+
: List.of()
385+
);
384386
}
385387

386388
@Override
387389
public void writeTo(StreamOutput out) throws IOException {
388390
out.writeCollection(drivers);
389-
}
390-
391-
@Override
392-
public boolean equals(Object o) {
393-
if (this == o) {
394-
return true;
395-
}
396-
if (o == null || getClass() != o.getClass()) {
397-
return false;
391+
if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_INCLUDE_PLAN_8_19)) {
392+
out.writeCollection(plans);
398393
}
399-
Profile profile = (Profile) o;
400-
return Objects.equals(drivers, profile.drivers);
401-
}
402-
403-
@Override
404-
public int hashCode() {
405-
return Objects.hash(drivers);
406-
}
407-
408-
List<DriverProfile> drivers() {
409-
return drivers;
410394
}
411395
}
412396
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,10 @@ final class ComputeResponse extends TransportResponse {
6060
ComputeResponse(StreamInput in) throws IOException {
6161
super(in);
6262
if (in.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED_8_19)) {
63-
completionInfo = new DriverCompletionInfo(in);
63+
completionInfo = DriverCompletionInfo.readFrom(in);
6464
} else if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
6565
if (in.readBoolean()) {
66-
completionInfo = new DriverCompletionInfo(0, 0, in.readCollectionAsImmutableList(DriverProfile::new));
66+
completionInfo = new DriverCompletionInfo(0, 0, in.readCollectionAsImmutableList(DriverProfile::new), List.of());
6767
} else {
6868
completionInfo = DriverCompletionInfo.EMPTY;
6969
}
@@ -96,7 +96,7 @@ public void writeTo(StreamOutput out) throws IOException {
9696
completionInfo.writeTo(out);
9797
} else if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
9898
out.writeBoolean(true);
99-
out.writeCollection(completionInfo.collectedProfiles());
99+
out.writeCollection(completionInfo.driverProfiles());
100100
}
101101
if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_16_0)) {
102102
out.writeOptionalTimeValue(took);

0 commit comments

Comments
 (0)