diff --git a/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java b/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java index 69ed1b4fe..67541d000 100644 --- a/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java +++ b/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java @@ -16,6 +16,7 @@ import tech.ydb.query.QuerySession; import tech.ydb.query.QueryStream; import tech.ydb.query.settings.ExecuteQuerySettings; +import tech.ydb.query.settings.QueryStatsMode; import tech.ydb.query.tools.QueryReader; import tech.ydb.table.Session; import tech.ydb.table.SessionPoolStats; @@ -23,12 +24,18 @@ import tech.ydb.table.impl.BaseSession; import tech.ydb.table.query.DataQueryResult; import tech.ydb.table.query.Params; +import tech.ydb.table.query.stats.CompilationStats; +import tech.ydb.table.query.stats.OperationStats; +import tech.ydb.table.query.stats.QueryPhaseStats; import tech.ydb.table.query.stats.QueryStats; +import tech.ydb.table.query.stats.TableAccessStats; import tech.ydb.table.result.ResultSetReader; import tech.ydb.table.rpc.TableRpc; import tech.ydb.table.rpc.grpc.GrpcTableRpc; import tech.ydb.table.settings.ExecuteDataQuerySettings; +import static java.util.stream.Collectors.toList; + /** * * @author Aleksandr Gorshenin @@ -88,14 +95,18 @@ private YdbQuery.TransactionControl mapTxControl(YdbTable.TransactionControl tc) return TxControl.txModeCtrl(TxMode.NONE, tc.getCommitTx()); } - private class ProxedDataQueryResult extends DataQueryResult { + private static class ProxedDataQueryResult extends DataQueryResult { private final String txID; private final QueryReader reader; + private final QueryStats queryStats; - ProxedDataQueryResult(String txID, QueryReader reader) { + private ProxedDataQueryResult(String txID, QueryReader reader) { super(YdbTable.ExecuteQueryResult.getDefaultInstance()); this.txID = txID; this.reader = reader; + + tech.ydb.query.result.QueryStats stats = reader.getQueryInfo().getStats(); + this.queryStats = stats == null ? null : queryStats(stats); } @Override @@ -130,12 +141,56 @@ public boolean isEmpty() { @Override public QueryStats getQueryStats() { - return null; + return this.queryStats; } @Override public boolean hasQueryStats() { - return false; + return this.queryStats != null; + } + + private static QueryStats queryStats(tech.ydb.query.result.QueryStats stats) { + return new QueryStats( + stats.getPhases().stream().map(qp -> queryPhaseStats(qp)).collect(toList()), + compilationStats(stats.getComplilationStats()), + stats.getProcessCpuTimeUs(), + stats.getQueryPlan(), + stats.getQueryAst(), + stats.getTotalDurationUs(), + stats.getTotalCpuTimeUs() + ); + } + + private static QueryPhaseStats queryPhaseStats(tech.ydb.query.result.QueryStats.QueryPhase queryPhase) { + return new QueryPhaseStats( + queryPhase.getDurationUs(), + queryPhase.getTableAccesses().stream().map(ta -> tableAccessStats(ta)).collect(toList()), + queryPhase.getCpuTimeUs(), + queryPhase.getAffectedShards(), + queryPhase.isLiteralPhase() + ); + } + + private static TableAccessStats tableAccessStats(tech.ydb.query.result.QueryStats.TableAccess tableAccess) { + return new TableAccessStats( + tableAccess.getTableName(), + operationStats(tableAccess.getReads()), + operationStats(tableAccess.getUpdates()), + operationStats(tableAccess.getDeletes()), + tableAccess.getPartitionsCount() + ); + } + + private static OperationStats operationStats(tech.ydb.query.result.QueryStats.Operation operation) { + return new OperationStats(operation.getRows(), operation.getBytes()); + } + + private static CompilationStats compilationStats(tech.ydb.query.result.QueryStats.Compilation compilation) { + return new CompilationStats( + compilation.isFromCache(), + compilation.getDurationUs(), + compilation.getCpuTimeUs() + ); } } @@ -154,6 +209,7 @@ public CompletableFuture> executeDataQueryInternal( ExecuteQuerySettings qs = ExecuteQuerySettings.newBuilder() .withTraceId(settings.getTraceId()) .withRequestTimeout(settings.getTimeoutDuration()) + .withStatsMode(QueryStatsMode.valueOf(settings.collectStats().name())) .build(); final AtomicReference txRef = new AtomicReference<>(""); diff --git a/query/src/main/java/tech/ydb/query/result/QueryStats.java b/query/src/main/java/tech/ydb/query/result/QueryStats.java index a31c736d4..697fc9532 100644 --- a/query/src/main/java/tech/ydb/query/result/QueryStats.java +++ b/query/src/main/java/tech/ydb/query/result/QueryStats.java @@ -117,6 +117,10 @@ public QueryPhase(YdbQueryStats.QueryPhaseStats stats) { this.isLiteralPhase = stats.getLiteralPhase(); } + public List getTableAccesses() { + return this.tableAccesses; + } + public long getDurationUs() { return this.durationUs; } diff --git a/query/src/test/java/tech/ydb/query/TableExampleTest.java b/query/src/test/java/tech/ydb/query/TableExampleTest.java index e42e2431b..59a212f6f 100644 --- a/query/src/test/java/tech/ydb/query/TableExampleTest.java +++ b/query/src/test/java/tech/ydb/query/TableExampleTest.java @@ -421,4 +421,34 @@ public void step09_tclTransaction() { return transaction.commit(); }).join().expectSuccess("tcl transaction problem"); } + + @Test + public void step10_queryStats() { + String query + = "DECLARE $seriesId AS Uint64; " + + "DECLARE $seasonId AS Uint64; " + + "SELECT sa.title AS season_title, sr.title AS series_title " + + "FROM seasons AS sa INNER JOIN series AS sr ON sa.series_id = sr.series_id " + + "WHERE sa.series_id = $seriesId AND sa.season_id = $seasonId"; + + tech.ydb.table.transaction.TxControl txControl = tech.ydb.table.transaction.TxControl.snapshotRo().setCommitTx(true); + Params params = Params.of( + "$seriesId", PrimitiveValue.newUint64(1), + "$seasonId", PrimitiveValue.newUint64(2) + ); + tech.ydb.table.settings.ExecuteDataQuerySettings settings = new tech.ydb.table.settings.ExecuteDataQuerySettings() + .disableQueryCache() + .setCollectStats(tech.ydb.table.query.stats.QueryStatsCollectionMode.FULL); + + DataQueryResult result = retryCtx + .supplyResult(session -> session.executeDataQuery(query, txControl, params, settings)) + .join().getValue(); + + Assert.assertTrue(result.hasQueryStats()); + Assert.assertNotNull(result.getQueryStats()); + Assert.assertNotEquals(0, result.getQueryStats().getQueryPhasesCount()); + Assert.assertNotEquals(0, result.getQueryStats().getQueryPhases(0).getTableAccessCount()); + Assert.assertFalse(result.getQueryStats().getQueryAst().isEmpty()); + Assert.assertFalse(result.getQueryStats().getQueryPlan().isEmpty()); + } } diff --git a/table/src/main/java/tech/ydb/table/query/stats/CompilationStats.java b/table/src/main/java/tech/ydb/table/query/stats/CompilationStats.java index cc647fe5c..82f49f718 100644 --- a/table/src/main/java/tech/ydb/table/query/stats/CompilationStats.java +++ b/table/src/main/java/tech/ydb/table/query/stats/CompilationStats.java @@ -8,9 +8,17 @@ public final class CompilationStats { private final long cpuTimeUs; public CompilationStats(tech.ydb.proto.YdbQueryStats.CompilationStats protoAutoGenCompilationStats) { - this.fromCache = protoAutoGenCompilationStats.getFromCache(); - this.durationUs = protoAutoGenCompilationStats.getDurationUs(); - this.cpuTimeUs = protoAutoGenCompilationStats.getCpuTimeUs(); + this( + protoAutoGenCompilationStats.getFromCache(), + protoAutoGenCompilationStats.getDurationUs(), + protoAutoGenCompilationStats.getCpuTimeUs() + ); + } + + public CompilationStats(boolean fromCache, long durationUs, long cpuTimeUs) { + this.fromCache = fromCache; + this.durationUs = durationUs; + this.cpuTimeUs = cpuTimeUs; } public boolean getFromCache() { diff --git a/table/src/main/java/tech/ydb/table/query/stats/OperationStats.java b/table/src/main/java/tech/ydb/table/query/stats/OperationStats.java index 97f9e13d1..807c061c4 100644 --- a/table/src/main/java/tech/ydb/table/query/stats/OperationStats.java +++ b/table/src/main/java/tech/ydb/table/query/stats/OperationStats.java @@ -7,8 +7,12 @@ public final class OperationStats { private final long bytes; public OperationStats(tech.ydb.proto.YdbQueryStats.OperationStats protoAutoGenOperationStats) { - this.rows = protoAutoGenOperationStats.getRows(); - this.bytes = protoAutoGenOperationStats.getBytes(); + this(protoAutoGenOperationStats.getRows(), protoAutoGenOperationStats.getBytes()); + } + + public OperationStats(long rows, long bytes) { + this.rows = rows; + this.bytes = bytes; } public long getRows() { diff --git a/table/src/main/java/tech/ydb/table/query/stats/QueryPhaseStats.java b/table/src/main/java/tech/ydb/table/query/stats/QueryPhaseStats.java index 20c5e8eff..bca9698b4 100644 --- a/table/src/main/java/tech/ydb/table/query/stats/QueryPhaseStats.java +++ b/table/src/main/java/tech/ydb/table/query/stats/QueryPhaseStats.java @@ -2,7 +2,8 @@ import java.util.List; import java.util.Objects; -import java.util.stream.Collectors; + +import static java.util.stream.Collectors.toList; public final class QueryPhaseStats { private final long durationUs; @@ -11,14 +12,28 @@ public final class QueryPhaseStats { private final long affectedShards; private final boolean literalPhase; - public QueryPhaseStats(tech.ydb.proto.YdbQueryStats.QueryPhaseStats protoAutoGenQueryPhaseStats) { - this.durationUs = protoAutoGenQueryPhaseStats.getDurationUs(); - this.tableAccess = protoAutoGenQueryPhaseStats.getTableAccessList().stream().map(TableAccessStats::new) - .collect(Collectors.toList()); - this.cpuTimeUs = protoAutoGenQueryPhaseStats.getCpuTimeUs(); - this.affectedShards = protoAutoGenQueryPhaseStats.getAffectedShards(); - this.literalPhase = protoAutoGenQueryPhaseStats.getLiteralPhase(); + this( + protoAutoGenQueryPhaseStats.getDurationUs(), + protoAutoGenQueryPhaseStats.getTableAccessList().stream().map(TableAccessStats::new).collect(toList()), + protoAutoGenQueryPhaseStats.getCpuTimeUs(), + protoAutoGenQueryPhaseStats.getAffectedShards(), + protoAutoGenQueryPhaseStats.getLiteralPhase() + ); + } + + public QueryPhaseStats( + long durationUs, + List tableAccess, + long cpuTimeUs, + long affectedShards, + boolean literalPhase + ) { + this.durationUs = durationUs; + this.tableAccess = tableAccess; + this.cpuTimeUs = cpuTimeUs; + this.affectedShards = affectedShards; + this.literalPhase = literalPhase; } public long getDurationUs() { diff --git a/table/src/main/java/tech/ydb/table/query/stats/QueryStats.java b/table/src/main/java/tech/ydb/table/query/stats/QueryStats.java index d464e4144..9781edb06 100644 --- a/table/src/main/java/tech/ydb/table/query/stats/QueryStats.java +++ b/table/src/main/java/tech/ydb/table/query/stats/QueryStats.java @@ -2,7 +2,8 @@ import java.util.List; import java.util.Objects; -import java.util.stream.Collectors; + +import static java.util.stream.Collectors.toList; public final class QueryStats { private final List queryPhases; @@ -14,14 +15,33 @@ public final class QueryStats { private final long totalCpuTimeUs; public QueryStats(tech.ydb.proto.YdbQueryStats.QueryStats protoAutoGenQueryStats) { - this.queryPhases = protoAutoGenQueryStats.getQueryPhasesList().stream().map(QueryPhaseStats::new) - .collect(Collectors.toList()); - this.compilation = new CompilationStats(protoAutoGenQueryStats.getCompilation()); - this.processCpuTimeUs = protoAutoGenQueryStats.getProcessCpuTimeUs(); - this.queryPlan = protoAutoGenQueryStats.getQueryPlan(); - this.queryAst = protoAutoGenQueryStats.getQueryAst(); - this.totalDurationUs = protoAutoGenQueryStats.getTotalDurationUs(); - this.totalCpuTimeUs = protoAutoGenQueryStats.getProcessCpuTimeUs(); + this( + protoAutoGenQueryStats.getQueryPhasesList().stream().map(QueryPhaseStats::new).collect(toList()), + new CompilationStats(protoAutoGenQueryStats.getCompilation()), + protoAutoGenQueryStats.getProcessCpuTimeUs(), + protoAutoGenQueryStats.getQueryPlan(), + protoAutoGenQueryStats.getQueryAst(), + protoAutoGenQueryStats.getTotalDurationUs(), + protoAutoGenQueryStats.getProcessCpuTimeUs() + ); + } + + public QueryStats( + List queryPhases, + CompilationStats compilation, + long processCpuTimeUs, + String queryPlan, + String queryAst, + long totalDurationUs, + long totalCpuTimeUs + ) { + this.queryPhases = queryPhases; + this.compilation = compilation; + this.processCpuTimeUs = processCpuTimeUs; + this.queryPlan = queryPlan; + this.queryAst = queryAst; + this.totalDurationUs = totalDurationUs; + this.totalCpuTimeUs = totalCpuTimeUs; } public List getQueryPhasesList() { diff --git a/table/src/main/java/tech/ydb/table/query/stats/TableAccessStats.java b/table/src/main/java/tech/ydb/table/query/stats/TableAccessStats.java index c1187c8b5..436184ad2 100644 --- a/table/src/main/java/tech/ydb/table/query/stats/TableAccessStats.java +++ b/table/src/main/java/tech/ydb/table/query/stats/TableAccessStats.java @@ -10,11 +10,27 @@ public final class TableAccessStats { private final long partitionsCount; public TableAccessStats(tech.ydb.proto.YdbQueryStats.TableAccessStats protoAutoGenTableAccessStats) { - this.name = protoAutoGenTableAccessStats.getName(); - this.reads = new OperationStats(protoAutoGenTableAccessStats.getReads()); - this.updates = new OperationStats(protoAutoGenTableAccessStats.getUpdates()); - this.deletes = new OperationStats(protoAutoGenTableAccessStats.getDeletes()); - this.partitionsCount = protoAutoGenTableAccessStats.getPartitionsCount(); + this( + protoAutoGenTableAccessStats.getName(), + new OperationStats(protoAutoGenTableAccessStats.getReads()), + new OperationStats(protoAutoGenTableAccessStats.getUpdates()), + new OperationStats(protoAutoGenTableAccessStats.getDeletes()), + protoAutoGenTableAccessStats.getPartitionsCount() + ); + } + + public TableAccessStats( + String name, + OperationStats reads, + OperationStats updates, + OperationStats deletes, + long partitionsCount + ) { + this.name = name; + this.reads = reads; + this.updates = updates; + this.deletes = deletes; + this.partitionsCount = partitionsCount; } public String getName() { @@ -45,7 +61,8 @@ public boolean equals(Object obj) { return super.equals(obj); } else { TableAccessStats other = (TableAccessStats) obj; - return Objects.equals(getName(), other.getName()) && Objects.equals(getReads(), other.getReads()) && + return Objects.equals(getName(), other.getName()) && + Objects.equals(getReads(), other.getReads()) && Objects.equals(getUpdates(), other.getUpdates()) && Objects.equals(getDeletes(), other.getDeletes()) && Objects.equals(getPartitionsCount(), other.getPartitionsCount());