feat: add support for BI Engine Statistics (#1723)

fixes b/205146044

Author: stephaniewang526

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BiEngineReason.java

@@ -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();
+ }
+}

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BiEngineStats.java

@@ -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();
+ }
+}

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

@@ -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;
@@ -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;
@@ -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();
@@ -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();
@@ -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;
@@ -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;
@@ -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)
@@ -722,6 +740,7 @@ public final boolean equals(Object obj) {
  public final int hashCode() {
  return Objects.hash(
  baseHashCode(),
+ biEngineStats,
  billingTier,
  cacheHit,
  totalBytesBilled,
@@ -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);

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

@@ -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";
@@ -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)
@@ -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());

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

@@ -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());
  }