Skip to content

Commit f313956

Browse files
alex268alex268
authored
Merge pull request #555 from alex268/master
Updated QueryStats class in QueryClient
2 parents fe31231 + 44cdc71 commit f313956

File tree

4 files changed

+453
-60
lines changed

4 files changed

+453
-60
lines changed

‎query/src/main/java/tech/ydb/query/impl/TableClientImpl.java‎

Lines changed: 4 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import tech.ydb.query.QueryStream;
2323
import tech.ydb.query.result.QueryInfo;
2424
import tech.ydb.query.result.QueryResultPart;
25+
import tech.ydb.query.result.QueryStats;
2526
import tech.ydb.query.settings.ExecuteQuerySettings;
2627
import tech.ydb.query.settings.QueryStatsMode;
2728
import tech.ydb.table.Session;
@@ -30,17 +31,10 @@
3031
import tech.ydb.table.impl.BaseSession;
3132
import tech.ydb.table.query.DataQueryResult;
3233
import tech.ydb.table.query.Params;
33-
import tech.ydb.table.query.stats.CompilationStats;
34-
import tech.ydb.table.query.stats.OperationStats;
35-
import tech.ydb.table.query.stats.QueryPhaseStats;
36-
import tech.ydb.table.query.stats.QueryStats;
37-
import tech.ydb.table.query.stats.TableAccessStats;
3834
import tech.ydb.table.rpc.TableRpc;
3935
import tech.ydb.table.rpc.grpc.GrpcTableRpc;
4036
import tech.ydb.table.settings.ExecuteDataQuerySettings;
4137

42-
import static java.util.stream.Collectors.toList;
43-
4438
/**
4539
*
4640
* @author Aleksandr Gorshenin
@@ -100,53 +94,6 @@ private YdbQuery.TransactionControl mapTxControl(YdbTable.TransactionControl tc)
10094
return TxControl.txModeCtrl(TxMode.NONE, tc.getCommitTx());
10195
}
10296

103-
private static QueryStats queryStats(tech.ydb.query.result.QueryStats stats) {
104-
if (stats == null) {
105-
return null;
106-
}
107-
return new QueryStats(
108-
stats.getPhases().stream().map(qp -> queryPhaseStats(qp)).collect(toList()),
109-
compilationStats(stats.getCompilationStats()),
110-
stats.getProcessCpuTimeUs(),
111-
stats.getQueryPlan(),
112-
stats.getQueryAst(),
113-
stats.getTotalDurationUs(),
114-
stats.getTotalCpuTimeUs()
115-
);
116-
}
117-
118-
private static QueryPhaseStats queryPhaseStats(tech.ydb.query.result.QueryStats.QueryPhase queryPhase) {
119-
return new QueryPhaseStats(
120-
queryPhase.getDurationUs(),
121-
queryPhase.getTableAccesses().stream().map(ta -> tableAccessStats(ta)).collect(toList()),
122-
queryPhase.getCpuTimeUs(),
123-
queryPhase.getAffectedShards(),
124-
queryPhase.isLiteralPhase()
125-
);
126-
}
127-
128-
private static TableAccessStats tableAccessStats(tech.ydb.query.result.QueryStats.TableAccess tableAccess) {
129-
return new TableAccessStats(
130-
tableAccess.getTableName(),
131-
operationStats(tableAccess.getReads()),
132-
operationStats(tableAccess.getUpdates()),
133-
operationStats(tableAccess.getDeletes()),
134-
tableAccess.getPartitionsCount()
135-
);
136-
}
137-
138-
private static OperationStats operationStats(tech.ydb.query.result.QueryStats.Operation operation) {
139-
return new OperationStats(operation.getRows(), operation.getBytes());
140-
}
141-
142-
private static CompilationStats compilationStats(tech.ydb.query.result.QueryStats.Compilation compilation) {
143-
return new CompilationStats(
144-
compilation.isFromCache(),
145-
compilation.getDurationUs(),
146-
compilation.getCpuTimeUs()
147-
);
148-
}
149-
15097
private class TableSession extends BaseSession {
15198
private final SessionImpl querySession;
15299

@@ -204,10 +151,11 @@ public void onNextRawPart(long index, ValueProtos.ResultSet rs) {
204151
if (!res.isSuccess()) {
205152
return res.map(v -> null);
206153
}
207-
QueryStats info = queryStats(res.getValue().getStats());
154+
QueryStats stats = res.getValue().getStats();
208155
String txId = txRef.get();
209156
Status status = res.getStatus().withIssues(issues.toArray(new Issue[0]));
210-
return Result.success(new DataQueryResult(txId, results, info), status);
157+
DataQueryResult value = new DataQueryResult(txId, results, stats != null ? stats.toProtobuf() : null);
158+
return Result.success(value, status);
211159
});
212160
}
213161

‎query/src/main/java/tech/ydb/query/result/QueryStats.java‎

Lines changed: 163 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package tech.ydb.query.result;
22

33
import java.util.List;
4+
import java.util.Objects;
45
import java.util.stream.Collectors;
56

67
import tech.ydb.proto.YdbQueryStats;
@@ -30,11 +31,23 @@ public QueryStats(YdbQueryStats.QueryStats stats) {
3031
this.totalCpuTimeUs = stats.getTotalCpuTimeUs();
3132
}
3233

34+
public YdbQueryStats.QueryStats toProtobuf() {
35+
return YdbQueryStats.QueryStats.newBuilder()
36+
.setQueryAst(queryAst)
37+
.setQueryPlan(queryPlan)
38+
.setTotalCpuTimeUs(totalCpuTimeUs)
39+
.setTotalDurationUs(totalDurationUs)
40+
.setProcessCpuTimeUs(processCpuTimeUs)
41+
.setCompilation(compilationStats.toProtobuf())
42+
.addAllQueryPhases(queryPhases.stream().map(QueryPhase::toProtobuf).collect(Collectors.toList()))
43+
.build();
44+
}
45+
3346
public List<QueryPhase> getPhases() {
3447
return this.queryPhases;
3548
}
3649

37-
/**
50+
/*
3851
* @deprecated Use {{@link #getCompilationStats()}} instead
3952
*/
4053
@Deprecated
@@ -66,6 +79,34 @@ public long getProcessCpuTimeUs() {
6679
return this.processCpuTimeUs;
6780
}
6881

