Skip to content

Commit

Permalink
move StructObjectWrapper from converters to SFArrowResultSet
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-mkubik committed Dec 9, 2024
1 parent c9bf0e7 commit 0c3b19a
Show file tree
Hide file tree
Showing 7 changed files with 32 additions and 16 deletions.
23 changes: 16 additions & 7 deletions src/main/java/net/snowflake/client/core/SFArrowResultSet.java
Original file line number Diff line number Diff line change
Expand Up @@ -574,19 +574,23 @@ public Object getObject(int columnIndex) throws SFException {
converter.setTreatNTZAsUTC(treatNTZAsUTC);
converter.setUseSessionTimezone(useSessionTimezone);
converter.setSessionTimeZone(sessionTimeZone);
Object obj = converter.toObject(index);
boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
if (isVarcharConvertedStruct(type, isStructuredType, converter)) {
if (obj != null) {
Object obj = converter.toObject(index);
if (obj == null) {
return null;
}
if (isStructuredType) {
if (type == Types.STRUCT && converter instanceof VarCharConverter) {
return new StructObjectWrapper((String) obj, createJsonSqlInput(columnIndex, obj));
}
return new StructObjectWrapper(converter.toString(index), obj);
}
return obj;
}

private boolean isVarcharConvertedStruct(
int type, boolean isStructuredType, ArrowVectorConverter converter) {
return type == Types.STRUCT && isStructuredType && converter instanceof VarCharConverter;
@Override
public <T> Object getObject(int columnIndex, Class<T> type) throws SFException {
return getObject(columnIndex);
}

private Object createJsonSqlInput(int columnIndex, Object obj) throws SFException {
Expand Down Expand Up @@ -620,7 +624,7 @@ public Array getArray(int columnIndex) throws SFException {
if (converter instanceof VarCharConverter) {
return getJsonArray((String) obj, columnIndex);
} else if (converter instanceof ArrayConverter || converter instanceof VectorTypeConverter) {
StructObjectWrapper structObjectWrapper = (StructObjectWrapper) obj;
StructObjectWrapper structObjectWrapper = new StructObjectWrapper(converter.toString(), obj);
return getArrowArray(
structObjectWrapper.getJsonString(),
(List<Object>) structObjectWrapper.getObject(),
Expand All @@ -630,6 +634,11 @@ public Array getArray(int columnIndex) throws SFException {
}
}

public <T> Array getArray(int columnIndex, Class<T> type) throws SFException {
// TODO: don't calculate toString when not needed
return getArray(columnIndex);
}

private SfSqlArray getArrowArray(String text, List<Object> elements, int columnIndex)
throws SFException {
try {
Expand Down
6 changes: 6 additions & 0 deletions src/main/java/net/snowflake/client/core/SFBaseResultSet.java
Original file line number Diff line number Diff line change
Expand Up @@ -113,10 +113,16 @@ public abstract class SFBaseResultSet {

public abstract Object getObject(int columnIndex) throws SFException;

public abstract <T> Object getObject(int columnIndex, Class<T> object) throws SFException;

public Array getArray(int columnIndex) throws SFException {
throw new UnsupportedOperationException();
}

public <T> Array getArray(int columnIndex, Class<T> type) throws SFException {
throw new UnsupportedOperationException();
}

public abstract BigDecimal getBigDecimal(int columnIndex) throws SFException;

public abstract BigDecimal getBigDecimal(int columnIndex, int scale) throws SFException;
Expand Down
6 changes: 6 additions & 0 deletions src/main/java/net/snowflake/client/core/SFJsonResultSet.java
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,12 @@ public Object getObject(int columnIndex) throws SFException {
}
}

@Override
public <T> Object getObject(int columnIndex, Class<T> object) throws SFException {
// TODO: don't calculate string representation for structured types when not needed
return getObject(columnIndex);
}

/**
* Sometimes large BIGINTS overflow the java Long type. In these cases, return a BigDecimal type
* instead.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ public ArrayConverter(ListVector valueVector, int vectorIndex, DataConversionCon

@Override
public Object toObject(int index) throws SFException {
return isNull(index) ? null : new StructObjectWrapper(toString(index), vector.getObject(index));
return isNull(index) ? null : vector.getObject(index);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,12 +36,10 @@ public Object toObject(int index) throws SFException {

List<JsonStringHashMap<String, Object>> entriesList =
(List<JsonStringHashMap<String, Object>>) vector.getObject(index);
Map<String, Object> map =
entriesList.stream()
return entriesList.stream()
.collect(
Collectors.toMap(
entry -> entry.get("key").toString(), entry -> entry.get("value")));
return new StructObjectWrapper(toString(index), map);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,7 @@ public StructConverter(StructVector vector, int columnIndex, DataConversionConte

@Override
public Object toObject(int index) throws SFException {
return isNull(index)
? null
: new StructObjectWrapper(toString(index), structVector.getObject(index));
return isNull(index) ? null : structVector.getObject(index);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,7 @@ public Object toObject(int index) throws SFException {
if (isNull(index)) {
return null;
}
Object object = vector.getObject(index);
return new StructObjectWrapper(object.toString(), object);
return vector.getObject(index);
}

@Override
Expand Down

0 comments on commit 0c3b19a

Please sign in to comment.