Skip to content

Commit

Permalink
feat: Add ExportDataStats to QueryStatistics (#3244)
Browse files Browse the repository at this point in the history
  • Loading branch information
@PhongChuong
PhongChuong committed Apr 17, 2024
1 parent 2c3399d commit e91be80
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 0 deletions.
97 changes: 97 additions & 0 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
View file
Expand Up @@ -17,12 +17,14 @@
package com.google.cloud.bigquery;

import com.google.api.core.ApiFunction;
import com.google.api.services.bigquery.model.ExportDataStatistics;
import com.google.api.services.bigquery.model.JobConfiguration;
import com.google.api.services.bigquery.model.JobStatistics2;
import com.google.api.services.bigquery.model.JobStatistics3;
import com.google.api.services.bigquery.model.JobStatistics4;
import com.google.api.services.bigquery.model.JobStatistics5;
import com.google.api.services.bigquery.model.QueryParameter;
import com.google.auto.value.AutoValue;
import com.google.cloud.StringEnumType;
import com.google.cloud.StringEnumValue;
import com.google.common.base.Function;
Expand All @@ -32,6 +34,7 @@
import java.io.Serializable;
import java.util.List;
import java.util.Objects;
import javax.annotation.Nullable;
import org.checkerframework.checker.nullness.compatqual.NullableDecl;

/** A Google BigQuery Job statistics. */
Expand Down Expand Up @@ -398,6 +401,7 @@ public static class QueryStatistics extends JobStatistics {
private final Long estimatedBytesProcessed;
private final Long numDmlAffectedRows;
private final DmlStats dmlStats;
private final ExportDataStats exportDataStats;
private final List<TableId> referencedTables;
private final StatementType statementType;
private final Long totalBytesBilled;
Expand Down Expand Up @@ -472,6 +476,80 @@ public static StatementType[] values() {
}
}

/**
* Statistics for the EXPORT DATA statement as part of Query Job. EXTRACT JOB statistics are
* populated in ExtractStatistics.
*/
@AutoValue
public abstract static class ExportDataStats implements Serializable {
private static final long serialVersionUID = 1L;

/**
* Returns number of destination files generated in case of EXPORT DATA statement only.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract Long getFileCount();

/**
* Returns number of destination rows generated in case of EXPORT DATA statement only.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract Long getRowCount();

public abstract Builder toBuilder();

public static Builder newBuilder() {
return new AutoValue_JobStatistics_QueryStatistics_ExportDataStats.Builder();
}

static ExportDataStats fromPb(ExportDataStatistics exportDataStatisticsPb) {
Builder builder = newBuilder();
if (exportDataStatisticsPb.getFileCount() != null) {
builder.setFileCount(exportDataStatisticsPb.getFileCount());
}
if (exportDataStatisticsPb.getRowCount() != null) {
builder.setRowCount(exportDataStatisticsPb.getRowCount());
}
return builder.build();
}

ExportDataStatistics toPb() {
ExportDataStatistics exportDataStatisticsPb = new ExportDataStatistics();
if (getFileCount() != null) {
exportDataStatisticsPb.setFileCount(getFileCount());
}
if (getRowCount() != null) {
exportDataStatisticsPb.setRowCount(getRowCount());
}
return exportDataStatisticsPb;
}

@AutoValue.Builder
public abstract static class Builder {

/**
* Number of destination files generated in case of EXPORT DATA statement only.
*
* @param fileCount fileCount or {@code null} for none
*/
public abstract Builder setFileCount(Long fileCount);

/**
* Number of destination rows generated in case of EXPORT DATA statement only.
*
* @param rowCount rowCount or {@code null} for none
*/
public abstract Builder setRowCount(Long rowCount);

/** Creates a {@code ExportDataStats} object. */
public abstract ExportDataStats build();
}
}

static final class Builder extends JobStatistics.Builder<QueryStatistics, Builder> {

private BiEngineStats biEngineStats;
Expand All @@ -483,6 +561,7 @@ static final class Builder extends JobStatistics.Builder<QueryStatistics, Builde
private Long estimatedBytesProcessed;
private Long numDmlAffectedRows;
private DmlStats dmlStats;
private ExportDataStats exportDataStats;
private List<TableId> referencedTables;
private StatementType statementType;
private Long totalBytesBilled;
Expand Down Expand Up @@ -553,6 +632,10 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
if (statisticsPb.getQuery().getDmlStats() != null) {
this.dmlStats = DmlStats.fromPb(statisticsPb.getQuery().getDmlStats());
}
if (statisticsPb.getQuery().getExportDataStatistics() != null) {
this.exportDataStats =
ExportDataStats.fromPb(statisticsPb.getQuery().getExportDataStatistics());
}
}
}

Expand Down Expand Up @@ -601,6 +684,11 @@ Builder setDmlStats(DmlStats dmlStats) {
return self();
}

Builder setExportDataStats(ExportDataStats exportDataStats) {
this.exportDataStats = exportDataStats;
return self();
}

