Skip to content

Commit e91be80

Browse files
PhongChuongPhongChuong
authored
feat: Add ExportDataStats to QueryStatistics (#3244)
1 parent 2c3399d commit e91be80

File tree

3 files changed

+133
-0
lines changed

3 files changed

+133
-0
lines changed

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,14 @@
1717
package com.google.cloud.bigquery;
1818

1919
import com.google.api.core.ApiFunction;
20+
import com.google.api.services.bigquery.model.ExportDataStatistics;
2021
import com.google.api.services.bigquery.model.JobConfiguration;
2122
import com.google.api.services.bigquery.model.JobStatistics2;
2223
import com.google.api.services.bigquery.model.JobStatistics3;
2324
import com.google.api.services.bigquery.model.JobStatistics4;
2425
import com.google.api.services.bigquery.model.JobStatistics5;
2526
import com.google.api.services.bigquery.model.QueryParameter;
27+
import com.google.auto.value.AutoValue;
2628
import com.google.cloud.StringEnumType;
2729
import com.google.cloud.StringEnumValue;
2830
import com.google.common.base.Function;
@@ -32,6 +34,7 @@
3234
import java.io.Serializable;
3335
import java.util.List;
3436
import java.util.Objects;
37+
import javax.annotation.Nullable;
3538
import org.checkerframework.checker.nullness.compatqual.NullableDecl;
3639

3740
/** A Google BigQuery Job statistics. */
@@ -398,6 +401,7 @@ public static class QueryStatistics extends JobStatistics {
398401
private final Long estimatedBytesProcessed;
399402
private final Long numDmlAffectedRows;
400403
private final DmlStats dmlStats;
404+
private final ExportDataStats exportDataStats;
401405
private final List<TableId> referencedTables;
402406
private final StatementType statementType;
403407
private final Long totalBytesBilled;
@@ -472,6 +476,80 @@ public static StatementType[] values() {
472476
}
473477
}
474478

479+
/**
480+
* Statistics for the EXPORT DATA statement as part of Query Job. EXTRACT JOB statistics are
481+
* populated in ExtractStatistics.
482+
*/
483+
@AutoValue
484+
public abstract static class ExportDataStats implements Serializable {
485+
private static final long serialVersionUID = 1L;
486+
487+
/**
488+
* Returns number of destination files generated in case of EXPORT DATA statement only.
489+
*
490+
* @return value or {@code null} for none
491+
*/
492+
@Nullable
493+
public abstract Long getFileCount();
494+
495+
/**
496+
* Returns number of destination rows generated in case of EXPORT DATA statement only.
497+
*
498+
* @return value or {@code null} for none
499+
*/
500+
@Nullable
501+
public abstract Long getRowCount();
502+
503+
public abstract Builder toBuilder();
504+
505+
public static Builder newBuilder() {
506+
return new AutoValue_JobStatistics_QueryStatistics_ExportDataStats.Builder();
507+
}
508+
509+
static ExportDataStats fromPb(ExportDataStatistics exportDataStatisticsPb) {
510+
Builder builder = newBuilder();
511+
if (exportDataStatisticsPb.getFileCount() != null) {
512+
builder.setFileCount(exportDataStatisticsPb.getFileCount());
513+
}
514+
if (exportDataStatisticsPb.getRowCount() != null) {
515+
builder.setRowCount(exportDataStatisticsPb.getRowCount());
516+
}
517+
return builder.build();
518+
}
519+
520+
ExportDataStatistics toPb() {
521+
ExportDataStatistics exportDataStatisticsPb = new ExportDataStatistics();
522+
if (getFileCount() != null) {
523+
exportDataStatisticsPb.setFileCount(getFileCount());
524+
}
525+
if (getRowCount() != null) {
526+
exportDataStatisticsPb.setRowCount(getRowCount());
527+
}
528+
return exportDataStatisticsPb;
529+
}
530+
531+
@AutoValue.Builder
532+
public abstract static class Builder {
533+
534+
/**
535+
* Number of destination files generated in case of EXPORT DATA statement only.
536+
*
537+
* @param fileCount fileCount or {@code null} for none
538+
*/
539+
public abstract Builder setFileCount(Long fileCount);
540+
541+
/**
542+
* Number of destination rows generated in case of EXPORT DATA statement only.
543+
*
544+
* @param rowCount rowCount or {@code null} for none
545+
*/
546+
public abstract Builder setRowCount(Long rowCount);
547+
548+
/** Creates a {@code ExportDataStats} object. */
549+
public abstract ExportDataStats build();
550+
}
551+
}
552+
475553
static final class Builder extends JobStatistics.Builder<QueryStatistics, Builder> {
476554

477555
private BiEngineStats biEngineStats;
@@ -483,6 +561,7 @@ static final class Builder extends JobStatistics.Builder<QueryStatistics, Builde
483561
private Long estimatedBytesProcessed;
484562
private Long numDmlAffectedRows;
485563
private DmlStats dmlStats;
564+
private ExportDataStats exportDataStats;
486565
private List<TableId> referencedTables;
487566
private StatementType statementType;
488567
private Long totalBytesBilled;
@@ -553,6 +632,10 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
553632
if (statisticsPb.getQuery().getDmlStats() != null) {
554633
this.dmlStats = DmlStats.fromPb(statisticsPb.getQuery().getDmlStats());
555634
}
635+
if (statisticsPb.getQuery().getExportDataStatistics() != null) {
636+
this.exportDataStats =
637+
ExportDataStats.fromPb(statisticsPb.getQuery().getExportDataStatistics());
638+
}
556639
}
557640
}
558641