82+
@Override
83+
public int hashCode() {
84+
int hash = Objects.hash(queryPlan, queryAst, compilationStats, queryPhases);
85+
hash = 31 * hash + (int) (processCpuTimeUs ^ (processCpuTimeUs >>> 32));
86+
hash = 31 * hash + (int) (totalDurationUs ^ (totalDurationUs >>> 32));
87+
hash = 31 * hash + (int) (totalCpuTimeUs ^ (totalCpuTimeUs >>> 32));
88+
return hash;
89+
}
90+
91+
@Override
92+
public boolean equals(Object other) {
93+
if (this == other) {
94+
return true;
95+
}
96+
if (other == null || getClass() != other.getClass()) {
97+
return false;
98+
}
99+
100+
QueryStats o = (QueryStats) other;
101+
return Objects.equals(queryPlan, o.queryPlan)
102+
&& Objects.equals(queryAst, o.queryAst)
103+
&& Objects.equals(compilationStats, o.compilationStats)
104+
&& Objects.equals(queryPhases, o.queryPhases)
105+
&& processCpuTimeUs == o.processCpuTimeUs
106+
&& totalDurationUs == o.totalDurationUs
107+
&& totalCpuTimeUs == o.totalCpuTimeUs;
108+
}
109+
69110
@Override
70111
public String toString() {
71112
StringBuilder sb = new StringBuilder("QueryStats{");
@@ -104,10 +145,39 @@ public boolean isFromCache() {
104145
return this.isFromCache;
105146
}
106147

148+
public YdbQueryStats.CompilationStats toProtobuf() {
149+
return YdbQueryStats.CompilationStats.newBuilder()
150+
.setCpuTimeUs(cpuTimeUs)
151+
.setDurationUs(durationUs)
152+
.setFromCache(isFromCache)
153+
.build();
154+
}
155+
107156
@Override
108157
public String toString() {
109158
return "Compilation{durationUs=" + durationUs + ", cpuTimeUs=" + cpuTimeUs + ", cache=" + isFromCache + "}";
110159
}
160+
161+
@Override
162+
public int hashCode() {
163+
int hash = Boolean.hashCode(isFromCache);
164+
hash = 31 * hash + (int) (durationUs ^ (durationUs >>> 32));
165+
hash = 31 * hash + (int) (cpuTimeUs ^ (cpuTimeUs >>> 32));
166+
return hash;
167+
}
168+
169+
@Override
170+
public boolean equals(Object other) {
171+
if (this == other) {
172+
return true;
173+
}
174+
if (other == null || getClass() != other.getClass()) {
175+
return false;
176+
}
177+
178+
Compilation o = (Compilation) other;
179+
return isFromCache == o.isFromCache && durationUs == o.durationUs && cpuTimeUs == o.cpuTimeUs;
180+
}
111181
}
112182

