Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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;
}

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand All @@ -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("<init>");
Expand All @@ -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);
Expand Down Expand Up @@ -158,7 +162,8 @@ static BigQueryArrowResultSet of(
-1,
ownedTask,
bigQuery,
job);
job,
statement != null && statement.isEnableTimestampPicos());
}

BigQueryArrowResultSet() throws SQLException {
Expand All @@ -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,
Expand All @@ -189,7 +199,8 @@ static BigQueryArrowResultSet getNestedResultSet(
toIndexExclusive,
null,
null,
null);
null,
enableTimestampPicos);
}

private class ArrowDeserializer implements AutoCloseable {
Expand Down Expand Up @@ -345,6 +356,9 @@ private Object getObjectInternal(int columnIndex) throws SQLException {
value = LocalDate.ofEpochDay(((Integer) value).longValue());
}
}
if (value instanceof Text) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why special case for Text here? Is it related to picoseconds in some way?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, so in Storage Read API:

  • standard timestamps (<= 6 digits) are transmitted in Arrow as 64-bit integers (TimeStampMicroVector).
  • picosecond timestamps are transmitted as as VarCharVectors

So, calling VarCharVector.getObject() returns an org.apache.arrow.vector.util.Text and we convert that to String

value = value.toString();
}
setWasNull(value);
return value;
}
Expand All @@ -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);
}

Expand All @@ -388,61 +406,87 @@ public Object getObject(int columnIndex) throws SQLException {
|| elementTypeName == StandardSQLTypeName.BIGNUMERIC) {
JsonStringArrayList<BigDecimal> 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<String> 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);
}
}
Comment thread
keshavdandeva marked this conversation as resolved.
Comment thread
keshavdandeva marked this conversation as resolved.

Expand All @@ -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)
Expand All @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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;
}
Expand Down Expand Up @@ -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);
}
}
Loading
Loading