Skip to content

Commit 11bb8ba

Browse files
jinseopkim0jinseopkim0
committed
test(bigquery): add ArrowPojoUtilsTest and fix RECORD field construction
1 parent 939cccb commit 11bb8ba

2 files changed

Lines changed: 227 additions & 12 deletions

File tree

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

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -48,30 +48,36 @@ static com.google.cloud.bigquery.Field arrowFieldToBigQueryField(Field arrowFiel
4848
"Arrow List field must have at least one child field: " + name);
4949
}
5050
Field innerField = arrowField.getChildren().get(0);
51-
LegacySQLTypeName innerType = arrowTypeToLegacySQLTypeName(innerField.getType());
52-
builder = com.google.cloud.bigquery.Field.newBuilder(name, innerType);
53-
builder.setMode(com.google.cloud.bigquery.Field.Mode.REPEATED);
5451
if (!innerField.getChildren().isEmpty()) {
5552
List<com.google.cloud.bigquery.Field> subFields = new ArrayList<>();
5653
for (Field childField : innerField.getChildren()) {
5754
subFields.add(arrowFieldToBigQueryField(childField));
5855
}
59-
builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields));
60-
}
61-
} else {
62-
LegacySQLTypeName bqType = arrowTypeToLegacySQLTypeName(type);
63-
builder = com.google.cloud.bigquery.Field.newBuilder(name, bqType);
64-
if (arrowField.isNullable()) {
65-
builder.setMode(com.google.cloud.bigquery.Field.Mode.NULLABLE);
56+
builder =
57+
com.google.cloud.bigquery.Field.newBuilder(
58+
name, LegacySQLTypeName.RECORD, FieldList.of(subFields));
6659
} else {
67-
builder.setMode(com.google.cloud.bigquery.Field.Mode.REQUIRED);
60+
LegacySQLTypeName innerType = arrowTypeToLegacySQLTypeName(innerField.getType());
61+
builder = com.google.cloud.bigquery.Field.newBuilder(name, innerType);
6862
}
63+
builder.setMode(com.google.cloud.bigquery.Field.Mode.REPEATED);
64+
} else {
6965
if (!arrowField.getChildren().isEmpty()) {
7066
List<com.google.cloud.bigquery.Field> subFields = new ArrayList<>();
7167
for (Field childField : arrowField.getChildren()) {
7268
subFields.add(arrowFieldToBigQueryField(childField));
7369
}
74-
builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields));
70+
builder =
71+
com.google.cloud.bigquery.Field.newBuilder(
72+
name, LegacySQLTypeName.RECORD, FieldList.of(subFields));
73+
} else {
74+
LegacySQLTypeName bqType = arrowTypeToLegacySQLTypeName(type);
75+
builder = com.google.cloud.bigquery.Field.newBuilder(name, bqType);
76+
}
77+
if (arrowField.isNullable()) {
78+
builder.setMode(com.google.cloud.bigquery.Field.Mode.NULLABLE);
79+
} else {
80+
builder.setMode(com.google.cloud.bigquery.Field.Mode.REQUIRED);
7581
}
7682
}
7783
return builder.build();

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

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
/*
2+
* Copyright 2026 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 org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertThrows;
21+
import static org.junit.jupiter.api.Assertions.assertTrue;
22+
23+
import com.google.common.collect.ImmutableList;
24+
import java.util.List;
25+
import org.apache.arrow.memory.BufferAllocator;
26+
import org.apache.arrow.memory.RootAllocator;
27+
import org.apache.arrow.vector.FieldVector;
28+
import org.apache.arrow.vector.types.DateUnit;
29+
import org.apache.arrow.vector.types.FloatingPointPrecision;
30+
import org.apache.arrow.vector.types.TimeUnit;
31+
import org.apache.arrow.vector.types.pojo.ArrowType;
32+
import org.apache.arrow.vector.types.pojo.Field;
33+
import org.apache.arrow.vector.types.pojo.FieldType;
34+
import org.apache.arrow.vector.types.pojo.Schema;
35+
import org.junit.jupiter.api.Test;
36+
37+
public class ArrowPojoUtilsTest {
38+
39+
@Test
40+
public void testArrowSchemaToBigQuerySchema_Primitives() {
41+
Field intField = new Field("int_col", FieldType.nullable(new ArrowType.Int(64, true)), null);
42+
Field strField = new Field("str_col", FieldType.notNullable(new ArrowType.Utf8()), null);
43+
Field boolField = new Field("bool_col", FieldType.nullable(new ArrowType.Bool()), null);
44+
Field bytesField = new Field("bytes_col", FieldType.nullable(new ArrowType.Binary()), null);
45+
Field floatField =
46+
new Field(
47+
"float_col",
48+
FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)),
49+
null);
50+
Field decimalField =
51+
new Field("num_col", FieldType.nullable(new ArrowType.Decimal(38, 9, 128)), null);
52+
Field dateField =
53+
new Field("date_col", FieldType.nullable(new ArrowType.Date(DateUnit.DAY)), null);
54+
Field timeField =
55+
new Field(
56+
"time_col", FieldType.nullable(new ArrowType.Time(TimeUnit.MICROSECOND, 64)), null);
57+
Field tsField =
58+
new Field(
59+
"ts_col",
60+
FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")),
61+
null);
62+
63+
Schema arrowSchema =
64+
new Schema(
65+
ImmutableList.of(
66+
intField,
67+
strField,
68+
boolField,
69+
bytesField,
70+
floatField,
71+
decimalField,
72+
dateField,
73+
timeField,
74+
tsField));
75+
76+
com.google.cloud.bigquery.Schema bqSchema =
77+
ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema);
78+
79+
assertEquals(9, bqSchema.getFields().size());
80+
assertEquals(LegacySQLTypeName.INTEGER, bqSchema.getFields().get(0).getType());
81+
assertEquals(
82+
com.google.cloud.bigquery.Field.Mode.NULLABLE, bqSchema.getFields().get(0).getMode());
83+
84+
assertEquals(LegacySQLTypeName.STRING, bqSchema.getFields().get(1).getType());
85+
assertEquals(
86+
com.google.cloud.bigquery.Field.Mode.REQUIRED, bqSchema.getFields().get(1).getMode());
87+
88+
assertEquals(LegacySQLTypeName.BOOLEAN, bqSchema.getFields().get(2).getType());
89+
assertEquals(LegacySQLTypeName.BYTES, bqSchema.getFields().get(3).getType());
90+
assertEquals(LegacySQLTypeName.FLOAT, bqSchema.getFields().get(4).getType());
91+
assertEquals(LegacySQLTypeName.NUMERIC, bqSchema.getFields().get(5).getType());
92+
assertEquals(LegacySQLTypeName.DATE, bqSchema.getFields().get(6).getType());
93+
assertEquals(LegacySQLTypeName.TIME, bqSchema.getFields().get(7).getType());
94+
assertEquals(LegacySQLTypeName.TIMESTAMP, bqSchema.getFields().get(8).getType());
95+
}
96+
97+
@Test
98+
public void testArrowSchemaToBigQuerySchema_NestedStruct() {
99+
Field innerInt = new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null);
100+
Field innerStr = new Field("name", FieldType.nullable(new ArrowType.Utf8()), null);
101+
Field structField =
102+
new Field(
103+
"person",
104+
FieldType.nullable(new ArrowType.Struct()),
105+
ImmutableList.of(innerInt, innerStr));
106+
107+
Schema arrowSchema = new Schema(ImmutableList.of(structField));
108+
com.google.cloud.bigquery.Schema bqSchema =
109+
ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema);
110+
111+
assertEquals(1, bqSchema.getFields().size());
112+
com.google.cloud.bigquery.Field personField = bqSchema.getFields().get(0);
113+
assertEquals("person", personField.getName());
114+
assertEquals(LegacySQLTypeName.RECORD, personField.getType());
115+
assertEquals(2, personField.getSubFields().size());
116+
assertEquals("id", personField.getSubFields().get(0).getName());
117+
assertEquals(LegacySQLTypeName.INTEGER, personField.getSubFields().get(0).getType());
118+
assertEquals("name", personField.getSubFields().get(1).getName());
119+
assertEquals(LegacySQLTypeName.STRING, personField.getSubFields().get(1).getType());
120+
}
121+
122+
@Test
123+
public void testArrowSchemaToBigQuerySchema_ListPrimitives() {
124+
Field itemField = new Field("item", FieldType.notNullable(new ArrowType.Utf8()), null);
125+
Field listField =
126+
new Field("tags", FieldType.nullable(new ArrowType.List()), ImmutableList.of(itemField));
127+
128+
Schema arrowSchema = new Schema(ImmutableList.of(listField));
129+
com.google.cloud.bigquery.Schema bqSchema =
130+
ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema);
131+
132+
assertEquals(1, bqSchema.getFields().size());
133+
com.google.cloud.bigquery.Field tagsField = bqSchema.getFields().get(0);
134+
assertEquals("tags", tagsField.getName());
135+
assertEquals(LegacySQLTypeName.STRING, tagsField.getType());
136+
assertEquals(com.google.cloud.bigquery.Field.Mode.REPEATED, tagsField.getMode());
137+
}
138+
139+
@Test
140+
public void testArrowSchemaToBigQuerySchema_ListOfStruct() {
141+
Field innerKey = new Field("key", FieldType.nullable(new ArrowType.Utf8()), null);
142+
Field innerVal = new Field("value", FieldType.nullable(new ArrowType.Int(64, true)), null);
143+
Field structField =
144+
new Field(
145+
"item",
146+
FieldType.nullable(new ArrowType.Struct()),
147+
ImmutableList.of(innerKey, innerVal));
148+
Field listField =
149+
new Field(
150+
"entries", FieldType.nullable(new ArrowType.List()), ImmutableList.of(structField));
151+
152+
Schema arrowSchema = new Schema(ImmutableList.of(listField));
153+
com.google.cloud.bigquery.Schema bqSchema =
154+
ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema);
155+
156+
assertEquals(1, bqSchema.getFields().size());
157+
com.google.cloud.bigquery.Field entriesField = bqSchema.getFields().get(0);
158+
assertEquals("entries", entriesField.getName());
159+
assertEquals(LegacySQLTypeName.RECORD, entriesField.getType());
160+
assertEquals(com.google.cloud.bigquery.Field.Mode.REPEATED, entriesField.getMode());
161+
assertEquals(2, entriesField.getSubFields().size());
162+
assertEquals("key", entriesField.getSubFields().get(0).getName());
163+
assertEquals("value", entriesField.getSubFields().get(1).getName());
164+
}
165+
166+
@Test
167+
public void testArrowSchemaToBigQuerySchema_EmptyListThrowsException() {
168+
Field emptyList =
169+
new Field("empty_list", FieldType.nullable(new ArrowType.List()), ImmutableList.of());
170+
Schema arrowSchema = new Schema(ImmutableList.of(emptyList));
171+
172+
IllegalArgumentException thrown =
173+
assertThrows(
174+
IllegalArgumentException.class,
175+
() -> ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema));
176+
assertTrue(thrown.getMessage().contains("must have at least one child field"));
177+
}
178+
179+
@Test
180+
public void testArrowSchemaToBigQuerySchema_UnsupportedTypeThrowsException() {
181+
Field unsupportedField =
182+
new Field("unsupported", FieldType.nullable(new ArrowType.Null()), null);
183+
Schema arrowSchema = new Schema(ImmutableList.of(unsupportedField));
184+
185+
IllegalArgumentException thrown =
186+
assertThrows(
187+
IllegalArgumentException.class,
188+
() -> ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema));
189+
assertTrue(thrown.getMessage().contains("Unsupported Arrow type"));
190+
}
191+
192+
@Test
193+
public void testCreateVectors_Success() {
194+
try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
195+
Field intField = new Field("int_col", FieldType.nullable(new ArrowType.Int(32, true)), null);
196+
Field strField = new Field("str_col", FieldType.nullable(new ArrowType.Utf8()), null);
197+
Schema arrowSchema = new Schema(ImmutableList.of(intField, strField));
198+
199+
List<FieldVector> vectors = ArrowPojoUtils.createVectors(arrowSchema, allocator);
200+
assertEquals(2, vectors.size());
201+
assertEquals("int_col", vectors.get(0).getName());
202+
assertEquals("str_col", vectors.get(1).getName());
203+
204+
for (FieldVector v : vectors) {
205+
v.close();
206+
}
207+
}
208+
}
209+
}

0 commit comments

Comments
 (0)