113183
public static class QueryPhase {
@@ -157,6 +227,40 @@ public String toString() {
157227
sb.append("]}");
158228
return sb.toString();
159229
}
230+
231+
public YdbQueryStats.QueryPhaseStats toProtobuf() {
232+
return YdbQueryStats.QueryPhaseStats.newBuilder()
233+
.setAffectedShards(affectedShards)
234+
.setCpuTimeUs(cpuTimeUs)
235+
.setDurationUs(durationUs)
236+
.setLiteralPhase(isLiteralPhase)
237+
.addAllTableAccess(tableAccesses.stream().map(TableAccess::toProtobuf).collect(Collectors.toList()))
238+
.build();
239+
}
240+
241+
@Override
242+
public int hashCode() {
243+
int hash = Boolean.hashCode(isLiteralPhase);
244+
hash = 31 * hash + (int) (durationUs ^ (durationUs >>> 32));
245+
hash = 31 * hash + (int) (cpuTimeUs ^ (cpuTimeUs >>> 32));
246+
hash = 31 * hash + (int) (affectedShards ^ (affectedShards >>> 32));
247+
hash = 31 * hash + tableAccesses.hashCode();
248+
return hash;
249+
}
250+
251+
@Override
252+
public boolean equals(Object other) {
253+
if (this == other) {
254+
return true;
255+
}
256+
if (other == null || getClass() != other.getClass()) {
257+
return false;
258+
}
259+
260+
QueryPhase o = (QueryPhase) other;
261+
return durationUs == o.durationUs && cpuTimeUs == o.cpuTimeUs && affectedShards == o.affectedShards
262+
&& isLiteralPhase == o.isLiteralPhase && Objects.equals(tableAccesses, o.tableAccesses);
263+
}
160264
}
161265

162266
public static class TableAccess {
@@ -203,6 +307,40 @@ public String toString() {
203307
+ ", deletes={rows=" + deletes.rows + ", byte=" + deletes.bytes + "}"
204308
+ "}";
205309
}
310+
311+
public YdbQueryStats.TableAccessStats toProtobuf() {
312+
return YdbQueryStats.TableAccessStats.newBuilder()
313+
.setName(name)
314+
.setPartitionsCount(partitionsCount)
315+
.setReads(reads.toProtobuf())
316+
.setDeletes(deletes.toProtobuf())
317+
.setUpdates(updates.toProtobuf())
318+
.build();
319+
}
320+
321+
@Override
322+
public int hashCode() {
323+
int hash = Objects.hash(name, reads, updates, deletes);
324+
hash = 31 * hash + (int) (partitionsCount ^ (partitionsCount >>> 32));
325+
return hash;
326+
}
327+
328+
@Override
329+
public boolean equals(Object other) {
330+
if (this == other) {
331+
return true;
332+
}
333+
if (other == null || getClass() != other.getClass()) {
334+
return false;
335+
}
336+
337+
TableAccess o = (TableAccess) other;
338+
return Objects.equals(name, o.name)
339+
&& Objects.equals(reads, o.reads)
340+
&& Objects.equals(updates, o.updates)
341+
&& Objects.equals(deletes, o.deletes)
342+
&& partitionsCount == o.partitionsCount;
343+
}
206344
}
207345

208346
public static class Operation {
@@ -226,5 +364,29 @@ public long getBytes() {
226364
public String toString() {
227365
return "OperationStats{rows=" + rows + ", bytes=" + bytes + "}";
228366
}
367+
368+
public YdbQueryStats.OperationStats toProtobuf() {
369+
return YdbQueryStats.OperationStats.newBuilder().setRows(rows).setBytes(bytes).build();
370+
}
371+
372+
@Override
373+
public int hashCode() {
374+
int hash = 31 + (int) (rows ^ (rows >>> 32));
375+
hash = 31 * hash + (int) (bytes ^ (bytes >>> 32));
376+
return hash;
377+
}
378+
379+
@Override
380+
public boolean equals(Object other) {
381+
if (this == other) {
382+
return true;
383+
}
384+
if (other == null || getClass() != other.getClass()) {
385+
return false;
386+
}
387+
388+
Operation o = (Operation) other;
389+
return rows == o.rows && bytes == o.bytes;
390+
}
229391
}
230392
}

0 commit comments

Comments
 (0)