From c1fb3904c648d4db17b7aa8bbcd970d0e0374bb6 Mon Sep 17 00:00:00 2001 From: gsabhnani <42009274+gsabhnani@users.noreply.github.com> Date: Fri, 3 Aug 2018 12:53:03 -0700 Subject: [PATCH] Add model table type for the new BigQuery ML models (#3529) Fixes listTables on datasets containing models. --- .../cloud/bigquery/ModelTableDefinition.java | 88 +++++++++++++++++++ .../cloud/bigquery/TableDefinition.java | 10 +++ .../com/google/cloud/bigquery/TableInfo.java | 3 +- .../cloud/bigquery/BigQueryImplTest.java | 7 +- 4 files changed, 106 insertions(+), 2 deletions(-) create mode 100644 google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelTableDefinition.java diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelTableDefinition.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelTableDefinition.java new file mode 100644 index 000000000000..3f7bfccab356 --- /dev/null +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelTableDefinition.java @@ -0,0 +1,88 @@ +/* + * Copyright 2018 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.core.BetaApi; +import com.google.api.services.bigquery.model.Table; +import com.google.auto.value.AutoValue; +import javax.annotation.Nullable; + +/** + * A Google BigQuery Model table definition. This definition is used to represent a BigQuery + * ML model. + * + * @see BigQuery ML Model + */ +@AutoValue +@BetaApi +public abstract class ModelTableDefinition extends TableDefinition { + + private static final long serialVersionUID = 2113445776046717900L; + + @AutoValue.Builder + public abstract static class Builder + extends TableDefinition.Builder { + + public abstract Builder setNumBytes(Long numBytes); + + public abstract Builder setLocation(String location); + + public abstract Builder setType(Type type); + + /** Creates a {@code ModelTableDefinition} object. */ + public abstract ModelTableDefinition build(); + } + + /** Returns the size of this table in bytes, excluding any data in the streaming buffer. */ + @Nullable + public abstract Long getNumBytes(); + + /** + * Returns the geographic location where the table should reside. This value is inherited from the + * dataset. + * + * @see + * Dataset Location + */ + @Nullable + public abstract String getLocation(); + + /** + * Returns a builder for a BigQuery ML model table definition. + */ + public static Builder newBuilder() { + return new AutoValue_ModelTableDefinition.Builder().setType(Type.MODEL); + } + + /** Returns a builder for the {@code + * BigQuery ML Model + */ + public static final Type MODEL = type.createAndRegister("MODEL"); + private Type(String constant) { super(constant); } @@ -157,6 +165,8 @@ static T fromPb(Table tablePb) { return (T) ViewDefinition.fromPb(tablePb); case "EXTERNAL": return (T) ExternalTableDefinition.fromPb(tablePb); + case "MODEL": + return (T) ModelTableDefinition.fromPb(tablePb); default: // never reached throw new IllegalArgumentException("Format " + tablePb.getType() + " is not supported"); diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java index ac741967c117..55c950cd8122 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java @@ -32,7 +32,8 @@ /** * Google BigQuery table information. Use {@link StandardTableDefinition} to create simple BigQuery * table. Use {@link ViewDefinition} to create a BigQuery view. Use {@link ExternalTableDefinition} - * to create a BigQuery a table backed by external data. + * to create a BigQuery a table backed by external data. Use {@link ModelDefinition} to create a + * BigQuery ML model. * * @see Managing Tables */ diff --git a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 7fd91024c8eb..4c695b22995a 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -108,10 +108,14 @@ public class BigQueryImplTest { private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA1, FIELD_SCHEMA2); private static final StandardTableDefinition TABLE_DEFINITION = StandardTableDefinition.of(TABLE_SCHEMA); + private static final ModelTableDefinition MODEL_TABLE_DEFINITION = + ModelTableDefinition.newBuilder().build(); private static final TableInfo TABLE_INFO = TableInfo.of(TABLE_ID, TABLE_DEFINITION); private static final TableInfo OTHER_TABLE_INFO = TableInfo.of(OTHER_TABLE_ID, TABLE_DEFINITION); private static final TableInfo TABLE_INFO_WITH_PROJECT = TableInfo.of(TABLE_ID_WITH_PROJECT, TABLE_DEFINITION); + private static final TableInfo MODEL_TABLE_INFO_WITH_PROJECT = + TableInfo.of(TABLE_ID_WITH_PROJECT, MODEL_TABLE_DEFINITION); private static final LoadJobConfiguration LOAD_JOB_CONFIGURATION = LoadJobConfiguration.of(TABLE_ID, "URI"); private static final LoadJobConfiguration LOAD_JOB_CONFIGURATION_WITH_PROJECT = @@ -613,7 +617,8 @@ public void testListTables() { ImmutableList tableList = ImmutableList.of( new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), - new Table(bigquery, new TableInfo.BuilderImpl(OTHER_TABLE_INFO))); + new Table(bigquery, new TableInfo.BuilderImpl(OTHER_TABLE_INFO)), + new Table(bigquery, new TableInfo.BuilderImpl(MODEL_TABLE_INFO_WITH_PROJECT))); Tuple> result = Tuple.of(CURSOR, Iterables.transform(tableList, TableInfo.TO_PB_FUNCTION)); EasyMock.expect(bigqueryRpcMock.listTables(PROJECT, DATASET, EMPTY_RPC_OPTIONS))