Skip to content

Commit 7235227

Browse files
keshavdandevakeshavdandeva
authored
feat(bigquery-jdbc): support picosecond timestamps in Arrow Storage Read API and nested types
1 parent 5b769ed commit 7235227

11 files changed

Lines changed: 499 additions & 68 deletions

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArray.java‎

Lines changed: 43 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,22 +23,28 @@
2323
import java.sql.SQLException;
2424
import org.apache.arrow.vector.util.JsonStringArrayList;
2525
import org.apache.arrow.vector.util.JsonStringHashMap;
26+
import org.apache.arrow.vector.util.Text;
2627

2728
/**
2829
* An implementation of {@link BigQueryBaseArray} used to represent Array values from Arrow data.
2930
*/
3031
class BigQueryArrowArray extends BigQueryBaseArray {
3132

3233
private JsonStringArrayList<?> values;
34+
private final boolean enableTimestampPicos;
3335

34-
public BigQueryArrowArray(Field schema, JsonStringArrayList<?> values) {
35-
this(schema, values, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowArray.class));
36+
BigQueryArrowArray(Field schema, JsonStringArrayList<?> values) {
37+
this(schema, values, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowArray.class), false);
3638
}
3739

38-
public BigQueryArrowArray(
39-
Field schema, JsonStringArrayList<?> values, BigQueryJdbcResultSetLogger log) {
40+
BigQueryArrowArray(
41+
Field schema,
42+
JsonStringArrayList<?> values,
43+
BigQueryJdbcResultSetLogger log,
44+
boolean enableTimestampPicos) {
4045
super(schema, log);
4146
this.values = values;
47+
this.enableTimestampPicos = enableTimestampPicos;
4248
}
4349

4450
@Override
@@ -72,7 +78,11 @@ public ResultSet getResultSet() throws SQLException {
7278
BigQueryArrowBatchWrapper arrowBatchWrapper =
7379
BigQueryArrowBatchWrapper.getNestedFieldValueListWrapper(values);
7480
return BigQueryArrowResultSet.getNestedResultSet(
75-
Schema.of(singleElementSchema()), arrowBatchWrapper, 0, this.values.size());
81+
Schema.of(singleElementSchema()),
82+
arrowBatchWrapper,
83+
0,
84+
this.values.size(),
85+
this.enableTimestampPicos);
7686
}
7787

7888
@Override
@@ -86,7 +96,11 @@ public ResultSet getResultSet(long index, int count) throws SQLException {
8696
BigQueryArrowBatchWrapper arrowBatchWrapper =
8797
BigQueryArrowBatchWrapper.getNestedFieldValueListWrapper(values);
8898
return BigQueryArrowResultSet.getNestedResultSet(
89-
Schema.of(singleElementSchema()), arrowBatchWrapper, range.x(), range.y());
99+
Schema.of(singleElementSchema()),
100+
arrowBatchWrapper,
101+
range.x(),
102+
range.y(),
103+
this.enableTimestampPicos);
90104
}
91105

92106
@Override
@@ -96,13 +110,32 @@ public void free() {
96110
markInvalid();
97111
}
98112

113+
@Override
114+
protected Class<?> getTargetClass() {
115+
LOG.finestTrace("getTargetClass");
116+
if (this.enableTimestampPicos && BigQueryArrowResultSet.isPicosecondTimestamp(this.schema)) {
117+
return String.class;
118+
}
119+
return super.getTargetClass();
120+
}
121+
99122
@Override
100123
Object getCoercedValue(int index) throws SQLException {
101124
LOG.finestTrace("getCoercedValue");
102125
Object value = this.values.get(index);
103-
return this.arrayOfStruct
104-
? new BigQueryArrowStruct(
105-
schema.getSubFields(), (JsonStringHashMap<?, ?>) value, this.LOG.getArrowStructLogger())
106-
: BigQueryTypeRegistry.convert(value, this.schema.getType().getStandardType(), null);
126+
if (value instanceof Text) {
127+
value = value.toString();
128+
}
129+
if (this.arrayOfStruct) {
130+
return new BigQueryArrowStruct(
131+
schema.getSubFields(),
132+
(JsonStringHashMap<?, ?>) value,
133+
this.LOG.getArrowStructLogger(),
134+
this.enableTimestampPicos);
135+
}
136+
if (this.enableTimestampPicos && BigQueryArrowResultSet.isPicosecondTimestamp(this.schema)) {
137+
return BigQueryTemporalUtility.formatTimestampValue(value, true);
138+
}
139+
return BigQueryTypeRegistry.convert(value, this.schema.getType().getStandardType(), null);
107140
}
108141
}

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java‎

