Merge pull request #523 from mziccard/support-query-plan

Add support for BigQuery query plan explanation
googleapis · Jan 7, 2016 · b1ff62b · b1ff62b
2 parents 5da46e8 + 13b1b0e
commit b1ff62b
Show file tree

Hide file tree

Showing 5 changed files with 636 additions and 2 deletions.
diff --git a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/JobStatistics.java b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/JobStatistics.java
@@ -5,6 +5,7 @@
 import com.google.api.services.bigquery.model.JobStatistics4;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.MoreObjects.ToStringHelper;
+import com.google.common.collect.Lists;
 
 import java.io.Serializable;
 import java.util.List;
@@ -242,13 +243,15 @@ public static class QueryStatistics extends JobStatistics {
     private final Boolean cacheHit;
     private final Long totalBytesBilled;
     private final Long totalBytesProcessed;
+    private final List<QueryStage> queryPlan;
 
     static final class Builder extends JobStatistics.Builder<QueryStatistics, Builder> {
 
       private Integer billingTier;
       private Boolean cacheHit;
       private Long totalBytesBilled;
       private Long totalBytesProcessed;
+      private List<QueryStage> queryPlan;
 
       private Builder() {}
 
@@ -258,6 +261,10 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
         this.cacheHit = statisticsPb.getQuery().getCacheHit();
         this.totalBytesBilled = statisticsPb.getQuery().getTotalBytesBilled();
         this.totalBytesProcessed = statisticsPb.getQuery().getTotalBytesProcessed();
+        if (statisticsPb.getQuery().getQueryPlan() != null) {
+          this.queryPlan =
+              Lists.transform(statisticsPb.getQuery().getQueryPlan(), QueryStage.FROM_PB_FUNCTION);
+        }
       }
 
       Builder billingTier(Integer billingTier) {
@@ -280,6 +287,11 @@ Builder totalBytesProcessed(Long totalBytesProcessed) {
         return self();
       }
 
+      Builder queryPlan(List<QueryStage> queryPlan) {
+        this.queryPlan = queryPlan;
+        return self();
+      }
+
       @Override
       QueryStatistics build() {
         return new QueryStatistics(this);
@@ -292,6 +304,7 @@ private QueryStatistics(Builder builder) {
       this.cacheHit = builder.cacheHit;
       this.totalBytesBilled = builder.totalBytesBilled;
       this.totalBytesProcessed = builder.totalBytesProcessed;
+      this.queryPlan = builder.queryPlan;
     }
 
     /**
@@ -325,13 +338,26 @@ public Long totalBytesProcessed() {
       return totalBytesProcessed;
     }
 
+    /**
+     * Returns the query plan as a list of stages or {@code null} if a query plan is not available.
+     * Each stage involves a number of steps that read from data sources, perform a series of
+     * transformations on the input, and emit an output to a future stage (or the final result). The
+     * query plan is available for a completed query job and is retained for 7 days.
+     *
+     * @see <a href="https://cloud.google.com/bigquery/query-plan-explanation">Query Plan</a>
+     */
+    public List<QueryStage> queryPlan() {
+      return queryPlan;
+    }
+
     @Override
     ToStringHelper toStringHelper() {
       return super.toStringHelper()
           .add("billingTier", billingTier)
           .add("cacheHit", cacheHit)
           .add("totalBytesBilled", totalBytesBilled)
-          .add("totalBytesProcessed", totalBytesProcessed);
+          .add("totalBytesProcessed", totalBytesProcessed)
+          .add("queryPlan", queryPlan);
     }
 
     @Override
@@ -343,7 +369,7 @@ public boolean equals(Object obj) {
     @Override
     public int hashCode() {
       return Objects.hash(super.hashCode(), billingTier, cacheHit, totalBytesBilled,
-          totalBytesProcessed);
+          totalBytesProcessed, queryPlan);
     }
 
     @Override
@@ -353,6 +379,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
       queryStatisticsPb.setCacheHit(cacheHit);
       queryStatisticsPb.setTotalBytesBilled(totalBytesBilled);
       queryStatisticsPb.setTotalBytesProcessed(totalBytesProcessed);
+      if (queryPlan != null) {
+        queryStatisticsPb.setQueryPlan(Lists.transform(queryPlan, QueryStage.TO_PB_FUNCTION));
+      }
       return super.toPb().setQuery(queryStatisticsPb);
     }