Bigquery: add range partitioning to tables, load jobs, and query jobs. (#6552)

* added RangePartitioning feature to LoadJobConfiguration

* add range partitioning to tables, load job and query job

* fix javadoc

Author: Praful Makani

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

@@ -53,6 +53,7 @@ public final class LoadJobConfiguration extends JobConfiguration implements Load
   private final Boolean useAvroLogicalTypes;
   private final Map<String, String> labels;
   private final Long jobTimeoutMs;
+  private final RangePartitioning rangePartitioning;
 
   public static final class Builder extends JobConfiguration.Builder<LoadJobConfiguration, Builder>
       implements LoadConfiguration.Builder {
@@ -75,6 +76,7 @@ public static final class Builder extends JobConfiguration.Builder<LoadJobConfig
     private Boolean useAvroLogicalTypes;
     private Map<String, String> labels;
     private Long jobTimeoutMs;
+    private RangePartitioning rangePartitioning;
 
     private Builder() {
       super(Type.LOAD);
@@ -100,6 +102,7 @@ private Builder(LoadJobConfiguration loadConfiguration) {
       this.useAvroLogicalTypes = loadConfiguration.useAvroLogicalTypes;
       this.labels = loadConfiguration.labels;
       this.jobTimeoutMs = loadConfiguration.jobTimeoutMs;
+      this.rangePartitioning = loadConfiguration.rangePartitioning;
     }
 
     private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -179,6 +182,10 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
       if (configurationPb.getJobTimeoutMs() != null) {
         this.jobTimeoutMs = configurationPb.getJobTimeoutMs();
       }
+      if (loadConfigurationPb.getRangePartitioning() != null) {
+        this.rangePartitioning =
+            RangePartitioning.fromPb(loadConfigurationPb.getRangePartitioning());
+      }
     }
 
     @Override
@@ -301,6 +308,17 @@ public Builder setJobTimeoutMs(Long jobTimeoutMs) {
       return this;
     }
 
+    /**
+     * Range partitioning specification for this table. Only one of timePartitioning and
+     * rangePartitioning should be specified.
+     *
+     * @param rangePartitioning rangePartitioning or {@code null} for none
+     */
+    public Builder setRangePartitioning(RangePartitioning rangePartitioning) {
+      this.rangePartitioning = rangePartitioning;
+      return this;
+    }
+
     @Override
     public LoadJobConfiguration build() {
       return new LoadJobConfiguration(this);
@@ -326,6 +344,7 @@ private LoadJobConfiguration(Builder builder) {
     this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
     this.labels = builder.labels;
     this.jobTimeoutMs = builder.jobTimeoutMs;
+    this.rangePartitioning = builder.rangePartitioning;
   }
 
   @Override
@@ -428,6 +447,11 @@ public Long getJobTimeoutMs() {
     return jobTimeoutMs;
   }
 
+  /** Returns the range partitioning specification for the table */
+  public RangePartitioning getRangePartitioning() {
+    return rangePartitioning;
+  }
+
   @Override
   public Builder toBuilder() {
     return new Builder(this);
@@ -452,7 +476,8 @@ ToStringHelper toStringHelper() {
         .add("clustering", clustering)
         .add("useAvroLogicalTypes", useAvroLogicalTypes)
         .add("labels", labels)
-        .add("jobTimeoutMs", jobTimeoutMs);
+        .add("jobTimeoutMs", jobTimeoutMs)
+        .add("rangePartitioning", rangePartitioning);
   }
 
   @Override
@@ -542,6 +567,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
     if (jobTimeoutMs != null) {
       jobConfiguration.setJobTimeoutMs(jobTimeoutMs);
     }
+    if (rangePartitioning != null) {
+      loadConfigurationPb.setRangePartitioning(rangePartitioning.toPb());
+    }
     jobConfiguration.setLoad(loadConfigurationPb);
     return jobConfiguration;
   }

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

@@ -67,6 +67,7 @@ public final class QueryJobConfiguration extends JobConfiguration {
   private final Clustering clustering;
   private final Long jobTimeoutMs;
   private final Map<String, String> labels;
+  private final RangePartitioning rangePartitioning;
 
   /**
    * Priority levels for a query. If not specified the priority is assumed to be {@link
@@ -114,6 +115,7 @@ public static final class Builder
     private Clustering clustering;
     private Long jobTimeoutMs;
     private Map<String, String> labels;
+    private RangePartitioning rangePartitioning;
 
     private Builder() {
       super(Type.QUERY);
@@ -144,6 +146,7 @@ private Builder(QueryJobConfiguration jobConfiguration) {
       this.clustering = jobConfiguration.clustering;
       this.jobTimeoutMs = jobConfiguration.jobTimeoutMs;
       this.labels = jobConfiguration.labels;
+      this.rangePartitioning = jobConfiguration.rangePartitioning;
     }
 
     private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -233,6 +236,10 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
       if (configurationPb.getLabels() != null) {
         this.labels = configurationPb.getLabels();
       }
+      if (queryConfigurationPb.getRangePartitioning() != null) {
+        this.rangePartitioning =
+            RangePartitioning.fromPb(queryConfigurationPb.getRangePartitioning());
+      }
     }
 
     /** Sets the BigQuery SQL query to execute. */
@@ -561,6 +568,17 @@ public Builder setLabels(Map<String, String> labels) {
       return this;
     }
 
+    /**
+     * Range partitioning specification for this table. Only one of timePartitioning and
+     * rangePartitioning should be specified.
+     *
+     * @param rangePartitioning rangePartitioning or {@code null} for none
+     */
+    public Builder setRangePartitioning(RangePartitioning rangePartitioning) {
+      this.rangePartitioning = rangePartitioning;
+      return this;
+    }
+
     public QueryJobConfiguration build() {
       return new QueryJobConfiguration(this);
     }
@@ -600,6 +618,7 @@ private QueryJobConfiguration(Builder builder) {
     this.clustering = builder.clustering;
     this.jobTimeoutMs = builder.jobTimeoutMs;
     this.labels = builder.labels;
+    this.rangePartitioning = builder.rangePartitioning;
   }
 
   /**
@@ -779,6 +798,11 @@ public Map<String, String> getLabels() {
     return labels;
   }
 
+  /** Returns the range partitioning specification for the table */
+  public RangePartitioning getRangePartitioning() {
+    return rangePartitioning;
+  }
+
   @Override
   public Builder toBuilder() {
     return new Builder(this);
@@ -809,7 +833,8 @@ ToStringHelper toStringHelper() {
         .add("timePartitioning", timePartitioning)
         .add("clustering", clustering)
         .add("jobTimeoutMs", jobTimeoutMs)
-        .add("labels", labels);
+        .add("labels", labels)
+        .add("rangePartitioning", rangePartitioning);
   }
 
   @Override
@@ -843,7 +868,8 @@ public int hashCode() {
         timePartitioning,
         clustering,
         jobTimeoutMs,
-        labels);
+        labels,
+        rangePartitioning);
   }
 
   @Override
@@ -939,6 +965,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
     if (labels != null) {
       configurationPb.setLabels(labels);
     }
+    if (rangePartitioning != null) {
+      queryConfigurationPb.setRangePartitioning(rangePartitioning.toPb());
+    }
     configurationPb.setQuery(queryConfigurationPb);
     return configurationPb;
   }