Builder setReferenceTables(List<TableId> referencedTables) {
this.referencedTables = referencedTables;
return self();
Expand Down Expand Up @@ -683,6 +771,7 @@ private QueryStatistics(Builder builder) {
this.estimatedBytesProcessed = builder.estimatedBytesProcessed;
this.numDmlAffectedRows = builder.numDmlAffectedRows;
this.dmlStats = builder.dmlStats;
this.exportDataStats = builder.exportDataStats;
this.referencedTables = builder.referencedTables;
this.statementType = builder.statementType;
this.totalBytesBilled = builder.totalBytesBilled;
Expand Down Expand Up @@ -749,6 +838,11 @@ public DmlStats getDmlStats() {
return dmlStats;
}

/** Detailed statistics for EXPORT DATA statement. */
public ExportDataStats getExportDataStats() {
return exportDataStats;
}

/**
* Referenced tables for the job. Queries that reference more than 50 tables will not have a
* complete list.
Expand Down Expand Up @@ -900,6 +994,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
if (dmlStats != null) {
queryStatisticsPb.setDmlStats(dmlStats.toPb());
}
if (exportDataStats != null) {
queryStatisticsPb.setExportDataStatistics(exportDataStats.toPb());
}
if (referencedTables != null) {
queryStatisticsPb.setReferencedTables(
Lists.transform(referencedTables, TableId.TO_PB_FUNCTION));
Expand Down
12 changes: 12 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java
View file
Expand Up @@ -23,6 +23,7 @@
import com.google.cloud.bigquery.JobStatistics.ExtractStatistics;
import com.google.cloud.bigquery.JobStatistics.LoadStatistics;
import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.ExportDataStats;
import com.google.cloud.bigquery.JobStatistics.ReservationUsage;
import com.google.cloud.bigquery.JobStatistics.ScriptStatistics;
import com.google.cloud.bigquery.JobStatistics.ScriptStatistics.ScriptStackFrame;
Expand Down Expand Up @@ -64,6 +65,13 @@ public class JobStatisticsTest {
.setInsertedRowCount(INSERTED_ROW_COUNT)
.setUpdatedRowCount(UPDATED_ROW_COUNT)
.build();
private static final Long EXPORT_DATA_STATS_ROW_COUNT = 3L;
private static final Long EXPORT_DATA_STATS_FILE_COUNT = 2L;
private static final ExportDataStats EXPORT_DATA_STATS =
ExportDataStats.newBuilder()
.setRowCount(EXPORT_DATA_STATS_ROW_COUNT)
.setFileCount(EXPORT_DATA_STATS_FILE_COUNT)
.build();
private static final QueryStatistics.StatementType STATEMENT_TYPE =
QueryStatistics.StatementType.SELECT;
private static final Long TOTAL_BYTES_BILLED = 24L;
Expand Down Expand Up @@ -189,6 +197,7 @@ public class JobStatisticsTest {
.setEstimatedBytesProcessed(ESTIMATE_BYTES_PROCESSED)
.setNumDmlAffectedRows(NUM_DML_AFFECTED_ROWS)
.setDmlStats(DML_STATS)
.setExportDataStats(EXPORT_DATA_STATS)
.setReferenceTables(REFERENCED_TABLES)
.setStatementType(STATEMENT_TYPE)
.setTotalBytesBilled(TOTAL_BYTES_BILLED)
Expand Down Expand Up @@ -293,6 +302,7 @@ public void testBuilder() {
assertEquals(ESTIMATE_BYTES_PROCESSED, QUERY_STATISTICS.getEstimatedBytesProcessed());
assertEquals(NUM_DML_AFFECTED_ROWS, QUERY_STATISTICS.getNumDmlAffectedRows());
assertEquals(DML_STATS, QUERY_STATISTICS.getDmlStats());
assertEquals(EXPORT_DATA_STATS, QUERY_STATISTICS.getExportDataStats());
assertEquals(REFERENCED_TABLES, QUERY_STATISTICS.getReferencedTables());
assertEquals(STATEMENT_TYPE, QUERY_STATISTICS.getStatementType());
assertEquals(TOTAL_BYTES_BILLED, QUERY_STATISTICS.getTotalBytesBilled());
Expand Down Expand Up @@ -448,6 +458,8 @@ private void compareQueryStatistics(QueryStatistics expected, QueryStatistics va
assertEquals(expected.getMetadataCacheStats(), value.getMetadataCacheStats());
assertEquals(expected.getStatementType(), value.getStatementType());
assertEquals(expected.getTimeline(), value.getTimeline());
assertEquals(expected.getDmlStats(), value.getDmlStats());
assertEquals(expected.getExportDataStats(), value.getExportDataStats());
}

private void compareStatistics(JobStatistics expected, JobStatistics value) {
Expand Down
24 changes: 24 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
View file
Expand Up @@ -6734,4 +6734,28 @@ static GoogleCredentials loadCredentials(String credentialFile) {
}
return null;
}

@Test
public void testQueryExportStatistics() throws InterruptedException {
String query =
String.format(
"EXPORT DATA OPTIONS(\n"
+ " uri='gs://%s/*.csv',\n"
+ " format='CSV',\n"
+ " overwrite=true,\n"
+ " header=true,\n"
+ " field_delimiter=';') AS\n"
+ "SELECT num FROM UNNEST([1,2,3]) AS num",
BUCKET);
QueryJobConfiguration config =
QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build();
Job job = bigquery.create(JobInfo.of(JobId.of(), config));
job = job.waitFor();

QueryStatistics queryStatistics = job.getStatistics();
assertNotNull(queryStatistics);
assertNotNull(queryStatistics.getExportDataStats());
assertEquals(1L, queryStatistics.getExportDataStats().getFileCount().longValue());
assertEquals(3L, queryStatistics.getExportDataStats().getRowCount().longValue());
}
}

0 comments on commit e91be80

Please sign in to comment.