From e91be80ebdd39c2448914ff9aa1742f3079d0bb8 Mon Sep 17 00:00:00 2001 From: Phong Chuong <147636638+PhongChuong@users.noreply.github.com> Date: Wed, 17 Apr 2024 16:04:35 -0400 Subject: [PATCH] feat: Add ExportDataStats to QueryStatistics (#3244) --- .../google/cloud/bigquery/JobStatistics.java | 97 +++++++++++++++++++ .../cloud/bigquery/JobStatisticsTest.java | 12 +++ .../cloud/bigquery/it/ITBigQueryTest.java | 24 +++++ 3 files changed, 133 insertions(+) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java index c600f6a19..0200711d6 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java @@ -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; @@ -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. */ @@ -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 referencedTables; private final StatementType statementType; private final Long totalBytesBilled; @@ -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 { private BiEngineStats biEngineStats; @@ -483,6 +561,7 @@ static final class Builder extends JobStatistics.Builder referencedTables; private StatementType statementType; private Long totalBytesBilled; @@ -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()); + } } } @@ -601,6 +684,11 @@ Builder setDmlStats(DmlStats dmlStats) { return self(); } + Builder setExportDataStats(ExportDataStats exportDataStats) { + this.exportDataStats = exportDataStats; + return self(); + } + Builder setReferenceTables(List referencedTables) { this.referencedTables = referencedTables; return self(); @@ -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; @@ -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. @@ -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)); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java index 24344514e..aaf4aa2b1 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java @@ -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; @@ -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; @@ -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) @@ -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()); @@ -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) { diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index b385fe47b..430235fba 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -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()); + } }