Skip to content

Commit

Permalink
feat: add support for BI Engine Statistics (#1723)
Browse files Browse the repository at this point in the history
fixes b/205146044
  • Loading branch information
stephaniewang526 authored Dec 1, 2021
1 parent e7f1cf8 commit 13cc6e6
Show file tree
Hide file tree
Showing 5 changed files with 235 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
/*
* Copyright 2021 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.cloud.bigquery;

import com.google.auto.value.AutoValue;
import java.io.Serializable;
import javax.annotation.Nullable;

@AutoValue
public abstract class BiEngineReason implements Serializable {

@AutoValue.Builder
public abstract static class Builder {

/**
* High-level BI Engine reason for partial or disabled acceleration.
*
* @param code code or {@code null} for none
*/
public abstract Builder setCode(String code);

/**
* Free form human-readable reason for partial or disabled acceleration.
*
* @param message message or {@code null} for none
*/
public abstract Builder setMessage(String message);

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

/**
* High-level BI Engine reason for partial or disabled acceleration.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract String getCode();

/**
* Free form human-readable reason for partial or disabled acceleration.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract String getMessage();

public abstract Builder toBuilder();

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

com.google.api.services.bigquery.model.BiEngineReason toPb() {
com.google.api.services.bigquery.model.BiEngineReason biEngineReasonPb =
new com.google.api.services.bigquery.model.BiEngineReason();
if (getCode() != null) {
biEngineReasonPb.setCode(getCode());
}
if (getMessage() != null) {
biEngineReasonPb.setMessage(getMessage());
}
return biEngineReasonPb;
}

static BiEngineReason fromPb(
com.google.api.services.bigquery.model.BiEngineReason biEngineReasonPb) {
Builder builder = newBuilder();
if (biEngineReasonPb.getCode() != null) {
builder.setCode(biEngineReasonPb.getCode());
}
if (biEngineReasonPb.getMessage() != null) {
builder.setMessage(biEngineReasonPb.getMessage());
}
return builder.build();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
/*
* Copyright 2021 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.cloud.bigquery;

import com.google.api.services.bigquery.model.BiEngineStatistics;
import com.google.auto.value.AutoValue;
import java.io.Serializable;
import java.util.List;
import java.util.stream.Collectors;
import javax.annotation.Nullable;

/** BIEngineStatistics contains query statistics specific to the use of BI Engine. */
@AutoValue
public abstract class BiEngineStats implements Serializable {

@AutoValue.Builder
public abstract static class Builder {
/**
* Specifies which mode of BI Engine acceleration was performed (if any).
*
* @param biEngineMode biEngineMode or {@code null} for none
*/
public abstract Builder setBiEngineMode(String biEngineMode);

/**
* In case of DISABLED or PARTIAL bi_engine_mode, these contain the explanatory reasons as to
* why BI Engine could not accelerate. In case the full query was accelerated, this field is not
* populated.
*
* @param biEngineReasons biEngineReasons or {@code null} for none
*/
public abstract Builder setBiEngineReasons(List<BiEngineReason> biEngineReasons);

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

/**
* Specifies which mode of BI Engine acceleration was performed (if any).
*
* @return value or {@code null} for none
*/
@Nullable
public abstract String getBiEngineMode();

/**
* In case of DISABLED or PARTIAL bi_engine_mode, these contain the explanatory reasons as to why
* BI Engine could not accelerate. In case the full query was accelerated, this field is not
* populated.
*
* @return value or {@code null} for none
*/
@Nullable
public abstract List<BiEngineReason> getBiEngineReasons();

public abstract Builder toBuilder();

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

BiEngineStatistics toPb() {
BiEngineStatistics biEngineStatisticsPb = new BiEngineStatistics();
if (getBiEngineMode() != null) {
biEngineStatisticsPb.setBiEngineMode(getBiEngineMode());
}
if (getBiEngineReasons() != null) {
biEngineStatisticsPb.setBiEngineReasons(
getBiEngineReasons().stream().map(BiEngineReason::toPb).collect(Collectors.toList()));
}
return biEngineStatisticsPb;
}

static BiEngineStats fromPb(BiEngineStatistics biEngineStatisticsPb) {
Builder builder = newBuilder();
if (biEngineStatisticsPb.getBiEngineMode() != null) {
builder.setBiEngineMode(biEngineStatisticsPb.getBiEngineMode());
}
if (biEngineStatisticsPb.getBiEngineReasons() != null) {
builder.setBiEngineReasons(
biEngineStatisticsPb.getBiEngineReasons().stream()
.map(BiEngineReason::fromPb)
.collect(Collectors.toList()));
}
return builder.build();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -321,6 +321,7 @@ public static class QueryStatistics extends JobStatistics {

private static final long serialVersionUID = 7539354109226732353L;

private final BiEngineStats biEngineStats;
private final Integer billingTier;
private final Boolean cacheHit;
private final String ddlOperationPerformed;
Expand Down Expand Up @@ -402,6 +403,7 @@ public static StatementType[] values() {

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

private BiEngineStats biEngineStats;
private Integer billingTier;
private Boolean cacheHit;
private String ddlOperationPerformed;
Expand All @@ -425,6 +427,10 @@ private Builder() {}
private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsPb) {
super(statisticsPb);
if (statisticsPb.getQuery() != null) {
if (statisticsPb.getQuery().getBiEngineStatistics() != null) {
this.biEngineStats =
BiEngineStats.fromPb(statisticsPb.getQuery().getBiEngineStatistics());
}
this.billingTier = statisticsPb.getQuery().getBillingTier();
this.cacheHit = statisticsPb.getQuery().getCacheHit();
this.ddlOperationPerformed = statisticsPb.getQuery().getDdlOperationPerformed();
Expand Down Expand Up @@ -468,6 +474,11 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
}
}

Builder setBiEngineStats(BiEngineStats biEngineStats) {
this.biEngineStats = biEngineStats;
return self();
}

Builder setBillingTier(Integer billingTier) {
this.billingTier = billingTier;
return self();
Expand Down Expand Up @@ -566,6 +577,7 @@ QueryStatistics build() {

private QueryStatistics(Builder builder) {
super(builder);
this.biEngineStats = builder.biEngineStats;
this.billingTier = builder.billingTier;
this.cacheHit = builder.cacheHit;
this.ddlOperationPerformed = builder.ddlOperationPerformed;
Expand All @@ -585,6 +597,11 @@ private QueryStatistics(Builder builder) {
this.schema = builder.schema;
}

/** Returns query statistics specific to the use of BI Engine. */
public BiEngineStats getBiEngineStats() {
return biEngineStats;
}

/** Returns the billing tier for the job. */
public Integer getBillingTier() {
return billingTier;
Expand Down Expand Up @@ -701,6 +718,7 @@ public Schema getSchema() {
@Override
ToStringHelper toStringHelper() {
return super.toStringHelper()
.add("biEngineStats", biEngineStats)
.add("billingTier", billingTier)
.add("cacheHit", cacheHit)
.add("totalBytesBilled", totalBytesBilled)
Expand All @@ -722,6 +740,7 @@ public final boolean equals(Object obj) {
public final int hashCode() {
return Objects.hash(
baseHashCode(),
biEngineStats,
billingTier,
cacheHit,
totalBytesBilled,
Expand All @@ -733,6 +752,9 @@ public final int hashCode() {
@Override
com.google.api.services.bigquery.model.JobStatistics toPb() {
JobStatistics2 queryStatisticsPb = new JobStatistics2();
if (biEngineStats != null) {
queryStatisticsPb.setBiEngineStatistics(biEngineStats.toPb());
}
queryStatisticsPb.setBillingTier(billingTier);
queryStatisticsPb.setCacheHit(cacheHit);
queryStatisticsPb.setDdlOperationPerformed(ddlOperationPerformed);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,16 @@

public class JobStatisticsTest {

private static final BiEngineReason BI_ENGINE_REASON =
BiEngineReason.newBuilder()
.setMessage("Detected unsupported join type")
.setCode("UNSUPPORTED_SQL_TEXT")
.build();
private static final BiEngineStats BI_ENGINE_STATS =
BiEngineStats.newBuilder()
.setBiEngineReasons(ImmutableList.of(BI_ENGINE_REASON))
.setBiEngineMode("DISABLED")
.build();
private static final Integer BILLING_TIER = 42;
private static final Boolean CACHE_HIT = true;
private static final String DDL_OPERATION_PERFORMED = "SKIP";
Expand Down Expand Up @@ -154,6 +164,7 @@ public class JobStatisticsTest {
.setCreationTimestamp(CREATION_TIME)
.setEndTime(END_TIME)
.setStartTime(START_TIME)
.setBiEngineStats(BI_ENGINE_STATS)
.setBillingTier(BILLING_TIER)
.setCacheHit(CACHE_HIT)
.setDDLOperationPerformed(DDL_OPERATION_PERFORMED)
Expand Down Expand Up @@ -246,6 +257,7 @@ public void testBuilder() {
assertEquals(CREATION_TIME, QUERY_STATISTICS.getCreationTime());
assertEquals(START_TIME, QUERY_STATISTICS.getStartTime());
assertEquals(END_TIME, QUERY_STATISTICS.getEndTime());
assertEquals(BI_ENGINE_STATS, QUERY_STATISTICS.getBiEngineStats());
assertEquals(BILLING_TIER, QUERY_STATISTICS.getBillingTier());
assertEquals(CACHE_HIT, QUERY_STATISTICS.getCacheHit());
assertEquals(DDL_OPERATION_PERFORMED, QUERY_STATISTICS.getDdlOperationPerformed());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2900,6 +2900,14 @@ public void testQueryJob() throws InterruptedException, TimeoutException {
assertTrue(bigquery.delete(destinationTable));
Job queryJob = bigquery.getJob(remoteJob.getJobId());
JobStatistics.QueryStatistics statistics = queryJob.getStatistics();
if (statistics.getBiEngineStats() != null) {
assertEquals(statistics.getBiEngineStats().getBiEngineMode(), "DISABLED");
assertEquals(
statistics.getBiEngineStats().getBiEngineReasons().get(0).getCode(), "OTHER_REASON");
assertEquals(
statistics.getBiEngineStats().getBiEngineReasons().get(0).getMessage(),
"Query output to destination table is not supported.");
}
assertNotNull(statistics.getQueryPlan());
}

Expand Down

0 comments on commit 13cc6e6

Please sign in to comment.