Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,18 @@
import com.google.cloud.Tuple;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.Schema;
import com.google.cloud.bigquery.StandardSQLTypeName;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.time.LocalDate;
import org.apache.arrow.vector.util.JsonStringArrayList;
import org.apache.arrow.vector.util.JsonStringHashMap;

/**
* An implementation of {@link BigQueryBaseArray} used to represent Array values from Arrow data.
*/
class BigQueryArrowArray extends BigQueryBaseArray {
private static final BigQueryTypeCoercer BIGQUERY_TYPE_COERCER =
BigQueryTypeCoercionUtility.INSTANCE;

private JsonStringArrayList<?> values;

public BigQueryArrowArray(Field schema, JsonStringArrayList<?> values) {
Expand All @@ -43,7 +44,7 @@ public BigQueryArrowArray(
}

@Override
public Object getArray() {
public Object getArray() throws SQLException {
LOG.finestTrace("getArray");
ensureValid();
if (values == null) {
Expand All @@ -53,7 +54,7 @@ public Object getArray() {
}

@Override
public Object getArray(long index, int count) {
public Object getArray(long index, int count) throws SQLException {
LOG.finestTrace("getArray");
ensureValid();
if (values == null) {
Expand Down Expand Up @@ -98,12 +99,16 @@ public void free() {
}

@Override
Object getCoercedValue(int index) {
Object getCoercedValue(int index) throws SQLException {
LOG.finestTrace("getCoercedValue");
Object value = this.values.get(index);
if (value instanceof Integer
&& schema.getType().getStandardType() == StandardSQLTypeName.DATE) {
value = LocalDate.ofEpochDay(((Integer) value).longValue());
}
return this.arrayOfStruct
? new BigQueryArrowStruct(
schema.getSubFields(), (JsonStringHashMap<?, ?>) value, this.LOG.getArrowStructLogger())
: BIGQUERY_TYPE_COERCER.coerceTo(getTargetClass(), value, this.LOG);
: BigQueryTypeRegistry.convert(value, getTargetClass());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,18 @@
import io.opentelemetry.context.Scope;
import java.io.IOException;
import java.math.BigDecimal;
import java.sql.Date;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Future;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.VectorSchemaRoot;
Expand Down Expand Up @@ -340,6 +341,9 @@ private Object getObjectInternal(int columnIndex) throws SQLException {
FieldVector currentColumn = this.vectorSchemaRoot.getVector(columnIndex - 1);
// get the current row
value = currentColumn.getObject(this.currentBatchRowIndex);
if (value instanceof Integer && currentColumn instanceof DateDayVector) {
value = LocalDate.ofEpochDay(((Integer) value).longValue());
}
}
setWasNull(value);
return value;
Expand All @@ -357,7 +361,7 @@ public Object getObject(int columnIndex) throws SQLException {
}

if (this.isNested && columnIndex == 1) {
return this.bigQueryTypeCoercer.coerceTo(Integer.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Integer.class);
}

if (this.isNested && columnIndex == 2) {
Expand All @@ -368,10 +372,11 @@ public Object getObject(int columnIndex) throws SQLException {
(JsonStringHashMap<?, ?>) value,
this.LOG.getArrowStructLogger());
}
Class<?> targetClass =
BigQueryJdbcTypeMappings.standardSQLToJavaTypeMapping.get(
arrayField.getType().getStandardType());
return this.bigQueryTypeCoercer.coerceTo(targetClass, value, this.LOG);
if (value instanceof Integer
&& arrayField.getType().getStandardType() == StandardSQLTypeName.DATE) {
value = LocalDate.ofEpochDay(((Integer) value).longValue());
}
return BigQueryTypeRegistry.convert(value, arrayField.getType().getStandardType(), null);
}

int fieldIndex = this.isNested ? 0 : columnIndex - 1;
Expand Down Expand Up @@ -437,10 +442,7 @@ public Object getObject(int columnIndex) throws SQLException {
// Strip trailing zeros to match JSON API and CLI output
return ((BigDecimal) value).stripTrailingZeros();
}
Class<?> targetClass =
BigQueryJdbcTypeMappings.standardSQLToJavaTypeMapping.get(
fieldSchema.getType().getStandardType());
return this.bigQueryTypeCoercer.coerceTo(targetClass, value, this.LOG);
return BigQueryTypeRegistry.convert(value, fieldSchema.getType().getStandardType(), null);
}
}

Expand All @@ -460,24 +462,23 @@ private StandardSQLTypeName getElementTypeFromValue(Object element) {
return StandardSQLTypeName.STRING;
}

private String formatRangeElement(Object element, StandardSQLTypeName elementType) {
private String formatRangeElement(Object element, StandardSQLTypeName elementType)
throws SQLException {
if (element == null) {
return "UNBOUNDED";
}
switch (elementType) {
case DATE:
// Arrow gives DATE as an Integer (days since epoch)
Date date = this.bigQueryTypeCoercer.coerceTo(Date.class, (Integer) element, this.LOG);
return date.toString();
return LocalDate.ofEpochDay(((Integer) element).longValue()).toString();
case DATETIME:
// Arrow gives DATETIME as a LocalDateTime
Timestamp dtTs =
this.bigQueryTypeCoercer.coerceTo(Timestamp.class, (LocalDateTime) element, this.LOG);
return this.bigQueryTypeCoercer.coerceTo(String.class, dtTs, this.LOG);
Timestamp dtTs = BigQueryTypeRegistry.convert((LocalDateTime) element, Timestamp.class);
return BigQueryTypeRegistry.convert(dtTs, String.class);
case TIMESTAMP:
// Arrow gives TIMESTAMP as a Long (microseconds since epoch)
Timestamp ts = this.bigQueryTypeCoercer.coerceTo(Timestamp.class, (Long) element, this.LOG);
return this.bigQueryTypeCoercer.coerceTo(String.class, ts, this.LOG);
Timestamp ts = BigQueryTypeRegistry.convert((Long) element, Timestamp.class);
return BigQueryTypeRegistry.convert(ts, String.class);
default:
// Fallback for any other unexpected type
return element.toString();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,10 @@

import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.FieldList;
import com.google.cloud.bigquery.StandardSQLTypeName;
import java.lang.reflect.Array;
import java.sql.SQLException;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.List;
import org.apache.arrow.vector.util.JsonStringArrayList;
Expand All @@ -30,8 +33,6 @@
* An implementation of {@link BigQueryBaseStruct} used to represent Struct values from Arrow data.
*/
class BigQueryArrowStruct extends BigQueryBaseStruct {
private static final BigQueryTypeCoercer BIGQUERY_TYPE_COERCER =
BigQueryTypeCoercionUtility.INSTANCE;

private final FieldList schema;

Expand All @@ -54,7 +55,7 @@ FieldList getSchema() {
}

@Override
public Object[] getAttributes() {
public Object[] getAttributes() throws SQLException {
LOG.finestTrace("getAttributes");
int size = this.schema.size();
Object[] attributes = (Object[]) Array.newInstance(Object.class, size);
Expand All @@ -73,7 +74,7 @@ public Object[] getAttributes() {
return attributes;
}

private Object getValue(Field currentSchema, Object currentValue) {
private Object getValue(Field currentSchema, Object currentValue) throws SQLException {
LOG.finestTrace("getValue");
if (isArray(currentSchema)) {
return new BigQueryArrowArray(
Expand All @@ -84,10 +85,12 @@ private Object getValue(Field currentSchema, Object currentValue) {
(JsonStringHashMap<?, ?>) currentValue,
this.LOG.getArrowStructLogger());
} else {
Class<?> targetClass =
BigQueryJdbcTypeMappings.standardSQLToJavaTypeMapping.get(
currentSchema.getType().getStandardType());
return BIGQUERY_TYPE_COERCER.coerceTo(targetClass, currentValue, this.LOG);
if (currentValue instanceof Integer
&& currentSchema.getType().getStandardType() == StandardSQLTypeName.DATE) {
currentValue = LocalDate.ofEpochDay(((Integer) currentValue).longValue());
}
return BigQueryTypeRegistry.convert(
currentValue, currentSchema.getType().getStandardType(), null);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -65,8 +65,7 @@ public final String getBaseTypeName() {
public final int getBaseType() {
LOG.finestTrace("getBaseType");
ensureValid();
return BigQueryJdbcTypeMappings.standardSQLToJavaSqlTypesMapping.get(
schema.getType().getStandardType());
return BigQueryTypeRegistry.toJdbcType(schema.getType().getStandardType());
}

@Override
Expand All @@ -91,7 +90,7 @@ public final ResultSet getResultSet(long index, int count, Map<String, Class<?>>
throw new BigQueryJdbcSqlFeatureNotSupportedException(CUSTOMER_TYPE_MAPPING_NOT_SUPPORTED);
}

protected Object getArrayInternal(int fromIndex, int toIndexExclusive) {
protected Object getArrayInternal(int fromIndex, int toIndexExclusive) throws SQLException {
LOG.finestTrace("getArrayInternal");
Class<?> targetClass = getTargetClass();
int size = toIndexExclusive - fromIndex;
Expand Down Expand Up @@ -145,11 +144,10 @@ protected Class<?> getTargetClass() {
LOG.finestTrace("getTargetClass");
return this.arrayOfStruct
? Struct.class
: BigQueryJdbcTypeMappings.standardSQLToJavaTypeMapping.get(
this.schema.getType().getStandardType());
: BigQueryTypeRegistry.toJavaClass(this.schema.getType().getStandardType());
}

abstract Object getCoercedValue(int index);
abstract Object getCoercedValue(int index) throws SQLException;

static boolean isArray(Field currentSchema) {
return currentSchema.getMode() == REPEATED;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet
private Job job;
private SQLWarning warnings;
private boolean warningsLoaded = false;
protected final BigQueryTypeCoercer bigQueryTypeCoercer = BigQueryTypeCoercionUtility.INSTANCE;

protected final SpanContext originalSpanContext;

protected BigQueryBaseResultSet(
Expand Down Expand Up @@ -297,7 +297,7 @@ public <T> T getObject(int columnIndex, Class<T> type) throws SQLException {
if (value == null) {
return null;
}
return this.bigQueryTypeCoercer.coerceTo(type, value, this.LOG);
return BigQueryTypeRegistry.convert(value, type);
} catch (RuntimeException e) {
throw createCoercionException(columnIndex, type, e);
}
Comment thread
Neenu1995 marked this conversation as resolved.
Expand All @@ -323,7 +323,7 @@ public String getString(int columnIndex) throws SQLException {
LOG.finestTrace("getString");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(String.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, String.class);
} catch (BigQueryJdbcCoercionNotFoundException e) {
throw createCoercionException(columnIndex, String.class, e);
}
Expand All @@ -342,7 +342,7 @@ public boolean getBoolean(int columnIndex) throws SQLException {

try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(Boolean.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Boolean.class);
} catch (BigQueryJdbcCoercionNotFoundException e) {
throw createCoercionException(columnIndex, Boolean.class, e);
}
Expand All @@ -353,7 +353,7 @@ public byte getByte(int columnIndex) throws SQLException {
LOG.finestTrace("getByte");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(Byte.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Byte.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, Byte.class, e);
}
Expand All @@ -364,7 +364,7 @@ public short getShort(int columnIndex) throws SQLException {
LOG.finestTrace("getShort");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(Short.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Short.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, Short.class, e);
}
Expand All @@ -375,7 +375,7 @@ public int getInt(int columnIndex) throws SQLException {
LOG.finestTrace("getInt");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(Integer.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Integer.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, Integer.class, e);
}
Expand All @@ -386,7 +386,7 @@ public long getLong(int columnIndex) throws SQLException {
LOG.finestTrace("getLong");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(Long.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Long.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, Long.class, e);
}
Expand All @@ -397,7 +397,7 @@ public float getFloat(int columnIndex) throws SQLException {
LOG.finestTrace("getFloat");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(Float.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Float.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, Float.class, e);
}
Expand All @@ -408,7 +408,7 @@ public double getDouble(int columnIndex) throws SQLException {
LOG.finestTrace("getDouble");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(Double.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, Double.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, Double.class, e);
}
Expand All @@ -421,7 +421,7 @@ public BigDecimal getBigDecimal(int columnIndex, int scale) throws SQLException
LOG.finestTrace("getBigDecimal");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(BigDecimal.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, BigDecimal.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, BigDecimal.class, e);
}
Expand All @@ -432,7 +432,7 @@ public byte[] getBytes(int columnIndex) throws SQLException {
LOG.finestTrace("getBytes");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(byte[].class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, byte[].class);
} catch (BigQueryJdbcCoercionNotFoundException e) {
throw createCoercionException(columnIndex, byte[].class, e);
}
Expand All @@ -443,7 +443,7 @@ public Date getDate(int columnIndex) throws SQLException {
LOG.finestTrace("getDate");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(java.sql.Date.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, java.sql.Date.class);
} catch (BigQueryJdbcCoercionNotFoundException e) {
throw createCoercionException(columnIndex, java.sql.Date.class, e);
}
Expand All @@ -458,7 +458,7 @@ public Time getTime(int columnIndex) throws SQLException {
}
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(java.sql.Time.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, java.sql.Time.class);
} catch (BigQueryJdbcCoercionNotFoundException e) {
throw createCoercionException(columnIndex, java.sql.Time.class, e);
}
Expand All @@ -473,7 +473,7 @@ public Timestamp getTimestamp(int columnIndex) throws SQLException {
}
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(java.sql.Timestamp.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, java.sql.Timestamp.class);
} catch (BigQueryJdbcCoercionNotFoundException e) {
throw createCoercionException(columnIndex, java.sql.Timestamp.class, e);
}
Expand All @@ -484,7 +484,7 @@ public BigDecimal getBigDecimal(int columnIndex) throws SQLException {
LOG.finestTrace("getBigDecimal");
try {
Object value = getObject(columnIndex);
return this.bigQueryTypeCoercer.coerceTo(BigDecimal.class, value, this.LOG);
return BigQueryTypeRegistry.convert(value, BigDecimal.class);
} catch (BigQueryJdbcCoercionNotFoundException | BigQueryJdbcCoercionException e) {
throw createCoercionException(columnIndex, BigDecimal.class, e);
}
Expand Down
Loading
Loading