Skip to content

Commit 56d5009

Browse files
idegtiarenkoidegtiarenko
authored
Add query plans to profile output (#128828)
1 parent 8b62a55 commit 56d5009

File tree

14 files changed

+208
-94
lines changed

14 files changed

+208
-94
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -322,6 +322,7 @@ static TransportVersion def(int id) {
322322
public static final TransportVersion CLUSTER_STATE_PROJECTS_SETTINGS = def(9_108_0_00);
323323
public static final TransportVersion ML_INFERENCE_ELASTIC_DENSE_TEXT_EMBEDDINGS_ADDED = def(9_109_00_0);
324324
public static final TransportVersion ML_INFERENCE_COHERE_API_VERSION = def(9_110_0_00);
325+
public static final TransportVersion ESQL_PROFILE_INCLUDE_PLAN = def(9_111_0_00);
325326

326327
/*
327328
* 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
@@ -2713,9 +2713,11 @@ public static Request newXContentRequest(HttpMethod method, String endpoint, ToX
27132713
}
27142714

27152715
protected static MapMatcher getProfileMatcher() {
2716-
return matchesMap().entry("query", instanceOf(Map.class))
2716+
return matchesMap() //
2717+
.entry("query", instanceOf(Map.class))
27172718
.entry("planning", instanceOf(Map.class))
2718-
.entry("drivers", instanceOf(List.class));
2719+
.entry("drivers", instanceOf(List.class))
2720+
.entry("plans", instanceOf(List.class));
27192721
}
27202722

27212723
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::readFrom));
92+
public static DriverCompletionInfo readFrom(StreamInput in) throws IOException {
93+
return new DriverCompletionInfo(
94+
in.readVLong(),
95+
in.readVLong(),
96+
in.readCollectionAsImmutableList(DriverProfile::readFrom),
97+
in.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_INCLUDE_PLAN)
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)) {
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
@@ -112,7 +112,9 @@ private void testQuery(Double percent, String query, int documentsFound, boolean
112112
matchesMap().entry("documents_found", documentsFound)
113113
.entry(
114114
"profile",
115-
matchesMap().entry("drivers", instanceOf(List.class))
115+
matchesMap() //
116+
.entry("drivers", instanceOf(List.class))
117+
.entry("plans", instanceOf(List.class))
116118
.entry("planning", matchesMap().extraOk())
117119
.entry("query", matchesMap().extraOk())
118120
)

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

Lines changed: 12 additions & 2 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;
@@ -279,6 +280,7 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
279280
return b;
280281
}));
281282
content.add(ChunkedToXContentHelper.array("drivers", profile.drivers.iterator(), params));
283+
content.add(ChunkedToXContentHelper.array("plans", profile.plans.iterator()));
282284
content.add(ChunkedToXContentHelper.endObject());
283285
}
284286
content.add(ChunkedToXContentHelper.endObject());
@@ -387,15 +389,23 @@ public EsqlResponse responseInternal() {
387389
return esqlResponse;
388390
}
389391

390-
public record Profile(List<DriverProfile> drivers) implements Writeable {
392+
public record Profile(List<DriverProfile> drivers, List<PlanProfile> plans) implements Writeable {
391393

392394
public static Profile readFrom(StreamInput in) throws IOException {
393-
return new Profile(in.readCollectionAsImmutableList(DriverProfile::readFrom));
395+
return new Profile(
396+
in.readCollectionAsImmutableList(DriverProfile::readFrom),
397+
in.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_INCLUDE_PLAN)
398+
? in.readCollectionAsImmutableList(PlanProfile::readFrom)
399+
: List.of()
400+
);
394401
}
395402

396403
@Override
397404
public void writeTo(StreamOutput out) throws IOException {
398405
out.writeCollection(drivers);
406+
if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_INCLUDE_PLAN)) {
407+
out.writeCollection(plans);
408+
}
399409
}
400410
}
401411
}

‎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
@@ -59,10 +59,10 @@ final class ComputeResponse extends TransportResponse {
5959

6060
ComputeResponse(StreamInput in) throws IOException {
6161
if (in.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED)) {
62-
completionInfo = new DriverCompletionInfo(in);
62+
completionInfo = DriverCompletionInfo.readFrom(in);
6363
} else if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
6464
if (in.readBoolean()) {
65-
completionInfo = new DriverCompletionInfo(0, 0, in.readCollectionAsImmutableList(DriverProfile::readFrom));
65+
completionInfo = new DriverCompletionInfo(0, 0, in.readCollectionAsImmutableList(DriverProfile::readFrom), List.of());
6666
} else {
6767
completionInfo = DriverCompletionInfo.EMPTY;
6868
}
@@ -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)