diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArray.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArray.java index 7c60866cf15a..01fb49e64476 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArray.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArray.java @@ -23,6 +23,7 @@ import java.sql.SQLException; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.JsonStringHashMap; +import org.apache.arrow.vector.util.Text; /** * An implementation of {@link BigQueryBaseArray} used to represent Array values from Arrow data. @@ -31,13 +32,16 @@ class BigQueryArrowArray extends BigQueryBaseArray { private JsonStringArrayList values; - public BigQueryArrowArray(Field schema, JsonStringArrayList values) { - this(schema, values, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowArray.class)); + BigQueryArrowArray(Field schema, JsonStringArrayList values) { + this(schema, values, false, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowArray.class)); } - public BigQueryArrowArray( - Field schema, JsonStringArrayList values, BigQueryJdbcResultSetLogger log) { - super(schema, log); + BigQueryArrowArray( + Field schema, + JsonStringArrayList values, + boolean enableTimestampPicos, + BigQueryJdbcResultSetLogger log) { + super(schema, enableTimestampPicos, log); this.values = values; } @@ -72,7 +76,11 @@ public ResultSet getResultSet() throws SQLException { BigQueryArrowBatchWrapper arrowBatchWrapper = BigQueryArrowBatchWrapper.getNestedFieldValueListWrapper(values); return BigQueryArrowResultSet.getNestedResultSet( - Schema.of(singleElementSchema()), arrowBatchWrapper, 0, this.values.size()); + Schema.of(singleElementSchema()), + arrowBatchWrapper, + 0, + this.values.size(), + this.enableTimestampPicos); } @Override @@ -86,7 +94,11 @@ public ResultSet getResultSet(long index, int count) throws SQLException { BigQueryArrowBatchWrapper arrowBatchWrapper = BigQueryArrowBatchWrapper.getNestedFieldValueListWrapper(values); return BigQueryArrowResultSet.getNestedResultSet( - Schema.of(singleElementSchema()), arrowBatchWrapper, range.x(), range.y()); + Schema.of(singleElementSchema()), + arrowBatchWrapper, + range.x(), + range.y(), + this.enableTimestampPicos); } @Override @@ -100,9 +112,19 @@ public void free() { Object getCoercedValue(int index) throws SQLException { LOG.finestTrace("getCoercedValue"); Object value = this.values.get(index); - return this.arrayOfStruct - ? new BigQueryArrowStruct( - schema.getSubFields(), (JsonStringHashMap) value, this.LOG.getArrowStructLogger()) - : BigQueryTypeRegistry.convert(value, this.schema.getType().getStandardType(), null); + if (value instanceof Text) { + value = value.toString(); + } + if (this.arrayOfStruct) { + return new BigQueryArrowStruct( + schema.getSubFields(), + (JsonStringHashMap) value, + this.enableTimestampPicos, + this.LOG.getArrowStructLogger()); + } + if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(this.schema)) { + return BigQueryTemporalUtility.formatTimestampValue(value, true); + } + return BigQueryTypeRegistry.convert(value, this.schema.getType().getStandardType(), null); } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index be20b5628035..f76cfaba4d5e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -51,6 +51,7 @@ import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.JsonStringHashMap; +import org.apache.arrow.vector.util.Text; /** {@link ResultSet} Implementation for Arrow datasource (Using Storage Read APIs) */ class BigQueryArrowResultSet extends BigQueryBaseResultSet { @@ -85,6 +86,7 @@ class BigQueryArrowResultSet extends BigQueryBaseResultSet { private VectorLoader vectorLoader; // producer task's reference private final Future ownedTask; + private final boolean enableTimestampPicos; private BigQueryArrowResultSet( Schema schema, @@ -98,7 +100,8 @@ private BigQueryArrowResultSet( int toIndexExclusive, Future ownedTask, BigQuery bigQuery, - Job job) + Job job, + boolean enableTimestampPicos) throws SQLException { super(bigQuery, statement, schema, isNested, job); LOG.finestTrace(""); @@ -109,6 +112,7 @@ private BigQueryArrowResultSet( this.toIndexExclusive = toIndexExclusive; this.nestedRowIndex = fromIndex - 1; this.ownedTask = ownedTask; + this.enableTimestampPicos = enableTimestampPicos; if (!isNested && arrowSchema != null) { try { this.arrowDeserializer = new ArrowDeserializer(arrowSchema); @@ -158,7 +162,8 @@ static BigQueryArrowResultSet of( -1, ownedTask, bigQuery, - job); + job, + statement != null && statement.isEnableTimestampPicos()); } BigQueryArrowResultSet() throws SQLException { @@ -172,10 +177,15 @@ static BigQueryArrowResultSet of( this.arrowDeserializer = null; this.vectorSchemaRoot = null; this.vectorLoader = null; + this.enableTimestampPicos = false; } static BigQueryArrowResultSet getNestedResultSet( - Schema schema, BigQueryArrowBatchWrapper nestedBatch, int fromIndex, int toIndexExclusive) + Schema schema, + BigQueryArrowBatchWrapper nestedBatch, + int fromIndex, + int toIndexExclusive, + boolean enableTimestampPicos) throws SQLException { return new BigQueryArrowResultSet( schema, @@ -189,7 +199,8 @@ static BigQueryArrowResultSet getNestedResultSet( toIndexExclusive, null, null, - null); + null, + enableTimestampPicos); } private class ArrowDeserializer implements AutoCloseable { @@ -345,6 +356,9 @@ private Object getObjectInternal(int columnIndex) throws SQLException { value = LocalDate.ofEpochDay(((Integer) value).longValue()); } } + if (value instanceof Text) { + value = value.toString(); + } setWasNull(value); return value; } @@ -370,12 +384,16 @@ public Object getObject(int columnIndex) throws SQLException { return new BigQueryArrowStruct( arrayField.getSubFields(), (JsonStringHashMap) value, + this.enableTimestampPicos, this.LOG.getArrowStructLogger()); } if (value instanceof Integer && arrayField.getType().getStandardType() == StandardSQLTypeName.DATE) { value = LocalDate.ofEpochDay(((Integer) value).longValue()); } + if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(arrayField)) { + return BigQueryTemporalUtility.formatTimestampValue(value, true); + } return BigQueryTypeRegistry.convert(value, arrayField.getType().getStandardType(), null); } @@ -388,61 +406,87 @@ public Object getObject(int columnIndex) throws SQLException { || elementTypeName == StandardSQLTypeName.BIGNUMERIC) { JsonStringArrayList newList = new JsonStringArrayList<>(); for (Object item : originalList) { - if (item != null) { - newList.add(((BigDecimal) item).stripTrailingZeros()); - } else { + if (item == null) { newList.add(null); + continue; } + newList.add(((BigDecimal) item).stripTrailingZeros()); } - return new BigQueryArrowArray(fieldSchema, newList, this.LOG.getArrowArrayLogger()); - } else if (elementTypeName == StandardSQLTypeName.RANGE) { + return new BigQueryArrowArray( + fieldSchema, newList, this.enableTimestampPicos, this.LOG.getArrowArrayLogger()); + } + if (elementTypeName == StandardSQLTypeName.RANGE) { JsonStringArrayList newList = new JsonStringArrayList<>(); for (Object item : originalList) { - if (item != null) { - JsonStringHashMap rangeMap = (JsonStringHashMap) item; - Object start = rangeMap.get("start"); - Object end = rangeMap.get("end"); + if (item == null) { + newList.add(null); + continue; + } + JsonStringHashMap rangeMap = (JsonStringHashMap) item; + Object start = rangeMap.get("start"); + Object end = rangeMap.get("end"); - Object representativeElement = (start != null) ? start : end; - StandardSQLTypeName rangeElementType = getElementTypeFromValue(representativeElement); + Object representativeElement = (start != null) ? start : end; + StandardSQLTypeName rangeElementType = + getRangeElementType(fieldSchema, representativeElement); - String formattedStart = formatRangeElement(start, rangeElementType); - String formattedEnd = formatRangeElement(end, rangeElementType); + String formattedStart = formatRangeElement(start, rangeElementType); + String formattedEnd = formatRangeElement(end, rangeElementType); - newList.add(String.format("[%s, %s)", formattedStart, formattedEnd)); - } else { - newList.add(null); - } + newList.add(String.format("[%s, %s)", formattedStart, formattedEnd)); } - return new BigQueryArrowArray(fieldSchema, newList, this.LOG.getArrowArrayLogger()); + return new BigQueryArrowArray( + fieldSchema, newList, this.enableTimestampPicos, this.LOG.getArrowArrayLogger()); } - return new BigQueryArrowArray(fieldSchema, originalList, this.LOG.getArrowArrayLogger()); - } else if (isStruct(fieldSchema)) { + return new BigQueryArrowArray( + fieldSchema, originalList, this.enableTimestampPicos, this.LOG.getArrowArrayLogger()); + } + + if (isStruct(fieldSchema)) { return new BigQueryArrowStruct( fieldSchema.getSubFields(), (JsonStringHashMap) value, + this.enableTimestampPicos, this.LOG.getArrowStructLogger()); - } else if (fieldSchema.getType().getStandardType() == StandardSQLTypeName.RANGE) { + } + + if (fieldSchema.getType().getStandardType() == StandardSQLTypeName.RANGE) { JsonStringHashMap rangeMap = (JsonStringHashMap) value; Object start = rangeMap.get("start"); Object end = rangeMap.get("end"); Object representativeElement = (start != null) ? start : end; - StandardSQLTypeName elementType = getElementTypeFromValue(representativeElement); + StandardSQLTypeName elementType = getRangeElementType(fieldSchema, representativeElement); String formattedStart = formatRangeElement(start, elementType); String formattedEnd = formatRangeElement(end, elementType); return String.format("[%s, %s)", formattedStart, formattedEnd); - } else { - if ((fieldSchema.getType().getStandardType() == StandardSQLTypeName.NUMERIC - || fieldSchema.getType().getStandardType() == StandardSQLTypeName.BIGNUMERIC) - && value instanceof BigDecimal) { - // The Arrow DecimalVector may return a BigDecimal with a larger scale than necessary. - // Strip trailing zeros to match JSON API and CLI output - return ((BigDecimal) value).stripTrailingZeros(); - } - return BigQueryTypeRegistry.convert(value, fieldSchema.getType().getStandardType(), null); + } + + if ((fieldSchema.getType().getStandardType() == StandardSQLTypeName.NUMERIC + || fieldSchema.getType().getStandardType() == StandardSQLTypeName.BIGNUMERIC) + && value instanceof BigDecimal) { + // The Arrow DecimalVector may return a BigDecimal with a larger scale than necessary. + // Strip trailing zeros to match JSON API and CLI output + return ((BigDecimal) value).stripTrailingZeros(); + } + if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(fieldSchema)) { + return BigQueryTemporalUtility.formatTimestampValue(value, true); + } + return BigQueryTypeRegistry.convert(value, fieldSchema.getType().getStandardType(), null); + } + + private StandardSQLTypeName getRangeElementType(Field field, Object representativeElement) { + if (field == null + || field.getRangeElementType() == null + || field.getRangeElementType().getType() == null) { + return getElementTypeFromValue(representativeElement); + } + try { + return StandardSQLTypeName.valueOf(field.getRangeElementType().getType()); + } catch (IllegalArgumentException ignored) { + return getElementTypeFromValue(representativeElement); } } @@ -467,6 +511,9 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp if (element == null) { return "UNBOUNDED"; } + if (element instanceof Text) { + element = element.toString(); + } switch (elementType) { case DATE: // Arrow gives DATE as an Integer (days since epoch) @@ -476,9 +523,7 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp Timestamp dtTs = Timestamp.valueOf((LocalDateTime) element); return BigQueryTypeRegistry.convert(dtTs, String.class); case TIMESTAMP: - // Arrow gives TIMESTAMP as a Long (microseconds since epoch) - Timestamp ts = BigQueryTypeRegistry.convert((Long) element, Timestamp.class); - return BigQueryTypeRegistry.convert(ts, String.class); + return BigQueryTemporalUtility.formatTimestampValue(element, this.enableTimestampPicos); default: // Fallback for any other unexpected type return element.toString(); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStruct.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStruct.java index 375c0619703a..ef6c3f69294c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStruct.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStruct.java @@ -28,6 +28,7 @@ import java.util.List; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.JsonStringHashMap; +import org.apache.arrow.vector.util.Text; /** * An implementation of {@link BigQueryBaseStruct} used to represent Struct values from Arrow data. @@ -39,12 +40,15 @@ class BigQueryArrowStruct extends BigQueryBaseStruct { private final JsonStringHashMap values; BigQueryArrowStruct(FieldList schema, JsonStringHashMap values) { - this(schema, values, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowStruct.class)); + this(schema, values, false, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowStruct.class)); } BigQueryArrowStruct( - FieldList schema, JsonStringHashMap values, BigQueryJdbcResultSetLogger log) { - super(log); + FieldList schema, + JsonStringHashMap values, + boolean enableTimestampPicos, + BigQueryJdbcResultSetLogger log) { + super(enableTimestampPicos, log); this.schema = schema; this.values = values; } @@ -76,21 +80,31 @@ public Object[] getAttributes() throws SQLException { private Object getValue(Field currentSchema, Object currentValue) throws SQLException { LOG.finestTrace("getValue"); + if (currentValue instanceof Text) { + currentValue = currentValue.toString(); + } if (isArray(currentSchema)) { return new BigQueryArrowArray( - currentSchema, (JsonStringArrayList) currentValue, this.LOG.getArrowArrayLogger()); - } else if (isStruct(currentSchema)) { + currentSchema, + (JsonStringArrayList) currentValue, + this.enableTimestampPicos, + this.LOG.getArrowArrayLogger()); + } + if (isStruct(currentSchema)) { return new BigQueryArrowStruct( currentSchema.getSubFields(), (JsonStringHashMap) currentValue, + this.enableTimestampPicos, this.LOG.getArrowStructLogger()); - } else { - if (currentValue instanceof Integer - && currentSchema.getType().getStandardType() == StandardSQLTypeName.DATE) { - currentValue = LocalDate.ofEpochDay(((Integer) currentValue).longValue()); - } - return BigQueryTypeRegistry.convert( - currentValue, currentSchema.getType().getStandardType(), null); } + if (currentValue instanceof Integer + && currentSchema.getType().getStandardType() == StandardSQLTypeName.DATE) { + currentValue = LocalDate.ofEpochDay(((Integer) currentValue).longValue()); + } + if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(currentSchema)) { + return BigQueryTemporalUtility.formatTimestampValue(currentValue, true); + } + return BigQueryTypeRegistry.convert( + currentValue, currentSchema.getType().getStandardType(), null); } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseArray.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseArray.java index d677f3182e8e..ac207e78e261 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseArray.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseArray.java @@ -46,12 +46,18 @@ abstract class BigQueryBaseArray implements java.sql.Array { protected final boolean arrayOfStruct; private boolean valid; protected Field schema; + protected final boolean enableTimestampPicos; BigQueryBaseArray(Field schema, BigQueryJdbcResultSetLogger log) { + this(schema, false, log); + } + + BigQueryBaseArray(Field schema, boolean enableTimestampPicos, BigQueryJdbcResultSetLogger log) { this.LOG = log; this.schema = schema; this.arrayOfStruct = isStruct(schema); this.valid = true; + this.enableTimestampPicos = enableTimestampPicos; } @Override @@ -142,6 +148,9 @@ protected Tuple createRange(long index, int count, int size) protected Class getTargetClass() { LOG.finestTrace("getTargetClass"); + if (this.enableTimestampPicos && BigQueryTemporalUtility.isPicosecondTimestamp(this.schema)) { + return String.class; + } return this.arrayOfStruct ? Struct.class : BigQueryTypeRegistry.toJavaClass(this.schema.getType().getStandardType()); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseStruct.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseStruct.java index 524e8cf5cb81..7caf0cdfe9b2 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseStruct.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseStruct.java @@ -37,9 +37,15 @@ */ abstract class BigQueryBaseStruct implements java.sql.Struct { protected final BigQueryJdbcResultSetLogger LOG; + protected final boolean enableTimestampPicos; BigQueryBaseStruct(BigQueryJdbcResultSetLogger log) { + this(false, log); + } + + BigQueryBaseStruct(boolean enableTimestampPicos, BigQueryJdbcResultSetLogger log) { this.LOG = log; + this.enableTimestampPicos = enableTimestampPicos; } abstract FieldList getSchema(); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index c0e5820289bc..d1348f6d8091 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -50,12 +50,14 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcSqlSyntaxErrorException; import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch; import com.google.cloud.bigquery.storage.v1.ArrowSchema; +import com.google.cloud.bigquery.storage.v1.ArrowSerializationOptions; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; import com.google.cloud.bigquery.storage.v1.DataFormat; import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; import com.google.cloud.bigquery.storage.v1.ReadSession; +import com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions; import com.google.common.annotations.VisibleForTesting; import com.google.common.util.concurrent.Uninterruptibles; import io.grpc.Status; @@ -883,6 +885,10 @@ ResultSet processArrowResultSet(TableResult results, Job job) throws SQLExceptio // format ReadSession.Builder sessionBuilder = ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW); + TableReadOptions readOptions = buildTableReadOptions(); + if (readOptions != null) { + sessionBuilder.setReadOptions(readOptions); + } CreateReadSessionRequest.Builder builder = CreateReadSessionRequest.newBuilder() @@ -937,6 +943,19 @@ ResultSet processArrowResultSet(TableResult results, Job job) throws SQLExceptio } } + private TableReadOptions buildTableReadOptions() { + if (!isEnableTimestampPicos()) { + return null; + } + return TableReadOptions.newBuilder() + .setArrowSerializationOptions( + ArrowSerializationOptions.newBuilder() + .setPicosTimestampPrecision( + ArrowSerializationOptions.PicosTimestampPrecision.TIMESTAMP_PRECISION_PICOS) + .build()) + .build(); + } + /** Asynchronously reads results and populates an arrow record queue */ @InternalApi Future populateArrowBufferedQueue( diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtility.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtility.java index f493857eb9e0..735e672687f8 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtility.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtility.java @@ -16,6 +16,8 @@ package com.google.cloud.bigquery.jdbc; +import com.google.cloud.bigquery.Field; +import com.google.cloud.bigquery.StandardSQLTypeName; import com.google.cloud.bigquery.exception.BigQueryJdbcException; import java.math.BigDecimal; import java.math.RoundingMode; @@ -306,6 +308,34 @@ static String formatTimestampStringFromMicroseconds( return sb.toString(); } + /** + * Formats a timestamp value (which may be a {@link Long} epoch microsecond, an ISO-8601 string, a + * {@link java.sql.Timestamp}, or an epoch decimal string) into a standard UTC JDBC timestamp + * string with 6 or 12 fractional digits according to {@code enableTimestampPicos}. + */ + static String formatTimestampValue(Object value, boolean enableTimestampPicos) + throws BigQueryJdbcException { + if (value == null) { + return null; + } + if (value instanceof Long) { + return formatTimestampStringFromMicroseconds((Long) value, enableTimestampPicos); + } + String str = value.toString(); + if (str.indexOf(':') >= 0) { + return formatTimestampStringFromIso(str, enableTimestampPicos); + } + return formatTimestampStringFromEpochDecimal(str, enableTimestampPicos); + } + + static boolean isPicosecondTimestamp(Field field) { + return field != null + && field.getType() != null + && field.getType().getStandardType() == StandardSQLTypeName.TIMESTAMP + && field.getTimestampPrecision() != null + && field.getTimestampPrecision() > 6; + } + private static StringBuilder formatDateTimeBase(LocalDateTime dt, int scale) { StringBuilder sb = new StringBuilder(scale == 12 ? 32 : 26); BASE_FORMATTER.formatTo(dt, sb); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArrayOfPrimitivesTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArrayOfPrimitivesTest.java index 07b5d28a3d9c..c9b71da98fe9 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArrayOfPrimitivesTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowArrayOfPrimitivesTest.java @@ -63,6 +63,7 @@ import java.util.stream.Stream; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.Text; +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.function.Executable; import org.junit.jupiter.params.ParameterizedTest; @@ -409,4 +410,66 @@ private void ensureArrayIsInvalid(Executable block) { Exception exception = assertThrows(IllegalStateException.class, block); assertThat(exception.getMessage()).isEqualTo(INVALID_ARRAY); } + + @Test + public void testArrowArrayTimestampPicosEnabled() throws SQLException { + Field field = + Field.newBuilder("picosArray", StandardSQLTypeName.TIMESTAMP) + .setMode(Field.Mode.REPEATED) + .setTimestampPrecision(12L) + .build(); + JsonStringArrayList values = new JsonStringArrayList<>(); + values.add(new Text("2026-04-08T10:00:00.123456789123Z")); + values.add(new Text("2026-04-08T11:00:00.987654321012Z")); + + BigQueryArrowArray array = + new BigQueryArrowArray( + field, values, true, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowArray.class)); + + assertThat(array.getBaseTypeName()).isEqualTo("TIMESTAMP"); + assertThat(array.getBaseType()).isEqualTo(Types.TIMESTAMP); + + Object result = array.getArray(); + assertThat(result).isInstanceOf(String[].class); + assertThat((String[]) result) + .asList() + .containsExactly("2026-04-08 10:00:00.123456789123", "2026-04-08 11:00:00.987654321012") + .inOrder(); + + ResultSet rs = array.getResultSet(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getInt(1)).isEqualTo(1); + assertThat(rs.getString(2)).isEqualTo("2026-04-08 10:00:00.123456789123"); + assertThat(rs.getObject(2)).isEqualTo("2026-04-08 10:00:00.123456789123"); + + assertThat(rs.next()).isTrue(); + assertThat(rs.getInt(1)).isEqualTo(2); + assertThat(rs.getString(2)).isEqualTo("2026-04-08 11:00:00.987654321012"); + assertThat(rs.getObject(2)).isEqualTo("2026-04-08 11:00:00.987654321012"); + } + + @Test + public void testArrowArrayTimestampPicosDisabled() throws SQLException { + Field field = + Field.newBuilder("picosArray", StandardSQLTypeName.TIMESTAMP) + .setMode(Field.Mode.REPEATED) + .setTimestampPrecision(12L) + .build(); + JsonStringArrayList values = new JsonStringArrayList<>(); + values.add(new Text("2026-04-08T10:00:00.123456789123Z")); + + BigQueryArrowArray array = + new BigQueryArrowArray( + field, values, false, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowArray.class)); + + Object result = array.getArray(); + assertThat(result).isInstanceOf(Timestamp[].class); + Timestamp expectedTs = Timestamp.valueOf("2026-04-08 10:00:00.123456789"); + assertThat((Timestamp[]) result).asList().containsExactly(expectedTs); + + ResultSet rs = array.getResultSet(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString(2)).isEqualTo("2026-04-08 10:00:00.123456"); + assertThat(rs.getObject(2)).isEqualTo(expectedTs); + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java index 1ffd05ff4cd6..08887854f363 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java @@ -22,6 +22,7 @@ import static org.apache.arrow.vector.types.Types.MinorType.INT; import static org.apache.arrow.vector.types.Types.MinorType.VARCHAR; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import com.google.cloud.bigquery.Field; @@ -255,7 +256,8 @@ public void setUp() throws SQLException, IOException { arraySchema, BigQueryArrowBatchWrapper.getNestedFieldValueListWrapper(jsonStringArrayList), 0, - jsonStringArrayList.size()); + jsonStringArrayList.size(), + false); } @Test @@ -437,5 +439,111 @@ private int resultSetRowCount(BigQueryArrowResultSet resultSet) throws SQLExcept return rowCount; } - // TODO: Unit Test for iteration and getters + private VectorSchemaRoot getPicosVectorSchemaRoot() { + RootAllocator allocator = new RootAllocator(); + VarCharVector picosCol = new VarCharVector("picosTimestamp", allocator); + picosCol.allocateNew(1); + picosCol.set(0, new Text("2026-04-08T10:00:00.123456789123Z")); + picosCol.setValueCount(1); + return new VectorSchemaRoot(ImmutableList.of(picosCol)); + } + + @Test + public void testTimestampPicosEnabled() throws SQLException, IOException { + VectorSchemaRoot root = getPicosVectorSchemaRoot(); + ArrowRecordBatch batch = + ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(serializeVectorSchemaRoot(root)) + .build(); + BlockingQueue picosBuffer = new LinkedBlockingDeque<>(); + picosBuffer.add(BigQueryArrowBatchWrapper.of(batch)); + picosBuffer.add(BigQueryArrowBatchWrapper.of(null, true)); + + BigQueryStatement picosStmt = mock(BigQueryStatement.class); + doReturn(true).when(picosStmt).isEnableTimestampPicos(); + + Schema schema = + Schema.of( + Field.newBuilder("picosTimestamp", StandardSQLTypeName.TIMESTAMP) + .setTimestampPrecision(12L) + .build()); + ArrowSchema arrowSchema = + ArrowSchema.newBuilder().setSerializedSchema(serializeSchema(root.getSchema())).build(); + + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + schema, arrowSchema, 1, picosStmt, picosBuffer, mock(Future.class), null); + + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("picosTimestamp")).isEqualTo("2026-04-08 10:00:00.123456789123"); + assertThat(rs.getString(1)).isEqualTo("2026-04-08 10:00:00.123456789123"); + assertThat(rs.getObject("picosTimestamp")).isEqualTo("2026-04-08 10:00:00.123456789123"); + assertThat(rs.getObject(1)).isEqualTo("2026-04-08 10:00:00.123456789123"); + assertThat(rs.getObject("picosTimestamp", String.class)) + .isEqualTo("2026-04-08 10:00:00.123456789123"); + + Timestamp expectedTs = Timestamp.valueOf("2026-04-08 10:00:00.123456789"); + assertThat(rs.getTimestamp("picosTimestamp")).isEqualTo(expectedTs); + assertThat(rs.getObject("picosTimestamp", Timestamp.class)).isEqualTo(expectedTs); + } + + @Test + public void testTimestampPicosDisabled() throws SQLException, IOException { + VectorSchemaRoot root = getPicosVectorSchemaRoot(); + ArrowRecordBatch batch = + ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(serializeVectorSchemaRoot(root)) + .build(); + BlockingQueue picosBuffer = new LinkedBlockingDeque<>(); + picosBuffer.add(BigQueryArrowBatchWrapper.of(batch)); + picosBuffer.add(BigQueryArrowBatchWrapper.of(null, true)); + + BigQueryStatement disabledStmt = mock(BigQueryStatement.class); + doReturn(false).when(disabledStmt).isEnableTimestampPicos(); + + Schema schema = + Schema.of( + Field.newBuilder("picosTimestamp", StandardSQLTypeName.TIMESTAMP) + .setTimestampPrecision(12L) + .build()); + ArrowSchema arrowSchema = + ArrowSchema.newBuilder().setSerializedSchema(serializeSchema(root.getSchema())).build(); + + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + schema, arrowSchema, 1, disabledStmt, picosBuffer, mock(Future.class), null); + + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("picosTimestamp")).isEqualTo("2026-04-08 10:00:00.123456"); + assertThat(rs.getString(1)).isEqualTo("2026-04-08 10:00:00.123456"); + Timestamp expectedTs = Timestamp.valueOf("2026-04-08 10:00:00.123456789"); + assertThat(rs.getObject("picosTimestamp")).isEqualTo(expectedTs); + assertThat(rs.getTimestamp("picosTimestamp")).isEqualTo(expectedTs); + } + + @Test + public void testStandardMicrosecondTimestampWithPicosEnabled() throws SQLException, IOException { + VectorSchemaRoot root = getTestVectorSchemaRoot(); + ArrowRecordBatch batch = + ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(serializeVectorSchemaRoot(root)) + .build(); + BlockingQueue picosBuffer = new LinkedBlockingDeque<>(); + picosBuffer.add(BigQueryArrowBatchWrapper.of(batch)); + picosBuffer.add(BigQueryArrowBatchWrapper.of(null, true)); + + BigQueryStatement picosStmt = mock(BigQueryStatement.class); + doReturn(true).when(picosStmt).isEnableTimestampPicos(); + + ArrowSchema arrowSchema = + ArrowSchema.newBuilder().setSerializedSchema(serializeSchema(root.getSchema())).build(); + + BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + QUERY_SCHEMA, arrowSchema, 1, picosStmt, picosBuffer, mock(Future.class), null); + + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("timeStampField")).isEqualTo("1970-01-01 00:00:00.010000"); + assertThat(rs.getObject("timeStampField")).isEqualTo(new Timestamp(10L)); + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java index e1afd98f9f99..e63bfe79b397 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowStructTest.java @@ -244,4 +244,38 @@ public void getAttributesWithCustomTypeMappingsIsNotSupported() { () -> structWithPrimitiveValues.getAttributes(emptyMap())); assertThat(exception.getMessage()).isEqualTo(CUSTOMER_TYPE_MAPPING_NOT_SUPPORTED); } + + @Test + public void testArrowStructTimestampPicosEnabled() throws SQLException { + Field picosField = Field.newBuilder("picosTs", TIMESTAMP).setTimestampPrecision(12L).build(); + FieldList schema = FieldList.of(picosField); + JsonStringHashMap values = new JsonStringHashMap<>(); + values.put("picosTs", new Text("2026-04-08T10:00:00.123456789123Z")); + + BigQueryArrowStruct struct = + new BigQueryArrowStruct( + schema, values, true, BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowStruct.class)); + + Object[] attributes = struct.getAttributes(); + assertThat(attributes).isEqualTo(new Object[] {"2026-04-08 10:00:00.123456789123"}); + } + + @Test + public void testArrowStructTimestampPicosDisabled() throws SQLException { + Field picosField = Field.newBuilder("picosTs", TIMESTAMP).setTimestampPrecision(12L).build(); + FieldList schema = FieldList.of(picosField); + JsonStringHashMap values = new JsonStringHashMap<>(); + values.put("picosTs", new Text("2026-04-08T10:00:00.123456789123Z")); + + BigQueryArrowStruct struct = + new BigQueryArrowStruct( + schema, + values, + false, + BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowStruct.class)); + + Object[] attributes = struct.getAttributes(); + Timestamp expectedTs = Timestamp.valueOf("2026-04-08 10:00:00.123456789"); + assertThat(attributes).isEqualTo(new Object[] {expectedTs}); + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java index e2e90910edb0..7ce5e6fba529 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java @@ -389,7 +389,7 @@ public void testArrowArrayConnectionIdPropagation() { Field arraySchema = Field.of("name", LegacySQLTypeName.INTEGER); JsonStringArrayList arrayValues = new JsonStringArrayList<>(); arrayValues.add(123L); - BigQueryArrowArray array = new BigQueryArrowArray(arraySchema, arrayValues, logger); + BigQueryArrowArray array = new BigQueryArrowArray(arraySchema, arrayValues, false, logger); assertConnectionIdPropagated( logger, connectionId, "Log from Arrow Array", () -> array.LOG.fine("Log from Arrow Array")); @@ -404,7 +404,8 @@ public void testArrowStructConnectionIdPropagation() { FieldList structSchema = FieldList.of(Field.of("col", LegacySQLTypeName.INTEGER)); JsonStringHashMap structValues = new JsonStringHashMap<>(); structValues.put("col", 456L); - BigQueryArrowStruct struct = new BigQueryArrowStruct(structSchema, structValues, structLogger); + BigQueryArrowStruct struct = + new BigQueryArrowStruct(structSchema, structValues, false, structLogger); assertConnectionIdPropagated( structLogger, @@ -429,7 +430,7 @@ public void testNestedStructInArrowArrayConnectionIdPropagation() throws Excepti JsonStringArrayList nestedValues = new JsonStringArrayList<>(); nestedValues.add(map); BigQueryArrowArray arrayWithNested = - new BigQueryArrowArray(arrayNestedSchema, nestedValues, logger); + new BigQueryArrowArray(arrayNestedSchema, nestedValues, false, logger); Object[] result = (Object[]) arrayWithNested.getArray(); BigQueryArrowStruct nestedStruct = (BigQueryArrowStruct) result[0]; @@ -460,7 +461,7 @@ public void testNestedArrayInArrowStructConnectionIdPropagation() throws Excepti values.put("array_col", arrayVal); BigQueryArrowStruct structWithNestedArray = - new BigQueryArrowStruct(nestedSchema, values, structLogger); + new BigQueryArrowStruct(nestedSchema, values, false, structLogger); Object[] attributes = structWithNestedArray.getAttributes(); BigQueryArrowArray nestedArray = (BigQueryArrowArray) attributes[0]; diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index d72d38580498..f8dcb3ac6e67 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -62,6 +62,7 @@ import com.google.cloud.bigquery.jdbc.BigQueryStatement.JobIdWrapper; import com.google.cloud.bigquery.spi.BigQueryRpcFactory; import com.google.cloud.bigquery.storage.v1.ArrowSchema; +import com.google.cloud.bigquery.storage.v1.ArrowSerializationOptions; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; import com.google.cloud.bigquery.storage.v1.ReadSession; @@ -353,6 +354,63 @@ public void getArrowResultSetTest() throws SQLException { assertThat(resultSet).isNotNull(); assertThat(resultSet).isInstanceOf(BigQueryArrowResultSet.class); assertThat(resultSet.isLast()).isFalse(); // as we have 10 rows + + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(CreateReadSessionRequest.class); + verify(bigQueryStatementSpy).getReadSession(requestCaptor.capture()); + assertThat( + requestCaptor + .getValue() + .getReadSession() + .getReadOptions() + .hasArrowSerializationOptions()) + .isFalse(); + } + + @Test + public void testProcessArrowResultSetWithTimestampPicos() throws SQLException { + doReturn(true).when(bigQueryConnection).isEnableTimestampPicos(); + BigQueryStatement stmt = new BigQueryStatement(bigQueryConnection); + BigQueryStatement bigQueryStatementSpy = Mockito.spy(stmt); + BigQueryReadClient bigQueryReadClient = Mockito.spy(mock(BigQueryReadClient.class)); + Schema schema = Schema.of(fieldList); + ReadSession readSession = ReadSession.getDefaultInstance(); + doReturn(bigQueryReadClient).when(bigQueryStatementSpy).getBigQueryReadClient(); + doReturn(readSession) + .when(bigQueryStatementSpy) + .getReadSession(any(CreateReadSessionRequest.class)); + Future mockWorker = mock(Future.class); + doReturn(mockWorker) + .when(bigQueryStatementSpy) + .populateArrowBufferedQueue( + any(ReadSession.class), any(BlockingQueue.class), any(BigQueryReadClient.class)); + + doReturn(arrowSchema).when(bigQueryStatementSpy).getArrowSchema(any(ReadSession.class)); + + JobId jobId = JobId.of("123"); + TableResult result = Mockito.mock(TableResult.class); + doReturn(schema).when(result).getSchema(); + doReturn(10L).when(result).getTotalRows(); + doReturn(TABLE_ID).when(bigQueryStatementSpy).getDestinationTable(any()); + doReturn(jobId).when(result).getJobId(); + Job job = mock(Job.class); + doReturn(mock(QueryStatistics.class)).when(job).getStatistics(); + doReturn(job).when(bigquery).getJob(jobId); + + ResultSet resultSet = bigQueryStatementSpy.processArrowResultSet(result, null); + assertThat(resultSet).isNotNull(); + + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(CreateReadSessionRequest.class); + verify(bigQueryStatementSpy).getReadSession(requestCaptor.capture()); + assertThat( + requestCaptor + .getValue() + .getReadSession() + .getReadOptions() + .getArrowSerializationOptions() + .getPicosTimestampPrecision()) + .isEqualTo(ArrowSerializationOptions.PicosTimestampPrecision.TIMESTAMP_PRECISION_PICOS); } @Test diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtilityTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtilityTest.java index f89fe157c666..08e92c002287 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtilityTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryTemporalUtilityTest.java @@ -392,4 +392,29 @@ public void testFormatTimestampStringFromMicroseconds() { String negOne = BigQueryTemporalUtility.formatTimestampStringFromMicroseconds(-1L, false); assertThat(negOne).isEqualTo("1969-12-31 23:59:59.999999"); } + + @Test + public void testFormatTimestampValue() throws BigQueryJdbcException { + // Null handling + assertThat(BigQueryTemporalUtility.formatTimestampValue(null, true)).isNull(); + + // Long (epoch microseconds) + assertThat(BigQueryTemporalUtility.formatTimestampValue(1680174859820226L, false)) + .isEqualTo("2023-03-30 11:14:19.820226"); + assertThat(BigQueryTemporalUtility.formatTimestampValue(1680174859820226L, true)) + .isEqualTo("2023-03-30 11:14:19.820226000000"); + + // ISO string with picoseconds + assertThat( + BigQueryTemporalUtility.formatTimestampValue( + "2026-04-08T10:00:00.123456789123Z", false)) + .isEqualTo("2026-04-08 10:00:00.123456"); + assertThat( + BigQueryTemporalUtility.formatTimestampValue("2026-04-08T10:00:00.123456789123Z", true)) + .isEqualTo("2026-04-08 10:00:00.123456789123"); + + // Epoch decimal string + assertThat(BigQueryTemporalUtility.formatTimestampValue("1680174859.123456789123", true)) + .isEqualTo("2023-03-30 11:14:19.123456789123"); + } }