Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 9e8cd76

Browse files
author
Praful Makani
authoredMay 12, 2020
feat: add extract model for extractjobconfiguration (#227)
* feat: add extract model * feat: modified code
1 parent 510a80e commit 9e8cd76

File tree

3 files changed

+180
-4
lines changed

3 files changed

+180
-4
lines changed
 

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExtractJobConfiguration.java

+85-4
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ public final class ExtractJobConfiguration extends JobConfiguration {
3838
private static final long serialVersionUID = 4147749733166593761L;
3939

4040
private final TableId sourceTable;
41+
private final ModelId sourceModel;
4142
private final List<String> destinationUris;
4243
private final Boolean printHeader;
4344
private final String fieldDelimiter;
@@ -51,6 +52,7 @@ public static final class Builder
5152
extends JobConfiguration.Builder<ExtractJobConfiguration, Builder> {
5253

5354
private TableId sourceTable;
55+
private ModelId sourceModel;
5456
private List<String> destinationUris;
5557
private Boolean printHeader;
5658
private String fieldDelimiter;
@@ -67,6 +69,7 @@ private Builder() {
6769
private Builder(ExtractJobConfiguration jobInfo) {
6870
this();
6971
this.sourceTable = jobInfo.sourceTable;
72+
this.sourceModel = jobInfo.sourceModel;
7073
this.destinationUris = jobInfo.destinationUris;
7174
this.printHeader = jobInfo.printHeader;
7275
this.fieldDelimiter = jobInfo.fieldDelimiter;
@@ -80,7 +83,12 @@ private Builder(ExtractJobConfiguration jobInfo) {
8083
private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
8184
this();
8285
JobConfigurationExtract extractConfigurationPb = configurationPb.getExtract();
83-
this.sourceTable = TableId.fromPb(extractConfigurationPb.getSourceTable());
86+
if (extractConfigurationPb.getSourceTable() != null) {
87+
this.sourceTable = TableId.fromPb(extractConfigurationPb.getSourceTable());
88+
}
89+
if (extractConfigurationPb.getSourceModel() != null) {
90+
this.sourceModel = ModelId.fromPb(extractConfigurationPb.getSourceModel());
91+
}
8492
this.destinationUris = extractConfigurationPb.getDestinationUris();
8593
this.printHeader = extractConfigurationPb.getPrintHeader();
8694
this.fieldDelimiter = extractConfigurationPb.getFieldDelimiter();
@@ -101,6 +109,12 @@ public Builder setSourceTable(TableId sourceTable) {
101109
return this;
102110
}
103111

112+
/** Sets the model to export. */
113+
public Builder setSourceModel(ModelId sourceModel) {
114+
this.sourceModel = sourceModel;
115+
return this;
116+
}
117+
104118
/**
105119
* Sets the list of fully-qualified Google Cloud Storage URIs (e.g. gs://bucket/path) where the
106120
* extracted table should be written.
@@ -191,7 +205,8 @@ public ExtractJobConfiguration build() {
191205

192206
private ExtractJobConfiguration(Builder builder) {
193207
super(builder);
194-
this.sourceTable = checkNotNull(builder.sourceTable);
208+
this.sourceTable = builder.sourceTable;
209+
this.sourceModel = builder.sourceModel;
195210
this.destinationUris = checkNotNull(builder.destinationUris);
196211
this.printHeader = builder.printHeader;
197212
this.fieldDelimiter = builder.fieldDelimiter;
@@ -207,6 +222,11 @@ public TableId getSourceTable() {
207222
return sourceTable;
208223
}
209224

225+
/** Returns the model to export. */
226+
public ModelId getSourceModel() {
227+
return sourceModel;
228+
}
229+
210230
/**
211231
* Returns the list of fully-qualified Google Cloud Storage URIs where the extracted table should
212232
* be written.
@@ -263,6 +283,7 @@ public Builder toBuilder() {
263283
ToStringHelper toStringHelper() {
264284
return super.toStringHelper()
265285
.add("sourceTable", sourceTable)
286+
.add("sourceModel", sourceModel)
266287
.add("destinationUris", destinationUris)
267288
.add("format", format)
268289
.add("printHeader", printHeader)
@@ -284,6 +305,7 @@ public int hashCode() {
284305
return Objects.hash(
285306
baseHashCode(),
286307
sourceTable,
308+
sourceModel,
287309
destinationUris,
288310
printHeader,
289311
fieldDelimiter,
@@ -296,9 +318,12 @@ public int hashCode() {
296318

297319
@Override
298320
ExtractJobConfiguration setProjectId(String projectId) {
299-
if (Strings.isNullOrEmpty(getSourceTable().getProject())) {
321+
if (getSourceTable() != null && Strings.isNullOrEmpty(getSourceTable().getProject())) {
300322
return toBuilder().setSourceTable(getSourceTable().setProjectId(projectId)).build();
301323
}
324+
if (getSourceModel() != null && Strings.isNullOrEmpty(getSourceModel().getProject())) {
325+
return toBuilder().setSourceModel(getSourceModel().setProjectId(projectId)).build();
326+
}
302327
return this;
303328
}
304329

@@ -308,7 +333,12 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
308333
com.google.api.services.bigquery.model.JobConfiguration jobConfiguration =
309334
new com.google.api.services.bigquery.model.JobConfiguration();
310335
extractConfigurationPb.setDestinationUris(destinationUris);
311-
extractConfigurationPb.setSourceTable(sourceTable.toPb());
336+
if (sourceTable != null) {
337+
extractConfigurationPb.setSourceTable(sourceTable.toPb());
338+
}
339+
if (sourceModel != null) {
340+
extractConfigurationPb.setSourceModel(sourceModel.toPb());
341+
}
312342
extractConfigurationPb.setPrintHeader(printHeader);
313343
extractConfigurationPb.setFieldDelimiter(fieldDelimiter);
314344
extractConfigurationPb.setDestinationFormat(format);
@@ -333,6 +363,15 @@ public static Builder newBuilder(TableId sourceTable, String destinationUri) {
333363
return newBuilder(sourceTable, ImmutableList.of(destinationUri));
334364
}
335365

366+
/**
367+
* Creates a builder for a BigQuery Extract Job configuration given source model and destination
368+
* URI.
369+
*/
370+
public static Builder newBuilder(ModelId sourceModel, String destinationUri) {
371+
checkArgument(!isNullOrEmpty(destinationUri), "Provided destinationUri is null or empty");
372+
return newBuilder(sourceModel, ImmutableList.of(destinationUri));
373+
}
374+
336375
/**
337376
* Creates a builder for a BigQuery Extract Job configuration given source table and destination
338377
* URIs.
@@ -341,20 +380,42 @@ public static Builder newBuilder(TableId sourceTable, List<String> destinationUr
341380
return new Builder().setSourceTable(sourceTable).setDestinationUris(destinationUris);
342381
}
343382

383+
/**
384+
* Creates a builder for a BigQuery Extract Job configuration given source model and destination
385+
* URIs.
386+
*/
387+
public static Builder newBuilder(ModelId sourceModel, List<String> destinationUris) {
388+
return new Builder().setSourceModel(sourceModel).setDestinationUris(destinationUris);
389+
}
390+
344391
/**
345392
* Returns a BigQuery Extract Job configuration for the given source table and destination URI.
346393
*/
347394
public static ExtractJobConfiguration of(TableId sourceTable, String destinationUri) {
348395
return newBuilder(sourceTable, destinationUri).build();
349396
}
350397

398+
/**
399+
* Returns a BigQuery Extract Job configuration for the given source model and destination URI.
400+
*/
401+
public static ExtractJobConfiguration of(ModelId sourceModel, String destinationUri) {
402+
return newBuilder(sourceModel, destinationUri).build();
403+
}
404+
351405
/**
352406
* Returns a BigQuery Extract Job configuration for the given source table and destination URIs.
353407
*/
354408
public static ExtractJobConfiguration of(TableId sourceTable, List<String> destinationUris) {
355409
return newBuilder(sourceTable, destinationUris).build();
356410
}
357411

412+
/**
413+
* Returns a BigQuery Extract Job configuration for the given source model and destination URIs.
414+
*/
415+
public static ExtractJobConfiguration of(ModelId sourceModel, List<String> destinationUris) {
416+
return newBuilder(sourceModel, destinationUris).build();
417+
}
418+
358419
/**
359420
* Returns a BigQuery Extract Job configuration for the given source table, format and destination
360421
* URI.
@@ -365,6 +426,16 @@ public static ExtractJobConfiguration of(
365426
return newBuilder(sourceTable, destinationUri).setFormat(format).build();
366427
}
367428

429+
/**
430+
* Returns a BigQuery Extract Job configuration for the given source model, format and destination
431+
* URI.
432+
*/
433+
public static ExtractJobConfiguration of(
434+
ModelId sourceTable, String destinationUri, String format) {
435+
checkArgument(!isNullOrEmpty(format), "Provided format is null or empty");
436+
return newBuilder(sourceTable, destinationUri).setFormat(format).build();
437+
}
438+
368439
/**
369440
* Returns a BigQuery Extract Job configuration for the given source table, format and destination
370441
* URIs.
@@ -375,6 +446,16 @@ public static ExtractJobConfiguration of(
375446
return newBuilder(sourceTable, destinationUris).setFormat(format).build();
376447
}
377448

449+
/**
450+
* Returns a BigQuery Extract Job configuration for the given source table, format and destination
451+
* URIs.
452+
*/
453+
public static ExtractJobConfiguration of(
454+
ModelId sourceModel, List<String> destinationUris, String format) {
455+
checkArgument(!isNullOrEmpty(format), "Provided format is null or empty");
456+
return newBuilder(sourceModel, destinationUris).setFormat(format).build();
457+
}
458+
378459
@SuppressWarnings("unchecked")
379460
static ExtractJobConfiguration fromPb(
380461
com.google.api.services.bigquery.model.JobConfiguration confPb) {

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExtractJobConfigurationTest.java

+57
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ public class ExtractJobConfigurationTest {
3232
private static final List<String> DESTINATION_URIS = ImmutableList.of("uri1", "uri2");
3333
private static final String DESTINATION_URI = "uri1";
3434
private static final TableId TABLE_ID = TableId.of("dataset", "table");
35+
private static final ModelId MODEL_ID = ModelId.of("dataset", "model");
3536
private static final String FIELD_DELIMITER = ",";
3637
private static final String FORMAT = "CSV";
3738
private static final String AVRO_FORMAT = "AVRO";
@@ -70,6 +71,16 @@ public class ExtractJobConfigurationTest {
7071
.setLabels(LABELS)
7172
.setJobTimeoutMs(TIMEOUT)
7273
.build();
74+
private static final ExtractJobConfiguration EXTRACT_CONFIGURATION_MODEL =
75+
ExtractJobConfiguration.newBuilder(MODEL_ID, DESTINATION_URIS)
76+
.setPrintHeader(PRINT_HEADER)
77+
.setFieldDelimiter(FIELD_DELIMITER)
78+
.setCompression(COMPRESSION)
79+
.setFormat(FORMAT)
80+
.setUseAvroLogicalTypes(USEAVROLOGICALTYPES)
81+
.setLabels(LABELS)
82+
.setJobTimeoutMs(TIMEOUT)
83+
.build();
7384

7485
@Test
7586
public void testToBuilder() {
@@ -78,6 +89,14 @@ public void testToBuilder() {
7889
ExtractJobConfiguration job =
7990
EXTRACT_CONFIGURATION.toBuilder().setSourceTable(TableId.of("dataset", "newTable")).build();
8091
assertEquals("newTable", job.getSourceTable().getTable());
92+
compareExtractJobConfiguration(
93+
EXTRACT_CONFIGURATION_MODEL, EXTRACT_CONFIGURATION_MODEL.toBuilder().build());
94+
ExtractJobConfiguration modelJob =
95+
EXTRACT_CONFIGURATION_MODEL
96+
.toBuilder()
97+
.setSourceModel(ModelId.of("dataset", "newModel"))
98+
.build();
99+
assertEquals("newModel", modelJob.getSourceModel().getModel());
81100
job = job.toBuilder().setSourceTable(TABLE_ID).build();
82101
compareExtractJobConfiguration(EXTRACT_CONFIGURATION, job);
83102
compareExtractJobConfiguration(
@@ -108,12 +127,28 @@ public void testOf() {
108127
assertEquals(TABLE_ID, job.getSourceTable());
109128
assertEquals(ImmutableList.of(DESTINATION_URI), job.getDestinationUris());
110129
assertEquals(JSON_FORMAT, job.getFormat());
130+
ExtractJobConfiguration modelJob = ExtractJobConfiguration.of(MODEL_ID, DESTINATION_URIS);
131+
assertEquals(MODEL_ID, modelJob.getSourceModel());
132+
assertEquals(DESTINATION_URIS, modelJob.getDestinationUris());
133+
modelJob = ExtractJobConfiguration.of(MODEL_ID, DESTINATION_URI);
134+
assertEquals(MODEL_ID, modelJob.getSourceModel());
135+
assertEquals(ImmutableList.of(DESTINATION_URI), modelJob.getDestinationUris());
136+
modelJob = ExtractJobConfiguration.of(MODEL_ID, DESTINATION_URIS, JSON_FORMAT);
137+
assertEquals(MODEL_ID, modelJob.getSourceModel());
138+
assertEquals(DESTINATION_URIS, modelJob.getDestinationUris());
139+
assertEquals(JSON_FORMAT, modelJob.getFormat());
140+
modelJob = ExtractJobConfiguration.of(MODEL_ID, DESTINATION_URI, JSON_FORMAT);
141+
assertEquals(MODEL_ID, modelJob.getSourceModel());
142+
assertEquals(ImmutableList.of(DESTINATION_URI), modelJob.getDestinationUris());
143+
assertEquals(JSON_FORMAT, modelJob.getFormat());
111144
}
112145

113146
@Test
114147
public void testToBuilderIncomplete() {
115148
ExtractJobConfiguration job = ExtractJobConfiguration.of(TABLE_ID, DESTINATION_URIS);
116149
compareExtractJobConfiguration(job, job.toBuilder().build());
150+
ExtractJobConfiguration modelJob = ExtractJobConfiguration.of(MODEL_ID, DESTINATION_URIS);
151+
compareExtractJobConfiguration(modelJob, modelJob.toBuilder().build());
117152
}
118153

119154
@Test
@@ -144,6 +179,14 @@ public void testBuilder() {
144179
assertEquals(USEAVROLOGICALTYPES, EXTRACT_CONFIGURATION_AVRO.getUseAvroLogicalTypes());
145180
assertEquals(LABELS, EXTRACT_CONFIGURATION_AVRO.getLabels());
146181
assertEquals(TIMEOUT, EXTRACT_CONFIGURATION_AVRO.getJobTimeoutMs());
182+
assertEquals(MODEL_ID, EXTRACT_CONFIGURATION_MODEL.getSourceModel());
183+
assertEquals(DESTINATION_URIS, EXTRACT_CONFIGURATION_MODEL.getDestinationUris());
184+
assertEquals(FIELD_DELIMITER, EXTRACT_CONFIGURATION_MODEL.getFieldDelimiter());
185+
assertEquals(COMPRESSION, EXTRACT_CONFIGURATION_MODEL.getCompression());
186+
assertEquals(PRINT_HEADER, EXTRACT_CONFIGURATION_MODEL.printHeader());
187+
assertEquals(FORMAT, EXTRACT_CONFIGURATION_MODEL.getFormat());
188+
assertEquals(LABELS, EXTRACT_CONFIGURATION_MODEL.getLabels());
189+
assertEquals(TIMEOUT, EXTRACT_CONFIGURATION_MODEL.getJobTimeoutMs());
147190
}
148191

149192
@Test
@@ -164,12 +207,17 @@ public void testToPbAndFromPb() {
164207
ExtractJobConfiguration.fromPb(EXTRACT_CONFIGURATION_AVRO.toPb()));
165208
ExtractJobConfiguration job = ExtractJobConfiguration.of(TABLE_ID, DESTINATION_URIS);
166209
compareExtractJobConfiguration(job, ExtractJobConfiguration.fromPb(job.toPb()));
210+
ExtractJobConfiguration modelJob = ExtractJobConfiguration.of(MODEL_ID, DESTINATION_URIS);
211+
compareExtractJobConfiguration(modelJob, ExtractJobConfiguration.fromPb(modelJob.toPb()));
167212
}
168213

169214
@Test
170215
public void testSetProjectId() {
171216
ExtractJobConfiguration configuration = EXTRACT_CONFIGURATION.setProjectId(TEST_PROJECT_ID);
172217
assertEquals(TEST_PROJECT_ID, configuration.getSourceTable().getProject());
218+
ExtractJobConfiguration modelConfiguration =
219+
EXTRACT_CONFIGURATION_MODEL.setProjectId(TEST_PROJECT_ID);
220+
assertEquals(TEST_PROJECT_ID, modelConfiguration.getSourceModel().getProject());
173221
}
174222

175223
@Test
@@ -181,13 +229,21 @@ public void testSetProjectIdDoNotOverride() {
181229
.build()
182230
.setProjectId("do-not-update");
183231
assertEquals(TEST_PROJECT_ID, configuration.getSourceTable().getProject());
232+
ExtractJobConfiguration modelConfiguration =
233+
EXTRACT_CONFIGURATION_MODEL
234+
.toBuilder()
235+
.setSourceModel(MODEL_ID.setProjectId(TEST_PROJECT_ID))
236+
.build()
237+
.setProjectId("do-not-update");
238+
assertEquals(TEST_PROJECT_ID, modelConfiguration.getSourceModel().getProject());
184239
}
185240

186241
@Test
187242
public void testGetType() {
188243
assertEquals(JobConfiguration.Type.EXTRACT, EXTRACT_CONFIGURATION.getType());
189244
assertEquals(JobConfiguration.Type.EXTRACT, EXTRACT_CONFIGURATION_ONE_URI.getType());
190245
assertEquals(JobConfiguration.Type.EXTRACT, EXTRACT_CONFIGURATION_AVRO.getType());
246+
assertEquals(JobConfiguration.Type.EXTRACT, EXTRACT_CONFIGURATION_MODEL.getType());
191247
}
192248

193249
private void compareExtractJobConfiguration(
@@ -196,6 +252,7 @@ private void compareExtractJobConfiguration(
196252
assertEquals(expected.hashCode(), value.hashCode());
197253
assertEquals(expected.toString(), value.toString());
198254
assertEquals(expected.getSourceTable(), value.getSourceTable());
255+
assertEquals(expected.getSourceModel(), value.getSourceModel());
199256
assertEquals(expected.getDestinationUris(), value.getDestinationUris());
200257
assertEquals(expected.getCompression(), value.getCompression());
201258
assertEquals(expected.printHeader(), value.printHeader());

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

+38
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,7 @@ public class ITBigQueryTest {
254254
private static final String LOAD_FILE = "load.csv";
255255
private static final String JSON_LOAD_FILE = "load.json";
256256
private static final String EXTRACT_FILE = "extract.csv";
257+
private static final String EXTRACT_MODEL_FILE = "extract_model.csv";
257258
private static final String BUCKET = RemoteStorageHelper.generateBucketName();
258259
private static final TableId TABLE_ID = TableId.of(DATASET, "testing_table");
259260
private static final String CSV_CONTENT = "StringValue1\nStringValue2\n";
@@ -1945,6 +1946,43 @@ public void testExtractJob() throws InterruptedException, TimeoutException {
19451946
assertTrue(bigquery.delete(destinationTable));
19461947
}
19471948

1949+
@Test
1950+
public void testExtractJobWithModel() throws InterruptedException {
1951+
String modelName = RemoteBigQueryHelper.generateModelName();
1952+
String sql =
1953+
"CREATE MODEL `"
1954+
+ MODEL_DATASET
1955+
+ "."
1956+
+ modelName
1957+
+ "`"
1958+
+ "OPTIONS ( "
1959+
+ "model_type='linear_reg', "
1960+
+ "max_iteration=1, "
1961+
+ "learn_rate=0.4, "
1962+
+ "learn_rate_strategy='constant' "
1963+
+ ") AS ( "
1964+
+ " SELECT 'a' AS f1, 2.0 AS label "
1965+
+ "UNION ALL "
1966+
+ "SELECT 'b' AS f1, 3.8 AS label "
1967+
+ ")";
1968+
1969+
QueryJobConfiguration config = QueryJobConfiguration.newBuilder(sql).build();
1970+
Job job = bigquery.create(JobInfo.of(JobId.of(), config));
1971+
job.waitFor();
1972+
assertNull(job.getStatus().getError());
1973+
ModelId destinationModel = ModelId.of(MODEL_DATASET, modelName);
1974+
assertNotNull(destinationModel);
1975+
ExtractJobConfiguration extractConfiguration =
1976+
ExtractJobConfiguration.newBuilder(
1977+
destinationModel, "gs://" + BUCKET + "/" + EXTRACT_MODEL_FILE)
1978+
.setPrintHeader(false)
1979+
.build();
1980+
Job remoteExtractJob = bigquery.create(JobInfo.of(extractConfiguration));
1981+
remoteExtractJob = remoteExtractJob.waitFor();
1982+
assertNull(remoteExtractJob.getStatus().getError());
1983+
assertTrue(bigquery.delete(destinationModel));
1984+
}
1985+
19481986
@Test
19491987
public void testExtractJobWithLabels() throws InterruptedException, TimeoutException {
19501988
String tableName = "test_export_job_table_label";

0 commit comments

Comments
 (0)
Failed to load comments.