Skip to content

Commit 72b1715

Browse files
stephaniewang526stephaniewang526
authored
feat: add ParquetOptions support and expose it in LoadJobConfiguration and ExternalTableDefinition classes (#1318)
* feat: add ParquetOptions support and expose it in LoadJobConfiguration and ExternalTableDefinition classes Fixes #1302 * update IT
1 parent 1586b02 commit 72b1715

File tree

6 files changed

+219
-0
lines changed

6 files changed

+219
-0
lines changed

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,10 @@ static ExternalTableDefinition fromPb(Table tablePb) {
449449
builder.setFormatOptions(
450450
BigtableOptions.fromPb(externalDataConfiguration.getBigtableOptions()));
451451
}
452+
if (externalDataConfiguration.getParquetOptions() != null) {
453+
builder.setFormatOptions(
454+
ParquetOptions.fromPb(externalDataConfiguration.getParquetOptions()));
455+
}
452456
builder.setMaxBadRecords(externalDataConfiguration.getMaxBadRecords());
453457
builder.setAutodetect(externalDataConfiguration.getAutodetect());
454458
if (externalDataConfiguration.getHivePartitioningOptions() != null) {
@@ -491,6 +495,10 @@ static ExternalTableDefinition fromExternalDataConfiguration(
491495
builder.setFormatOptions(
492496
BigtableOptions.fromPb(externalDataConfiguration.getBigtableOptions()));
493497
}
498+
if (externalDataConfiguration.getParquetOptions() != null) {
499+
builder.setFormatOptions(
500+
ParquetOptions.fromPb(externalDataConfiguration.getParquetOptions()));
501+
}
494502
if (externalDataConfiguration.getMaxBadRecords() != null) {
495503
builder.setMaxBadRecords(externalDataConfiguration.getMaxBadRecords());
496504
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,8 @@ public static FormatOptions of(String format) {
126126
return googleSheets();
127127
} else if (format.equals(BIGTABLE)) {
128128
return bigtable();
129+
} else if (format.equals(PARQUET)) {
130+
return parquet();
129131
}
130132
return new FormatOptions(format);
131133
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,10 @@ public CsvOptions getCsvOptions() {
390390
return formatOptions instanceof CsvOptions ? (CsvOptions) formatOptions : null;
391391
}
392392

393+
public ParquetOptions getParquetOptions() {
394+
return formatOptions instanceof ParquetOptions ? (ParquetOptions) formatOptions : null;
395+
}
396+
393397
@Override
394398
public DatastoreBackupOptions getDatastoreBackupOptions() {
395399
return formatOptions instanceof DatastoreBackupOptions
@@ -545,6 +549,10 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
545549
loadConfigurationPb.setSkipLeadingRows(Ints.checkedCast(csvOptions.getSkipLeadingRows()));
546550
}
547551
}
552+
if (getParquetOptions() != null) {
553+
ParquetOptions parquetOptions = getParquetOptions();
554+
loadConfigurationPb.setParquetOptions(parquetOptions.toPb());
555+
}
548556
if (schema != null) {
549557
loadConfigurationPb.setSchema(schema.toPb());
550558
}

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

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import com.google.common.base.MoreObjects;
20+
import java.util.Objects;
21+
22+
public class ParquetOptions extends FormatOptions {
23+
24+
private static final long serialVersionUID = 1992L;
25+
26+
private final Boolean enableListInference;
27+
private final Boolean enumAsString;
28+
29+
public Boolean getEnableListInference() {
30+
return enableListInference;
31+
}
32+
33+
public Boolean getEnumAsString() {
34+
return enumAsString;
35+
}
36+
37+
/** A builder for {@code ParquetOptions} objects. */
38+
public static final class Builder {
39+
private Boolean enableListInference;
40+
private Boolean enumAsString;
41+
42+
private Builder() {}
43+
44+
private Builder(ParquetOptions parquetOptions) {
45+
this.enableListInference = parquetOptions.enableListInference;
46+
this.enumAsString = parquetOptions.enumAsString;
47+
}
48+
49+
public Builder setEnableListInference(Boolean enableListInference) {
50+
this.enableListInference = enableListInference;
51+
return this;
52+
}
53+
54+
public Builder setEnumAsString(Boolean enumAsString) {
55+
this.enumAsString = enumAsString;
56+
return this;
57+
}
58+
59+
public ParquetOptions build() {
60+
return new ParquetOptions(this);
61+
}
62+
}
63+
/** Returns a builder for the {@link ParquetOptions} object. */
64+
public Builder toBuilder() {
65+
return new Builder(this);
66+
}
67+
68+
ParquetOptions(Builder builder) {
69+
super(FormatOptions.PARQUET);
70+
enableListInference = builder.enableListInference;
71+
enumAsString = builder.enumAsString;
72+
}
73+
74+
@Override
75+
public String toString() {
76+
return MoreObjects.toStringHelper(this)
77+
.add("enableListInference", enableListInference)
78+
.add("enumAsString", enumAsString)
79+
.toString();
80+
}
81+
82+
@Override
83+
public final int hashCode() {
84+
return Objects.hash(enableListInference, enumAsString);
85+
}
86+
87+
@Override
88+
public final boolean equals(Object obj) {
89+
if (obj == this) {
90+
return true;
91+
}
92+
if (obj == null || !obj.getClass().equals(ParquetOptions.class)) {
93+
return false;
94+
}
95+
ParquetOptions other = (ParquetOptions) obj;
96+
return enableListInference == other.enableListInference && enumAsString == other.enumAsString;
97+
}
98+
99+
/** Returns a builder for a {@link ParquetOptions} object. */
100+
public static ParquetOptions.Builder newBuilder() {
101+
return new ParquetOptions.Builder();
102+
}
103+
104+
static ParquetOptions fromPb(
105+
com.google.api.services.bigquery.model.ParquetOptions parquetOptions) {
106+
Builder builder = newBuilder();
107+
if (parquetOptions.getEnableListInference() != null) {
108+
builder.setEnableListInference(parquetOptions.getEnableListInference());
109+
}
110+
if (parquetOptions.getEnumAsString() != null) {
111+
builder.setEnumAsString(parquetOptions.getEnumAsString());
112+
}
113+
return builder.build();
114+
}
115+
116+
com.google.api.services.bigquery.model.ParquetOptions toPb() {
117+
com.google.api.services.bigquery.model.ParquetOptions parquetOptions =
118+
new com.google.api.services.bigquery.model.ParquetOptions();
119+
if (enableListInference != null) {
120+
parquetOptions.setEnableListInference(enableListInference);
121+
}
122+
if (enumAsString != null) {
123+
parquetOptions.setEnumAsString(enumAsString);
124+
}
125+
return parquetOptions;
126+
}
127+
}

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import static com.google.common.truth.Truth.assertThat;
20+
import static org.junit.Assert.assertEquals;
21+
22+
import org.junit.Test;
23+
24+
public class ParquetOptionsTest {
25+
26+
private static final ParquetOptions OPTIONS =
27+
ParquetOptions.newBuilder().setEnableListInference(true).setEnumAsString(true).build();
28+
29+
@Test
30+
public void testToBuilder() {
31+
compareParquetOptions(OPTIONS, OPTIONS.toBuilder().build());
32+
ParquetOptions parquetOptions = OPTIONS.toBuilder().setEnableListInference(true).build();
33+
assertEquals(true, parquetOptions.getEnableListInference());
34+
parquetOptions = parquetOptions.toBuilder().setEnumAsString(true).build();
35+
compareParquetOptions(OPTIONS, parquetOptions);
36+
}
37+
38+
@Test
39+
public void testToBuilderIncomplete() {
40+
ParquetOptions parquetOptions =
41+
ParquetOptions.newBuilder().setEnableListInference(true).build();
42+
assertEquals(parquetOptions, parquetOptions.toBuilder().build());
43+
}
44+
45+
@Test
46+
public void testBuilder() {
47+
assertEquals(FormatOptions.PARQUET, OPTIONS.getType());
48+
assertEquals(true, OPTIONS.getEnableListInference());
49+
assertEquals(true, OPTIONS.getEnumAsString());
50+
}
51+
52+
@Test
53+
public void testToAndFromPb() {
54+
compareParquetOptions(OPTIONS, ParquetOptions.fromPb(OPTIONS.toPb()));
55+
ParquetOptions parquetOptions =
56+
ParquetOptions.newBuilder().setEnableListInference(true).build();
57+
compareParquetOptions(parquetOptions, ParquetOptions.fromPb(parquetOptions.toPb()));
58+
}
59+
60+
private void compareParquetOptions(ParquetOptions expected, ParquetOptions actual) {
61+
assertThat(expected).isEqualTo(actual);
62+
assertThat(expected.getEnableListInference()).isEqualTo(actual.getEnableListInference());
63+
assertThat(expected.getEnumAsString()).isEqualTo(actual.getEnumAsString());
64+
assertThat(expected.hashCode()).isEqualTo(actual.hashCode());
65+
assertThat(expected.toString()).isEqualTo(actual.toString());
66+
}
67+
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@
7777
import com.google.cloud.bigquery.Model;
7878
import com.google.cloud.bigquery.ModelId;
7979
import com.google.cloud.bigquery.ModelInfo;
80+
import com.google.cloud.bigquery.ParquetOptions;
8081
import com.google.cloud.bigquery.PolicyTags;
8182
import com.google.cloud.bigquery.QueryJobConfiguration;
8283
import com.google.cloud.bigquery.QueryParameterValue;
@@ -1840,10 +1841,13 @@ public void testQueryExternalHivePartitioningOptionAutoLayout() throws Interrupt
18401841
.setSourceUriPrefix(sourceUriPrefix)
18411842
.build();
18421843
TableId tableId = TableId.of(DATASET, tableName);
1844+
ParquetOptions parquetOptions =
1845+
ParquetOptions.newBuilder().setEnableListInference(true).setEnumAsString(true).build();
18431846
ExternalTableDefinition externalTable =
18441847
ExternalTableDefinition.newBuilder(sourceUri, FormatOptions.parquet())
18451848
.setAutodetect(true)
18461849
.setHivePartitioningOptions(hivePartitioningOptions)
1850+
.setFormatOptions(parquetOptions)
18471851
.build();
18481852
assertNotNull(bigquery.create(TableInfo.of(tableId, externalTable)));
18491853
String query =
@@ -1866,6 +1870,8 @@ public void testQueryExternalHivePartitioningOptionCustomLayout() throws Interru
18661870
"gs://"
18671871
+ CLOUD_SAMPLES_DATA
18681872
+ "/bigquery/hive-partitioning-samples/customlayout/{pkey:STRING}/";
1873+
ParquetOptions parquetOptions =
1874+
ParquetOptions.newBuilder().setEnableListInference(true).setEnumAsString(true).build();
18691875
HivePartitioningOptions hivePartitioningOptions =
18701876
HivePartitioningOptions.newBuilder()
18711877
.setMode("CUSTOM")
@@ -1877,6 +1883,7 @@ public void testQueryExternalHivePartitioningOptionCustomLayout() throws Interru
18771883
ExternalTableDefinition.newBuilder(sourceUri, FormatOptions.parquet())
18781884
.setAutodetect(true)
18791885
.setHivePartitioningOptions(hivePartitioningOptions)
1886+
.setFormatOptions(parquetOptions)
18801887
.build();
18811888
assertNotNull(bigquery.create(TableInfo.of(tableId, externalTable)));
18821889
String query =

0 commit comments

Comments
 (0)