Skip to content

Commit 0c8c6a6

Browse files
nvamelichevnvamelichev
committed
feat(QueryService): Support QueryStats in QueryService-based TableClient wrapper (#534)
1 parent 2984ceb commit 0c8c6a6

File tree

8 files changed

+186
-32
lines changed

8 files changed

+186
-32
lines changed

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

Lines changed: 60 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,26 @@
1616
import tech.ydb.query.QuerySession;
1717
import tech.ydb.query.QueryStream;
1818
import tech.ydb.query.settings.ExecuteQuerySettings;
19+
import tech.ydb.query.settings.QueryStatsMode;
1920
import tech.ydb.query.tools.QueryReader;
2021
import tech.ydb.table.Session;
2122
import tech.ydb.table.SessionPoolStats;
2223
import tech.ydb.table.TableClient;
2324
import tech.ydb.table.impl.BaseSession;
2425
import tech.ydb.table.query.DataQueryResult;
2526
import tech.ydb.table.query.Params;
27+
import tech.ydb.table.query.stats.CompilationStats;
28+
import tech.ydb.table.query.stats.OperationStats;
29+
import tech.ydb.table.query.stats.QueryPhaseStats;
2630
import tech.ydb.table.query.stats.QueryStats;
31+
import tech.ydb.table.query.stats.TableAccessStats;
2732
import tech.ydb.table.result.ResultSetReader;
2833
import tech.ydb.table.rpc.TableRpc;
2934
import tech.ydb.table.rpc.grpc.GrpcTableRpc;
3035
import tech.ydb.table.settings.ExecuteDataQuerySettings;
3136

37+
import static java.util.stream.Collectors.toList;
38+
3239
/**
3340
*
3441
* @author Aleksandr Gorshenin
@@ -88,14 +95,18 @@ private YdbQuery.TransactionControl mapTxControl(YdbTable.TransactionControl tc)
8895
return TxControl.txModeCtrl(TxMode.NONE, tc.getCommitTx());
8996
}
9097

91-
private class ProxedDataQueryResult extends DataQueryResult {
98+
private static class ProxedDataQueryResult extends DataQueryResult {
9299
private final String txID;
93100
private final QueryReader reader;
101+
private final QueryStats queryStats;
94102

95-
ProxedDataQueryResult(String txID, QueryReader reader) {
103+
private ProxedDataQueryResult(String txID, QueryReader reader) {
96104
super(YdbTable.ExecuteQueryResult.getDefaultInstance());
97105
this.txID = txID;
98106
this.reader = reader;
107+
108+
tech.ydb.query.result.QueryStats stats = reader.getQueryInfo().getStats();
109+
this.queryStats = stats == null ? null : queryStats(stats);
99110
}
100111

101112
@Override
@@ -130,12 +141,56 @@ public boolean isEmpty() {
130141

131142
@Override
132143
public QueryStats getQueryStats() {
133-
return null;
144+
return this.queryStats;
134145
}
135146

136147
@Override
137148
public boolean hasQueryStats() {
138-
return false;
149+
return this.queryStats != null;
150+
}
151+
152+
private static QueryStats queryStats(tech.ydb.query.result.QueryStats stats) {
153+
return new QueryStats(
154+
stats.getPhases().stream().map(qp -> queryPhaseStats(qp)).collect(toList()),
155+
compilationStats(stats.getComplilationStats()),
156+
stats.getProcessCpuTimeUs(),
157+
stats.getQueryPlan(),
158+
stats.getQueryAst(),
159+
stats.getTotalDurationUs(),
160+
stats.getTotalCpuTimeUs()
161+
);
162+
}
163+
164+
private static QueryPhaseStats queryPhaseStats(tech.ydb.query.result.QueryStats.QueryPhase queryPhase) {
165+
return new QueryPhaseStats(
166+
queryPhase.getDurationUs(),
167+
queryPhase.getTableAccesses().stream().map(ta -> tableAccessStats(ta)).collect(toList()),
168+
queryPhase.getCpuTimeUs(),
169+
queryPhase.getAffectedShards(),
170+
queryPhase.isLiteralPhase()
171+
);
172+
}
173+
174+
private static TableAccessStats tableAccessStats(tech.ydb.query.result.QueryStats.TableAccess tableAccess) {
175+
return new TableAccessStats(
176+
tableAccess.getTableName(),
177+
operationStats(tableAccess.getReads()),
178+
operationStats(tableAccess.getUpdates()),
179+
operationStats(tableAccess.getDeletes()),
180+
tableAccess.getPartitionsCount()
181+
);
182+
}
183+
184+
private static OperationStats operationStats(tech.ydb.query.result.QueryStats.Operation operation) {
185+
return new OperationStats(operation.getRows(), operation.getBytes());
186+
}
187+
188+
private static CompilationStats compilationStats(tech.ydb.query.result.QueryStats.Compilation compilation) {
189+
return new CompilationStats(
190+
compilation.isFromCache(),
191+
compilation.getDurationUs(),
192+
compilation.getCpuTimeUs()
193+
);
139194
}
140195
}
141196

@@ -154,6 +209,7 @@ public CompletableFuture<Result<DataQueryResult>> executeDataQueryInternal(
154209
ExecuteQuerySettings qs = ExecuteQuerySettings.newBuilder()
155210
.withTraceId(settings.getTraceId())
156211
.withRequestTimeout(settings.getTimeoutDuration())
212+
.withStatsMode(QueryStatsMode.valueOf(settings.collectStats().name()))
157213
.build();
158214

159215
final AtomicReference<String> txRef = new AtomicReference<>("");

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,10 @@ public QueryPhase(YdbQueryStats.QueryPhaseStats stats) {
117117
this.isLiteralPhase = stats.getLiteralPhase();
118118
}
119119

120+
public List<TableAccess> getTableAccesses() {
121+
return this.tableAccesses;
122+
}
123+
120124
public long getDurationUs() {
121125
return this.durationUs;
122126
}

‎query/src/test/java/tech/ydb/query/TableExampleTest.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,4 +421,34 @@ public void step09_tclTransaction() {
421421
return transaction.commit();
422422
}).join().expectSuccess("tcl transaction problem");
423423
}
424+
425+
@Test
426+
public void step10_queryStats() {
427+
String query
428+
= "DECLARE $seriesId AS Uint64; "
429+
+ "DECLARE $seasonId AS Uint64; "
430+
+ "SELECT sa.title AS season_title, sr.title AS series_title "
431+
+ "FROM seasons AS sa INNER JOIN series AS sr ON sa.series_id = sr.series_id "
432+
+ "WHERE sa.series_id = $seriesId AND sa.season_id = $seasonId";
433+
434+
tech.ydb.table.transaction.TxControl<?> txControl = tech.ydb.table.transaction.TxControl.snapshotRo().setCommitTx(true);
435+
Params params = Params.of(
436+
"$seriesId", PrimitiveValue.newUint64(1),
437+
"$seasonId", PrimitiveValue.newUint64(2)
438+
);
439+
tech.ydb.table.settings.ExecuteDataQuerySettings settings = new tech.ydb.table.settings.ExecuteDataQuerySettings()
440+
.disableQueryCache()
441+
.setCollectStats(tech.ydb.table.query.stats.QueryStatsCollectionMode.FULL);
442+
443+
DataQueryResult result = retryCtx
444+
.supplyResult(session -> session.executeDataQuery(query, txControl, params, settings))
445+
.join().getValue();
446+
447+
Assert.assertTrue(result.hasQueryStats());
448+
Assert.assertNotNull(result.getQueryStats());
449+
Assert.assertNotEquals(0, result.getQueryStats().getQueryPhasesCount());
450+
Assert.assertNotEquals(0, result.getQueryStats().getQueryPhases(0).getTableAccessCount());
451+
Assert.assertFalse(result.getQueryStats().getQueryAst().isEmpty());
452+
Assert.assertFalse(result.getQueryStats().getQueryPlan().isEmpty());
453+
}
424454
}

‎table/src/main/java/tech/ydb/table/query/stats/CompilationStats.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,17 @@ public final class CompilationStats {
88
private final long cpuTimeUs;
99

1010
public CompilationStats(tech.ydb.proto.YdbQueryStats.CompilationStats protoAutoGenCompilationStats) {
11-
this.fromCache = protoAutoGenCompilationStats.getFromCache();
12-
this.durationUs = protoAutoGenCompilationStats.getDurationUs();
13-
this.cpuTimeUs = protoAutoGenCompilationStats.getCpuTimeUs();
11+
this(
12+
protoAutoGenCompilationStats.getFromCache(),
13+
protoAutoGenCompilationStats.getDurationUs(),
14+
protoAutoGenCompilationStats.getCpuTimeUs()
15+
);
16+
}
17+
18+
public CompilationStats(boolean fromCache, long durationUs, long cpuTimeUs) {
19+
this.fromCache = fromCache;
20+
this.durationUs = durationUs;
21+
this.cpuTimeUs = cpuTimeUs;
1422
}
1523

1624
public boolean getFromCache() {

‎table/src/main/java/tech/ydb/table/query/stats/OperationStats.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,12 @@ public final class OperationStats {
77
private final long bytes;
88

99
public OperationStats(tech.ydb.proto.YdbQueryStats.OperationStats protoAutoGenOperationStats) {
10-
this.rows = protoAutoGenOperationStats.getRows();
11-
this.bytes = protoAutoGenOperationStats.getBytes();
10+
this(protoAutoGenOperationStats.getRows(), protoAutoGenOperationStats.getBytes());
11+
}
12+
13+
public OperationStats(long rows, long bytes) {
14+
this.rows = rows;
15+
this.bytes = bytes;
1216
}
1317

1418
public long getRows() {

‎table/src/main/java/tech/ydb/table/query/stats/QueryPhaseStats.java‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22

33
import java.util.List;
44
import java.util.Objects;
5-
import java.util.stream.Collectors;
5+
6+
import static java.util.stream.Collectors.toList;
67

78
public final class QueryPhaseStats {
89
private final long durationUs;
@@ -11,14 +12,28 @@ public final class QueryPhaseStats {
1112
private final long affectedShards;
1213
private final boolean literalPhase;
1314

14-
1515
public QueryPhaseStats(tech.ydb.proto.YdbQueryStats.QueryPhaseStats protoAutoGenQueryPhaseStats) {
16-
this.durationUs = protoAutoGenQueryPhaseStats.getDurationUs();
17-
this.tableAccess = protoAutoGenQueryPhaseStats.getTableAccessList().stream().map(TableAccessStats::new)
18-
.collect(Collectors.toList());
19-
this.cpuTimeUs = protoAutoGenQueryPhaseStats.getCpuTimeUs();
20-
this.affectedShards = protoAutoGenQueryPhaseStats.getAffectedShards();
21-
this.literalPhase = protoAutoGenQueryPhaseStats.getLiteralPhase();
16+
this(
17+
protoAutoGenQueryPhaseStats.getDurationUs(),
18+
protoAutoGenQueryPhaseStats.getTableAccessList().stream().map(TableAccessStats::new).collect(toList()),
19+
protoAutoGenQueryPhaseStats.getCpuTimeUs(),
20+
protoAutoGenQueryPhaseStats.getAffectedShards(),
21+
protoAutoGenQueryPhaseStats.getLiteralPhase()
22+
);
23+
}
24+
25+
public QueryPhaseStats(
26+
long durationUs,
27+
List<TableAccessStats> tableAccess,
28+
long cpuTimeUs,
29+
long affectedShards,
30+
boolean literalPhase
31+
) {
32+
this.durationUs = durationUs;
33+
this.tableAccess = tableAccess;
34+
this.cpuTimeUs = cpuTimeUs;
35+
this.affectedShards = affectedShards;
36+
this.literalPhase = literalPhase;
2237
}
2338

2439
public long getDurationUs() {

‎table/src/main/java/tech/ydb/table/query/stats/QueryStats.java‎

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22

33
import java.util.List;
44
import java.util.Objects;
5-
import java.util.stream.Collectors;
5+
6+
import static java.util.stream.Collectors.toList;
67

78
public final class QueryStats {
89
private final List<QueryPhaseStats> queryPhases;
@@ -14,14 +15,33 @@ public final class QueryStats {
1415
private final long totalCpuTimeUs;
1516

1617
public QueryStats(tech.ydb.proto.YdbQueryStats.QueryStats protoAutoGenQueryStats) {
17-
this.queryPhases = protoAutoGenQueryStats.getQueryPhasesList().stream().map(QueryPhaseStats::new)
18-
.collect(Collectors.toList());
19-
this.compilation = new CompilationStats(protoAutoGenQueryStats.getCompilation());
20-
this.processCpuTimeUs = protoAutoGenQueryStats.getProcessCpuTimeUs();
21-
this.queryPlan = protoAutoGenQueryStats.getQueryPlan();
22-
this.queryAst = protoAutoGenQueryStats.getQueryAst();
23-
this.totalDurationUs = protoAutoGenQueryStats.getTotalDurationUs();
24-
this.totalCpuTimeUs = protoAutoGenQueryStats.getProcessCpuTimeUs();
18+
this(
19+
protoAutoGenQueryStats.getQueryPhasesList().stream().map(QueryPhaseStats::new).collect(toList()),
20+
new CompilationStats(protoAutoGenQueryStats.getCompilation()),
21+
protoAutoGenQueryStats.getProcessCpuTimeUs(),
22+
protoAutoGenQueryStats.getQueryPlan(),
23+
protoAutoGenQueryStats.getQueryAst(),
24+
protoAutoGenQueryStats.getTotalDurationUs(),
25+
protoAutoGenQueryStats.getProcessCpuTimeUs()
26+
);
27+
}
28+
29+
public QueryStats(
30+
List<QueryPhaseStats> queryPhases,
31+
CompilationStats compilation,
32+
long processCpuTimeUs,
33+
String queryPlan,
34+
String queryAst,
35+
long totalDurationUs,
36+
long totalCpuTimeUs
37+
) {
38+
this.queryPhases = queryPhases;
39+
this.compilation = compilation;
40+
this.processCpuTimeUs = processCpuTimeUs;
41+
this.queryPlan = queryPlan;
42+
this.queryAst = queryAst;
43+
this.totalDurationUs = totalDurationUs;
44+
this.totalCpuTimeUs = totalCpuTimeUs;
2545
}
2646

2747
public List<QueryPhaseStats> getQueryPhasesList() {

‎table/src/main/java/tech/ydb/table/query/stats/TableAccessStats.java‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,27 @@ public final class TableAccessStats {
1010
private final long partitionsCount;
1111

1212
public TableAccessStats(tech.ydb.proto.YdbQueryStats.TableAccessStats protoAutoGenTableAccessStats) {
13-
this.name = protoAutoGenTableAccessStats.getName();
14-
this.reads = new OperationStats(protoAutoGenTableAccessStats.getReads());
15-
this.updates = new OperationStats(protoAutoGenTableAccessStats.getUpdates());
16-
this.deletes = new OperationStats(protoAutoGenTableAccessStats.getDeletes());
17-
this.partitionsCount = protoAutoGenTableAccessStats.getPartitionsCount();
13+
this(
14+
protoAutoGenTableAccessStats.getName(),
15+
new OperationStats(protoAutoGenTableAccessStats.getReads()),
16+
new OperationStats(protoAutoGenTableAccessStats.getUpdates()),
17+
new OperationStats(protoAutoGenTableAccessStats.getDeletes()),
18+
protoAutoGenTableAccessStats.getPartitionsCount()
19+
);
20+
}
21+
22+
public TableAccessStats(
23+
String name,
24+
OperationStats reads,
25+
OperationStats updates,
26+
OperationStats deletes,
27+
long partitionsCount
28+
) {
29+
this.name = name;
30+
this.reads = reads;
31+
this.updates = updates;
32+
this.deletes = deletes;
33+
this.partitionsCount = partitionsCount;
1834
}
1935

2036
public String getName() {
@@ -45,7 +61,8 @@ public boolean equals(Object obj) {
4561
return super.equals(obj);
4662
} else {
4763
TableAccessStats other = (TableAccessStats) obj;
48-
return Objects.equals(getName(), other.getName()) && Objects.equals(getReads(), other.getReads()) &&
64+
return Objects.equals(getName(), other.getName()) &&
65+
Objects.equals(getReads(), other.getReads()) &&
4966
Objects.equals(getUpdates(), other.getUpdates()) &&
5067
Objects.equals(getDeletes(), other.getDeletes()) &&
5168
Objects.equals(getPartitionsCount(), other.getPartitionsCount());

0 commit comments

Comments
 (0)