@@ -601,6 +684,11 @@ Builder setDmlStats(DmlStats dmlStats) {
601684
return self();
602685
}
603686

687+
Builder setExportDataStats(ExportDataStats exportDataStats) {
688+
this.exportDataStats = exportDataStats;
689+
return self();
690+
}
691+
604692
Builder setReferenceTables(List<TableId> referencedTables) {
605693
this.referencedTables = referencedTables;
606694
return self();
@@ -683,6 +771,7 @@ private QueryStatistics(Builder builder) {
683771
this.estimatedBytesProcessed = builder.estimatedBytesProcessed;
684772
this.numDmlAffectedRows = builder.numDmlAffectedRows;
685773
this.dmlStats = builder.dmlStats;
774+
this.exportDataStats = builder.exportDataStats;
686775
this.referencedTables = builder.referencedTables;
687776
this.statementType = builder.statementType;
688777
this.totalBytesBilled = builder.totalBytesBilled;
@@ -749,6 +838,11 @@ public DmlStats getDmlStats() {
749838
return dmlStats;
750839
}
751840

841+
/** Detailed statistics for EXPORT DATA statement. */
842+
public ExportDataStats getExportDataStats() {
843+
return exportDataStats;
844+
}
845+
752846
/**
753847
* Referenced tables for the job. Queries that reference more than 50 tables will not have a
754848
* complete list.
@@ -900,6 +994,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
900994
if (dmlStats != null) {
901995
queryStatisticsPb.setDmlStats(dmlStats.toPb());
902996
}
997+
if (exportDataStats != null) {
998+
queryStatisticsPb.setExportDataStatistics(exportDataStats.toPb());
999+
}
9031000
if (referencedTables != null) {
9041001
queryStatisticsPb.setReferencedTables(
9051002
Lists.transform(referencedTables, TableId.TO_PB_FUNCTION));

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.cloud.bigquery.JobStatistics.ExtractStatistics;
2424
import com.google.cloud.bigquery.JobStatistics.LoadStatistics;
2525
import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
26+
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.ExportDataStats;
2627
import com.google.cloud.bigquery.JobStatistics.ReservationUsage;
2728
import com.google.cloud.bigquery.JobStatistics.ScriptStatistics;
2829
import com.google.cloud.bigquery.JobStatistics.ScriptStatistics.ScriptStackFrame;
@@ -64,6 +65,13 @@ public class JobStatisticsTest {
6465
.setInsertedRowCount(INSERTED_ROW_COUNT)
6566
.setUpdatedRowCount(UPDATED_ROW_COUNT)
6667
.build();
68+
private static final Long EXPORT_DATA_STATS_ROW_COUNT = 3L;
69+
private static final Long EXPORT_DATA_STATS_FILE_COUNT = 2L;
70+
private static final ExportDataStats EXPORT_DATA_STATS =
71+
ExportDataStats.newBuilder()
72+
.setRowCount(EXPORT_DATA_STATS_ROW_COUNT)
73+
.setFileCount(EXPORT_DATA_STATS_FILE_COUNT)
74+
.build();
6775
private static final QueryStatistics.StatementType STATEMENT_TYPE =
6876
QueryStatistics.StatementType.SELECT;
6977
private static final Long TOTAL_BYTES_BILLED = 24L;
@@ -189,6 +197,7 @@ public class JobStatisticsTest {
189197
.setEstimatedBytesProcessed(ESTIMATE_BYTES_PROCESSED)
190198
.setNumDmlAffectedRows(NUM_DML_AFFECTED_ROWS)
191199
.setDmlStats(DML_STATS)
200+
.setExportDataStats(EXPORT_DATA_STATS)
192201
.setReferenceTables(REFERENCED_TABLES)
193202
.setStatementType(STATEMENT_TYPE)
194203
.setTotalBytesBilled(TOTAL_BYTES_BILLED)
@@ -293,6 +302,7 @@ public void testBuilder() {
293302
assertEquals(ESTIMATE_BYTES_PROCESSED, QUERY_STATISTICS.getEstimatedBytesProcessed());
294303
assertEquals(NUM_DML_AFFECTED_ROWS, QUERY_STATISTICS.getNumDmlAffectedRows());
295304
assertEquals(DML_STATS, QUERY_STATISTICS.getDmlStats());
305+
assertEquals(EXPORT_DATA_STATS, QUERY_STATISTICS.getExportDataStats());
296306
assertEquals(REFERENCED_TABLES, QUERY_STATISTICS.getReferencedTables());
297307
assertEquals(STATEMENT_TYPE, QUERY_STATISTICS.getStatementType());
298308
assertEquals(TOTAL_BYTES_BILLED, QUERY_STATISTICS.getTotalBytesBilled());
@@ -448,6 +458,8 @@ private void compareQueryStatistics(QueryStatistics expected, QueryStatistics va
448458
assertEquals(expected.getMetadataCacheStats(), value.getMetadataCacheStats());
449459
assertEquals(expected.getStatementType(), value.getStatementType());
450460
assertEquals(expected.getTimeline(), value.getTimeline());
461+
assertEquals(expected.getDmlStats(), value.getDmlStats());
462+
assertEquals(expected.getExportDataStats(), value.getExportDataStats());
451463
}
452464

453465
private void compareStatistics(JobStatistics expected, JobStatistics value) {

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6734,4 +6734,28 @@ static GoogleCredentials loadCredentials(String credentialFile) {
67346734
}
67356735
return null;
67366736
}
6737+
6738+
@Test
6739+
public void testQueryExportStatistics() throws InterruptedException {
6740+
String query =
6741+
String.format(
6742+
"EXPORT DATA OPTIONS(\n"
6743+
+ " uri='gs://%s/*.csv',\n"
6744+
+ " format='CSV',\n"
6745+
+ " overwrite=true,\n"
6746+
+ " header=true,\n"
6747+
+ " field_delimiter=';') AS\n"
6748+
+ "SELECT num FROM UNNEST([1,2,3]) AS num",
6749+
BUCKET);
6750+
QueryJobConfiguration config =
6751+
QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build();
6752+
Job job = bigquery.create(JobInfo.of(JobId.of(), config));
6753+
job = job.waitFor();
6754+
6755+
QueryStatistics queryStatistics = job.getStatistics();
6756+
assertNotNull(queryStatistics);
6757+
assertNotNull(queryStatistics.getExportDataStats());
6758+
assertEquals(1L, queryStatistics.getExportDataStats().getFileCount().longValue());
6759+
assertEquals(3L, queryStatistics.getExportDataStats().getRowCount().longValue());
6760+
}
67376761
}

0 commit comments

Comments
 (0)