Skip to content

feat(QueryService): Support QueryStats in QueryService-based TableClient wrapper #539

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 3, 2025
Merged

feat(QueryService): Support QueryStats in QueryService-based TableClient wrapper #539

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 60 additions & 4 deletions query/src/main/java/tech/ydb/query/impl/TableClientImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,19 +16,26 @@
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;
import tech.ydb.table.TableClient;
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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
);
}
}

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

final AtomicReference<String> txRef = new AtomicReference<>("");
Expand Down
4 changes: 4 additions & 0 deletions query/src/main/java/tech/ydb/query/result/QueryStats.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,10 @@ public QueryPhase(YdbQueryStats.QueryPhaseStats stats) {
this.isLiteralPhase = stats.getLiteralPhase();
}

public List<TableAccess> getTableAccesses() {
return this.tableAccesses;
}

public long getDurationUs() {
return this.durationUs;
}
Expand Down
30 changes: 30 additions & 0 deletions query/src/test/java/tech/ydb/query/TableExampleTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}
14 changes: 11 additions & 3 deletions table/src/main/java/tech/ydb/table/query/stats/CompilationStats.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
8 changes: 6 additions & 2 deletions table/src/main/java/tech/ydb/table/query/stats/OperationStats.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
31 changes: 23 additions & 8 deletions table/src/main/java/tech/ydb/table/query/stats/QueryPhaseStats.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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<TableAccessStats> 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() {
Expand Down
38 changes: 29 additions & 9 deletions table/src/main/java/tech/ydb/table/query/stats/QueryStats.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<QueryPhaseStats> queryPhases;
Expand All @@ -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<QueryPhaseStats> 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<QueryPhaseStats> getQueryPhasesList() {
Expand Down
29 changes: 23 additions & 6 deletions table/src/main/java/tech/ydb/table/query/stats/TableAccessStats.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down Expand Up @@ -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());
Expand Down