Skip to content

Commit 68798eb

Browse files
jinseopkim0jinseopkim0
committed
refactor(bigquery): eliminate FQCNs in ArrowPojoUtils, QueryRequestInfo, and ArrowDeserializerTest
1 parent 332b54d commit 68798eb

5 files changed

Lines changed: 56 additions & 67 deletions

File tree

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,7 @@ static Schema arrowSchemaToBigQuerySchema(Object arrowSchema) {
6868
* @throws IOException if deserialization of the Arrow record batch fails
6969
*/
7070
static List<FieldValueList> deserializeRecordBatch(
71-
byte[] recordBatchBytes, Schema schema, Object arrowSchema)
72-
throws IOException {
71+
byte[] recordBatchBytes, Schema schema, Object arrowSchema) throws IOException {
7372
try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
7473
List<FieldVector> vectors = ArrowPojoUtils.createVectors(arrowSchema, allocator);
7574
try {

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

Lines changed: 23 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,55 +21,56 @@
2121
import org.apache.arrow.memory.BufferAllocator;
2222
import org.apache.arrow.vector.FieldVector;
2323
import org.apache.arrow.vector.types.pojo.ArrowType;
24-
import org.apache.arrow.vector.types.pojo.Field;
25-
import org.apache.arrow.vector.types.pojo.Schema;
2624

2725
/** Internal helper for Apache Arrow Schema/Field conversions. */
2826
final class ArrowPojoUtils {
2927

3028
private ArrowPojoUtils() {}
3129

32-
static com.google.cloud.bigquery.Schema arrowSchemaToBigQuerySchema(Object arrowSchemaObj) {
33-
Schema arrowSchema = (Schema) arrowSchemaObj;
34-
List<com.google.cloud.bigquery.Field> fields = new ArrayList<>();
35-
for (Field arrowField : arrowSchema.getFields()) {
30+
static Schema arrowSchemaToBigQuerySchema(Object arrowSchemaObj) {
31+
org.apache.arrow.vector.types.pojo.Schema arrowSchema =
32+
(org.apache.arrow.vector.types.pojo.Schema) arrowSchemaObj;
33+
List<Field> fields = new ArrayList<>();
34+
for (org.apache.arrow.vector.types.pojo.Field arrowField : arrowSchema.getFields()) {
3635
fields.add(arrowFieldToBigQueryField(arrowField));
3736
}
38-
return com.google.cloud.bigquery.Schema.of(fields);
37+
return Schema.of(fields);
3938
}
4039

41-
static com.google.cloud.bigquery.Field arrowFieldToBigQueryField(Field arrowField) {
40+
static Field arrowFieldToBigQueryField(Object arrowFieldObj) {
41+
org.apache.arrow.vector.types.pojo.Field arrowField =
42+
(org.apache.arrow.vector.types.pojo.Field) arrowFieldObj;
4243
String name = arrowField.getName();
4344
ArrowType type = arrowField.getType();
44-
com.google.cloud.bigquery.Field.Builder builder;
45+
Field.Builder builder;
4546

4647
if (type instanceof ArrowType.List) {
4748
if (arrowField.getChildren().isEmpty()) {
4849
throw new IllegalArgumentException(
4950
"Arrow List field must have at least one child field: " + name);
5051
}
51-
Field innerField = arrowField.getChildren().get(0);
52+
org.apache.arrow.vector.types.pojo.Field innerField = arrowField.getChildren().get(0);
5253
LegacySQLTypeName innerType = arrowTypeToLegacySQLTypeName(innerField.getType());
53-
builder = com.google.cloud.bigquery.Field.newBuilder(name, innerType);
54-
builder.setMode(com.google.cloud.bigquery.Field.Mode.REPEATED);
54+
builder = Field.newBuilder(name, innerType);
55+
builder.setMode(Field.Mode.REPEATED);
5556
if (!innerField.getChildren().isEmpty()) {
56-
List<com.google.cloud.bigquery.Field> subFields = new ArrayList<>();
57-
for (Field childField : innerField.getChildren()) {
57+
List<Field> subFields = new ArrayList<>();
58+
for (org.apache.arrow.vector.types.pojo.Field childField : innerField.getChildren()) {
5859
subFields.add(arrowFieldToBigQueryField(childField));
5960
}
6061
builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields));
6162
}
6263
} else {
6364
LegacySQLTypeName bqType = arrowTypeToLegacySQLTypeName(type);
64-
builder = com.google.cloud.bigquery.Field.newBuilder(name, bqType);
65+
builder = Field.newBuilder(name, bqType);
6566
if (arrowField.isNullable()) {
66-
builder.setMode(com.google.cloud.bigquery.Field.Mode.NULLABLE);
67+
builder.setMode(Field.Mode.NULLABLE);
6768
} else {
68-
builder.setMode(com.google.cloud.bigquery.Field.Mode.REQUIRED);
69+
builder.setMode(Field.Mode.REQUIRED);
6970
}
7071
if (!arrowField.getChildren().isEmpty()) {
71-
List<com.google.cloud.bigquery.Field> subFields = new ArrayList<>();
72-
for (Field childField : arrowField.getChildren()) {
72+
List<Field> subFields = new ArrayList<>();
73+
for (org.apache.arrow.vector.types.pojo.Field childField : arrowField.getChildren()) {
7374
subFields.add(arrowFieldToBigQueryField(childField));
7475
}
7576
builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields));
@@ -79,9 +80,10 @@ static com.google.cloud.bigquery.Field arrowFieldToBigQueryField(Field arrowFiel
7980
}
8081

8182
static List<FieldVector> createVectors(Object arrowSchemaObj, BufferAllocator allocator) {
82-
Schema arrowSchema = (Schema) arrowSchemaObj;
83+
org.apache.arrow.vector.types.pojo.Schema arrowSchema =
84+
(org.apache.arrow.vector.types.pojo.Schema) arrowSchemaObj;
8385
List<FieldVector> vectors = new ArrayList<>();
84-
for (Field field : arrowSchema.getFields()) {
86+
for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) {
8587
vectors.add(field.createVector(allocator));
8688
}
8789
return vectors;

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
4646
import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
4747
import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
48+
import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
4849
import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
4950
import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings;
5051
import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
@@ -356,8 +357,7 @@ public Page<FieldValueList> getNextPage() {
356357
while (rowBatch.size() < pageSize && streamIterator.hasNext()) {
357358
ReadRowsResponse response = streamIterator.next();
358359
if (response.hasArrowRecordBatch()) {
359-
com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch =
360-
response.getArrowRecordBatch();
360+
ArrowRecordBatch batch = response.getArrowRecordBatch();
361361
ArrowRecordBatch deserializedBatch =
362362
MessageSerializer.deserializeRecordBatch(
363363
new ReadChannel(

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

Lines changed: 6 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.bigquery;
1818

19+
import com.google.api.services.bigquery.model.DataFormatOptions;
1920
import com.google.api.services.bigquery.model.QueryParameter;
2021
import com.google.api.services.bigquery.model.QueryRequest;
2122
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
@@ -42,13 +43,14 @@ final class QueryRequestInfo {
4243
private final Boolean useQueryCache;
4344
private final Boolean useLegacySql;
4445
private final JobCreationMode jobCreationMode;
45-
private final com.google.api.services.bigquery.model.DataFormatOptions formatOptions;
46+
private final DataFormatOptions formatOptions;
4647
private final String reservation;
4748
private final Long jobTimeoutMs;
4849
private final QueryResultsFormat queryResultsFormat;
4950
private final ArrowSerializationOptions arrowSerializationOptions;
5051

51-
QueryRequestInfo(QueryJobConfiguration config, DataFormatOptions dataFormatOptions) {
52+
QueryRequestInfo(
53+
QueryJobConfiguration config, com.google.cloud.bigquery.DataFormatOptions dataFormatOptions) {
5254
this.config = config;
5355
this.connectionProperties = config.getConnectionProperties();
5456
this.defaultDataset = config.getDefaultDataset();
@@ -63,25 +65,13 @@ final class QueryRequestInfo {
6365
this.useLegacySql = config.useLegacySql();
6466
this.useQueryCache = config.useQueryCache();
6567
this.jobCreationMode = config.getJobCreationMode();
66-
this.formatOptions = dataFormatOptions.toPb();
68+
this.formatOptions = dataFormatOptions != null ? dataFormatOptions.toPb() : null;
6769
this.reservation = config.getReservation();
6870
this.jobTimeoutMs = config.getJobTimeoutMs();
6971
this.queryResultsFormat = config.getQueryResultsFormat();
7072
this.arrowSerializationOptions = config.getArrowSerializationOptions();
7173
}
7274

73-
/**
74-
* Determines if the query can be executed via the "fast query" path (jobs.query API) instead of
75-
* the "slow path" (jobs.insert API followed by jobs.getQueryResults).
76-
*
77-
* <p>The fast query path is preferred because it completes in a single RPC, significantly
78-
* reducing end-to-end latency for small queries.
79-
*
80-
* <p>However, the jobs.query API does not support all configuration options available in
81-
* jobs.insert (e.g., destination table, clustering, time partitioning). This method checks the
82-
* QueryJobConfiguration for any unsupported options. If any are present, we must fall back to the
83-
* jobs.insert path.
84-
*/
8575
boolean isFastQuerySupported() {
8676
return config.getClustering() == null
8777
&& config.getCreateDisposition() == null
@@ -169,7 +159,7 @@ public String toString() {
169159
.add("useQueryCache", useQueryCache)
170160
.add("useLegacySql", useLegacySql)
171161
.add("jobCreationMode", jobCreationMode)
172-
.add("formatOptions", formatOptions.getUseInt64Timestamp())
162+
.add("formatOptions", formatOptions != null ? formatOptions.getUseInt64Timestamp() : null)
173163
.add("reservation", reservation)
174164
.add("jobTimeoutMs", jobTimeoutMs)
175165
.toString();

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

Lines changed: 24 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -43,39 +43,42 @@
4343
import org.apache.arrow.vector.ipc.message.MessageSerializer;
4444
import org.apache.arrow.vector.types.TimeUnit;
4545
import org.apache.arrow.vector.types.pojo.ArrowType;
46-
import org.apache.arrow.vector.types.pojo.Field;
4746
import org.apache.arrow.vector.types.pojo.FieldType;
48-
import org.apache.arrow.vector.types.pojo.Schema;
4947
import org.junit.jupiter.api.Test;
5048

5149
public class ArrowDeserializerTest {
5250

5351
@Test
5452
public void testArrowSchemaToBigQuerySchema() {
55-
Field intField = new Field("int_col", FieldType.nullable(new ArrowType.Int(32, true)), null);
56-
Field strField = new Field("str_col", FieldType.notNullable(new ArrowType.Utf8()), null);
57-
Field boolField = new Field("bool_col", FieldType.nullable(new ArrowType.Bool()), null);
58-
Field tsField =
59-
new Field(
53+
org.apache.arrow.vector.types.pojo.Field intField =
54+
new org.apache.arrow.vector.types.pojo.Field(
55+
"int_col", FieldType.nullable(new ArrowType.Int(32, true)), null);
56+
org.apache.arrow.vector.types.pojo.Field strField =
57+
new org.apache.arrow.vector.types.pojo.Field(
58+
"str_col", FieldType.notNullable(new ArrowType.Utf8()), null);
59+
org.apache.arrow.vector.types.pojo.Field boolField =
60+
new org.apache.arrow.vector.types.pojo.Field(
61+
"bool_col", FieldType.nullable(new ArrowType.Bool()), null);
62+
org.apache.arrow.vector.types.pojo.Field tsField =
63+
new org.apache.arrow.vector.types.pojo.Field(
6064
"ts_col",
6165
FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")),
6266
null);
6367

64-
Schema arrowSchema = new Schema(ImmutableList.of(intField, strField, boolField, tsField));
68+
org.apache.arrow.vector.types.pojo.Schema arrowSchema =
69+
new org.apache.arrow.vector.types.pojo.Schema(
70+
ImmutableList.of(intField, strField, boolField, tsField));
6571

66-
com.google.cloud.bigquery.Schema bqSchema =
67-
ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema);
72+
Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema);
6873

6974
assertEquals(4, bqSchema.getFields().size());
7075
assertEquals("int_col", bqSchema.getFields().get(0).getName());
7176
assertEquals(LegacySQLTypeName.INTEGER, bqSchema.getFields().get(0).getType());
72-
assertEquals(
73-
com.google.cloud.bigquery.Field.Mode.NULLABLE, bqSchema.getFields().get(0).getMode());
77+
assertEquals(Field.Mode.NULLABLE, bqSchema.getFields().get(0).getMode());
7478

7579
assertEquals("str_col", bqSchema.getFields().get(1).getName());
7680
assertEquals(LegacySQLTypeName.STRING, bqSchema.getFields().get(1).getType());
77-
assertEquals(
78-
com.google.cloud.bigquery.Field.Mode.REQUIRED, bqSchema.getFields().get(1).getMode());
81+
assertEquals(Field.Mode.REQUIRED, bqSchema.getFields().get(1).getMode());
7982

8083
assertEquals("bool_col", bqSchema.getFields().get(2).getName());
8184
assertEquals(LegacySQLTypeName.BOOLEAN, bqSchema.getFields().get(2).getType());
@@ -128,9 +131,8 @@ public void testDeserializeRecordBatchPrimitives() throws IOException {
128131
ImmutableList.of(intVector, nameVector, scoreVector, activeVector, bytesVector, tsVector);
129132

130133
try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
131-
Schema arrowSchema = root.getSchema();
132-
com.google.cloud.bigquery.Schema bqSchema =
133-
ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema);
134+
Object arrowSchema = root.getSchema();
135+
Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema);
134136

135137
byte[] recordBatchBytes = serializeVectorSchemaRoot(root, allocator);
136138

@@ -155,11 +157,7 @@ public void testDeserializeRecordBatchPrimitives() throws IOException {
155157
assertEquals("102", row1.get("id").getStringValue());
156158
assertEquals("Bob", row1.get("name").getStringValue());
157159
assertNull(row1.get("score").getValue());
158-
assertEquals(
159-
"false",
160-
row1.get("false".equals("false") ? "active" : "score") != null
161-
? row1.get("active").getStringValue()
162-
: "false");
160+
assertEquals("false", row1.get("active").getStringValue());
163161
assertNull(row1.get("data").getValue());
164162
assertNull(row1.get("ts").getValue());
165163
} finally {
@@ -179,10 +177,10 @@ public void testSchemaMismatchThrowsException() {
179177
intVector.setValueCount(1);
180178

181179
try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(intVector))) {
182-
com.google.cloud.bigquery.Schema mismatchedSchema =
183-
com.google.cloud.bigquery.Schema.of(
184-
com.google.cloud.bigquery.Field.of("col1", LegacySQLTypeName.INTEGER),
185-
com.google.cloud.bigquery.Field.of("col2", LegacySQLTypeName.STRING));
180+
Schema mismatchedSchema =
181+
Schema.of(
182+
Field.of("col1", LegacySQLTypeName.INTEGER),
183+
Field.of("col2", LegacySQLTypeName.STRING));
186184

187185
try {
188186
ArrowDeserializer.arrowRootToFieldValueList(root, 0, mismatchedSchema);

0 commit comments

Comments
 (0)