Lines changed: 93 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
5252
import org.apache.arrow.vector.util.JsonStringArrayList;
5353
import org.apache.arrow.vector.util.JsonStringHashMap;
54+
import org.apache.arrow.vector.util.Text;
5455

5556
/** {@link ResultSet} Implementation for Arrow datasource (Using Storage Read APIs) */
5657
class BigQueryArrowResultSet extends BigQueryBaseResultSet {
@@ -85,6 +86,7 @@ class BigQueryArrowResultSet extends BigQueryBaseResultSet {
8586
private VectorLoader vectorLoader;
8687
// producer task's reference
8788
private final Future<?> ownedTask;
89+
private final boolean enableTimestampPicos;
8890

8991
private BigQueryArrowResultSet(
9092
Schema schema,
@@ -98,7 +100,8 @@ private BigQueryArrowResultSet(
98100
int toIndexExclusive,
99101
Future<?> ownedTask,
100102
BigQuery bigQuery,
101-
Job job)
103+
Job job,
104+
boolean enableTimestampPicos)
102105
throws SQLException {
103106
super(bigQuery, statement, schema, isNested, job);
104107
LOG.finestTrace("<init>");
@@ -109,6 +112,8 @@ private BigQueryArrowResultSet(
109112
this.toIndexExclusive = toIndexExclusive;
110113
this.nestedRowIndex = fromIndex - 1;
111114
this.ownedTask = ownedTask;
115+
this.enableTimestampPicos =
116+
statement != null ? statement.isEnableTimestampPicos() : enableTimestampPicos;
112117
if (!isNested && arrowSchema != null) {
113118
try {
114119
this.arrowDeserializer = new ArrowDeserializer(arrowSchema);
@@ -158,7 +163,8 @@ static BigQueryArrowResultSet of(
158163
-1,
159164
ownedTask,
160165
bigQuery,
161-
job);
166+
job,
167+
false);
162168
}
163169

164170
BigQueryArrowResultSet() throws SQLException {
@@ -172,10 +178,15 @@ static BigQueryArrowResultSet of(
172178
this.arrowDeserializer = null;
173179
this.vectorSchemaRoot = null;
174180
this.vectorLoader = null;
181+
this.enableTimestampPicos = false;
175182
}
176183

177184
static BigQueryArrowResultSet getNestedResultSet(
178-
Schema schema, BigQueryArrowBatchWrapper nestedBatch, int fromIndex, int toIndexExclusive)
185+
Schema schema,
186+
BigQueryArrowBatchWrapper nestedBatch,
187+
int fromIndex,
188+
int toIndexExclusive,
189+
boolean enableTimestampPicos)
179190
throws SQLException {
180191
return new BigQueryArrowResultSet(
181192
schema,
@@ -189,7 +200,8 @@ static BigQueryArrowResultSet getNestedResultSet(
189200
toIndexExclusive,
190201
null,
191202
null,
192-
null);
203+
null,
204+
enableTimestampPicos);
193205
}
194206

195207
private class ArrowDeserializer implements AutoCloseable {
@@ -345,10 +357,20 @@ private Object getObjectInternal(int columnIndex) throws SQLException {
345357
value = LocalDate.ofEpochDay(((Integer) value).longValue());
346358
}
347359
}
360+
if (value instanceof Text) {
361+
value = value.toString();
362+
}
348363
setWasNull(value);
349364
return value;
350365
}
351366

367+
static boolean isPicosecondTimestamp(Field field) {
368+
return field != null
369+
&& field.getType().getStandardType() == StandardSQLTypeName.TIMESTAMP
370+
&& field.getTimestampPrecision() != null
371+
&& field.getTimestampPrecision() > 6;
372+
}
373+
352374
@Override
353375
public Object getObject(int columnIndex) throws SQLException {
354376

@@ -370,12 +392,16 @@ public Object getObject(int columnIndex) throws SQLException {
370392
return new BigQueryArrowStruct(
371393
arrayField.getSubFields(),
372394
(JsonStringHashMap<?, ?>) value,
373-
this.LOG.getArrowStructLogger());
395+
this.LOG.getArrowStructLogger(),
396+
this.enableTimestampPicos);
374397
}
375398
if (value instanceof Integer
376399
&& arrayField.getType().getStandardType() == StandardSQLTypeName.DATE) {
377400
value = LocalDate.ofEpochDay(((Integer) value).longValue());
378401
}
402+
if (this.enableTimestampPicos && isPicosecondTimestamp(arrayField)) {
403+
return BigQueryTemporalUtility.formatTimestampValue(value, true);
404+
}
379405
return BigQueryTypeRegistry.convert(value, arrayField.getType().getStandardType(), null);
380406
}
381407

@@ -388,61 +414,87 @@ public Object getObject(int columnIndex) throws SQLException {
388414
|| elementTypeName == StandardSQLTypeName.BIGNUMERIC) {
389415
JsonStringArrayList<BigDecimal> newList = new JsonStringArrayList<>();
390416
for (Object item : originalList) {
391-
if (item != null) {
392-
newList.add(((BigDecimal) item).stripTrailingZeros());
393-
} else {
417+
if (item == null) {
394418
newList.add(null);
419+
continue;
395420
}
421+
newList.add(((BigDecimal) item).stripTrailingZeros());
396422
}
397-
return new BigQueryArrowArray(fieldSchema, newList, this.LOG.getArrowArrayLogger());
398-
} else if (elementTypeName == StandardSQLTypeName.RANGE) {
423+
return new BigQueryArrowArray(
424+
fieldSchema, newList, this.LOG.getArrowArrayLogger(), this.enableTimestampPicos);
425+
}
426+
if (elementTypeName == StandardSQLTypeName.RANGE) {
399427
JsonStringArrayList<String> newList = new JsonStringArrayList<>();
400428
for (Object item : originalList) {
401-
if (item != null) {
402-
JsonStringHashMap<?, ?> rangeMap = (JsonStringHashMap<?, ?>) item;
403-
Object start = rangeMap.get("start");
404-
Object end = rangeMap.get("end");
429+
if (item == null) {
430+
newList.add(null);
431+
continue;
432+
}
433+
JsonStringHashMap<?, ?> rangeMap = (JsonStringHashMap<?, ?>) item;
434+
Object start = rangeMap.get("start");
435+
Object end = rangeMap.get("end");
405436

406-
Object representativeElement = (start != null) ? start : end;
407-
StandardSQLTypeName rangeElementType = getElementTypeFromValue(representativeElement);
437+
Object representativeElement = (start != null) ? start : end;
438+
StandardSQLTypeName rangeElementType =
439+
getRangeElementType(fieldSchema, representativeElement);
408440

409-
String formattedStart = formatRangeElement(start, rangeElementType);
410-
String formattedEnd = formatRangeElement(end, rangeElementType);
441+
String formattedStart = formatRangeElement(start, rangeElementType);
442+
String formattedEnd = formatRangeElement(end, rangeElementType);
411443

412-
newList.add(String.format("[%s, %s)", formattedStart, formattedEnd));
413-
} else {
414-
newList.add(null);
415-
}
444+
newList.add(String.format("[%s, %s)", formattedStart, formattedEnd));
416445
}
417-
return new BigQueryArrowArray(fieldSchema, newList, this.LOG.getArrowArrayLogger());
446+
return new BigQueryArrowArray(
447+
fieldSchema, newList, this.LOG.getArrowArrayLogger(), this.enableTimestampPicos);
418448
}
419-
return new BigQueryArrowArray(fieldSchema, originalList, this.LOG.getArrowArrayLogger());
420-
} else if (isStruct(fieldSchema)) {
449+
return new BigQueryArrowArray(
450+
fieldSchema, originalList, this.LOG.getArrowArrayLogger(), this.enableTimestampPicos);
451+
}
452+
453+
if (isStruct(fieldSchema)) {
421454
return new BigQueryArrowStruct(
422455
fieldSchema.getSubFields(),
423456
(JsonStringHashMap<?, ?>) value,
424-
this.LOG.getArrowStructLogger());
425-
} else if (fieldSchema.getType().getStandardType() == StandardSQLTypeName.RANGE) {
457+
this.LOG.getArrowStructLogger(),
458+
this.enableTimestampPicos);
459+
}
460+
461+
if (fieldSchema.getType().getStandardType() == StandardSQLTypeName.RANGE) {
426462
JsonStringHashMap<?, ?> rangeMap = (JsonStringHashMap<?, ?>) value;
427463
Object start = rangeMap.get("start");
428464
Object end = rangeMap.get("end");
429465

430466
Object representativeElement = (start != null) ? start : end;
431-
StandardSQLTypeName elementType = getElementTypeFromValue(representativeElement);
467+
StandardSQLTypeName elementType = getRangeElementType(fieldSchema, representativeElement);
432468

433469
String formattedStart = formatRangeElement(start, elementType);
434470
String formattedEnd = formatRangeElement(end, elementType);
435471

436472
return String.format("[%s, %s)", formattedStart, formattedEnd);
437-
} else {
438-
if ((fieldSchema.getType().getStandardType() == StandardSQLTypeName.NUMERIC
439-
|| fieldSchema.getType().getStandardType() == StandardSQLTypeName.BIGNUMERIC)
440-
&& value instanceof BigDecimal) {
441-
// The Arrow DecimalVector may return a BigDecimal with a larger scale than necessary.
442-
// Strip trailing zeros to match JSON API and CLI output
443-
return ((BigDecimal) value).stripTrailingZeros();
444-
}
445-
return BigQueryTypeRegistry.convert(value, fieldSchema.getType().getStandardType(), null);
473+
}
474+
475+
if ((fieldSchema.getType().getStandardType() == StandardSQLTypeName.NUMERIC
476+
|| fieldSchema.getType().getStandardType() == StandardSQLTypeName.BIGNUMERIC)
477+
&& value instanceof BigDecimal) {
478+
// The Arrow DecimalVector may return a BigDecimal with a larger scale than necessary.
479+
// Strip trailing zeros to match JSON API and CLI output
480+
return ((BigDecimal) value).stripTrailingZeros();
481+
}
482+
if (this.enableTimestampPicos && isPicosecondTimestamp(fieldSchema)) {
483+
return BigQueryTemporalUtility.formatTimestampValue(value, true);
484+
}
485+
return BigQueryTypeRegistry.convert(value, fieldSchema.getType().getStandardType(), null);
486+
}
487+
488+
private StandardSQLTypeName getRangeElementType(Field field, Object representativeElement) {
489+
if (field == null
490+
|| field.getRangeElementType() == null
491+
|| field.getRangeElementType().getType() == null) {
492+
return getElementTypeFromValue(representativeElement);
493+
}
494+
try {
495+
return StandardSQLTypeName.valueOf(field.getRangeElementType().getType());
496+
} catch (IllegalArgumentException ignored) {
497+
return getElementTypeFromValue(representativeElement);
446498
}
447499
}
448500

@@ -467,6 +519,9 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp
467519
if (element == null) {
468520
return "UNBOUNDED";
469521
}
522+
if (element instanceof Text) {
523+
element = element.toString();
524+
}
470525
switch (elementType) {
471526
case DATE:
472527
// Arrow gives DATE as an Integer (days since epoch)
@@ -476,9 +531,7 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp
476531
Timestamp dtTs = Timestamp.valueOf((LocalDateTime) element);
477532
return BigQueryTypeRegistry.convert(dtTs, String.class);
478533
case TIMESTAMP:
479-
// Arrow gives TIMESTAMP as a Long (microseconds since epoch)
480-
Timestamp ts = BigQueryTypeRegistry.convert((Long) element, Timestamp.class);
481-
return BigQueryTypeRegistry.convert(ts, String.class);
534+
return BigQueryTemporalUtility.formatTimestampValue(element, this.enableTimestampPicos);
482535
default:
483536
// Fallback for any other unexpected type
484537
return element.toString();

0 commit comments

Comments
 (0)