Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1848734 Prevent string creation on getObject call when not necessary #1989

67 changes: 54 additions & 13 deletions src/main/java/net/snowflake/client/core/SFArrowResultSet.java
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import java.util.stream.Stream;
import net.snowflake.client.core.arrow.ArrayConverter;
import net.snowflake.client.core.arrow.ArrowVectorConverter;
import net.snowflake.client.core.arrow.StructConverter;
import net.snowflake.client.core.arrow.StructObjectWrapper;
import net.snowflake.client.core.arrow.VarCharConverter;
import net.snowflake.client.core.arrow.VectorTypeConverter;
Expand Down Expand Up @@ -574,22 +575,66 @@ 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 (type == Types.STRUCT && isStructuredType) {
if (converter instanceof VarCharConverter) {
return new StructObjectWrapper((String) obj, createJsonSqlInput(columnIndex, obj));
} else if (converter instanceof StructConverter) {
return new StructObjectWrapper(
converter.toString(index),
createArrowSqlInput(columnIndex, (Map<String, Object>) obj),
obj);
}
}
return new StructObjectWrapper(converter.toString(index), null, obj);
}

@Override
public <T> Object getObject(int columnIndex, Class<T> type) throws SFException {
if (type == String.class) {
return getObject(columnIndex);
}

int columnType = resultSetMetaData.getColumnType(columnIndex);
if (columnType == SnowflakeUtil.EXTRA_TYPES_VECTOR) {
return getString(columnIndex);
}
ArrowVectorConverter converter = currentChunkIterator.getCurrentConverter(columnIndex - 1);
int index = currentChunkIterator.getCurrentRowInRecordBatch();
wasNull = converter.isNull(index);
converter.setTreatNTZAsUTC(treatNTZAsUTC);
converter.setUseSessionTimezone(useSessionTimezone);
converter.setSessionTimeZone(sessionTimeZone);
Object obj = converter.toObject(index);
boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
if (obj == null) {
return null;
}
Copy link
Contributor

Choose a reason for hiding this comment

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

We have a lot of duplicated code lines between getObject(index) and getObject(index, type)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

moved some part of it to external method

if (columnType == Types.STRUCT && isStructuredType) {
if (converter instanceof VarCharConverter) {
return new StructObjectWrapper(null, createJsonSqlInput(columnIndex, obj));
} else if (converter instanceof StructConverter) {
return new StructObjectWrapper(
null, createArrowSqlInput(columnIndex, (Map<String, Object>) obj));
}
}
return obj;
return new StructObjectWrapper(null, null, obj);
}

private boolean isVarcharConvertedStruct(
int type, boolean isStructuredType, ArrowVectorConverter converter) {
return type == Types.STRUCT && isStructuredType && converter instanceof VarCharConverter;
private SQLInput createArrowSqlInput(int columnIndex, Map<String, Object> input)
throws SFException {
if (input == null) {
return null;
}
return new ArrowSqlInput(
input, session, converters, resultSetMetaData.getColumnFields(columnIndex));
}

private Object createJsonSqlInput(int columnIndex, Object obj) throws SFException {
private SQLInput createJsonSqlInput(int columnIndex, Object obj) throws SFException {
try {
if (obj == null) {
return null;
Expand Down Expand Up @@ -620,11 +665,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;
return getArrowArray(
structObjectWrapper.getJsonString(),
(List<Object>) structObjectWrapper.getObject(),
columnIndex);
return getArrowArray(converter.toString(), (List<Object>) obj, columnIndex);
} else {
throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,8 @@ public abstract class SFBaseResultSet {

public abstract Object getObject(int columnIndex) throws SFException;

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

Choose a reason for hiding this comment

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

this API looks strange - the result type should be T, not Object

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I agree, in reality this getObject mostly returns StructObjectWrapper which possibly could be generic. I may rename this method since it's used in higher level getObject anyway to better indicate it's purpose


public Array getArray(int columnIndex) throws SFException {
throw new UnsupportedOperationException();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,11 @@ public Object getObject(int columnIndex) throws SFException {
}
}

@Override
public <T> Object getObject(int columnIndex, Class<T> object) throws SFException {
return getObject(columnIndex);
}

/**
* Sometimes large BIGINTS overflow the java Long type. In these cases, return a BigDecimal type
* instead.
Expand Down Expand Up @@ -292,7 +297,7 @@ public Converters getConverters() {
return converters;
}

private Object getSqlInput(String input, int columnIndex) throws SFException {
private SQLInput getSqlInput(String input, int columnIndex) throws SFException {
try {
JsonNode jsonNode = OBJECT_MAPPER.readTree(input);
return new JsonSqlInput(
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 @@ -2,7 +2,6 @@

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import net.snowflake.client.core.DataConversionContext;
import net.snowflake.client.core.SFException;
Expand Down Expand Up @@ -36,12 +35,9 @@ 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()
.collect(
Collectors.toMap(
entry -> entry.get("key").toString(), entry -> entry.get("value")));
return new StructObjectWrapper(toString(index), map);
return entriesList.stream()
.collect(
Collectors.toMap(entry -> entry.get("key").toString(), entry -> entry.get("value")));
}

@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 @@ -3,22 +3,35 @@
*/
package net.snowflake.client.core.arrow;

import java.sql.SQLInput;
import net.snowflake.client.core.SnowflakeJdbcInternalApi;

@SnowflakeJdbcInternalApi
public class StructObjectWrapper {
private final String jsonString;
private final SQLInput sqlInput;
private final Object object;

public StructObjectWrapper(String jsonString, Object object) {
public StructObjectWrapper(String jsonString, SQLInput sqlInput) {
this.jsonString = jsonString;
this.sqlInput = sqlInput;
this.object = null;
}

public StructObjectWrapper(String jsonString, SQLInput sqlInput, Object object) {
this.jsonString = jsonString;
this.sqlInput = sqlInput;
this.object = object;
}

public String getJsonString() {
return jsonString;
}

public SQLInput getSqlInput() {
return sqlInput;
}

public Object getObject() {
return object;
}
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
Original file line number Diff line number Diff line change
Expand Up @@ -1399,7 +1399,10 @@ public <T> T getObject(int columnIndex, Class<T> type) throws SQLException {
() -> {
StructObjectWrapper structObjectWrapper =
(StructObjectWrapper) sfBaseResultSet.getObject(columnIndex);
return (SQLInput) createJsonSqlInput(columnIndex, structObjectWrapper);
if (structObjectWrapper == null) {
return null;
}
return structObjectWrapper.getSqlInput();
});
if (sqlInput == null) {
return null;
Expand Down Expand Up @@ -1637,23 +1640,25 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
int scale = valueFieldMetadata.getScale();
TimeZone tz = sfBaseResultSet.getSessionTimeZone();
StructObjectWrapper structObjectWrapper =
(StructObjectWrapper)
SnowflakeUtil.mapSFExceptionToSQLException(
() -> sfBaseResultSet.getObject(columnIndex));
SnowflakeUtil.mapSFExceptionToSQLException(
() -> (StructObjectWrapper) sfBaseResultSet.getObject(columnIndex, type));
if (structObjectWrapper == null) {
return null;
}

Map<String, Object> map =
mapSFExceptionToSQLException(
() -> prepareMapWithValues(structObjectWrapper.getObject(), type));
mapSFExceptionToSQLException(() -> prepareMapWithValues(structObjectWrapper, type));
Map<String, T> resultMap = new HashMap<>();
for (Map.Entry<String, Object> entry : map.entrySet()) {
if (SQLData.class.isAssignableFrom(type)) {
SQLData instance = (SQLData) SQLDataCreationHelper.create(type);
SQLInput parentSqlInput = structObjectWrapper.getSqlInput();
SQLInput sqlInput =
sfBaseResultSet.createSqlInputForColumn(
entry.getValue(),
structObjectWrapper.getObject().getClass(),
parentSqlInput == null
? structObjectWrapper.getObject().getClass()
: parentSqlInput.getClass(),
columnIndex,
session,
valueFieldMetadata.getFields());
Expand Down Expand Up @@ -1812,11 +1817,12 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
return iface.isInstance(this);
}

private <T> Map<String, Object> prepareMapWithValues(Object object, Class<T> type)
throws SFException {
if (object instanceof JsonSqlInput) {
private <T> Map<String, Object> prepareMapWithValues(
Copy link
Collaborator

Choose a reason for hiding this comment

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

where are we using type?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It was already there, in line 1830 to check if it's assignable to SQLData

StructObjectWrapper structObjectWrapper, Class<T> type) throws SFException {
SQLInput sqlInput = structObjectWrapper.getSqlInput();
if (sqlInput instanceof JsonSqlInput) {
Map map = new HashMap<>();
JsonNode jsonNode = ((JsonSqlInput) object).getInput();
JsonNode jsonNode = ((JsonSqlInput) sqlInput).getInput();
for (Iterator<String> it = jsonNode.fieldNames(); it.hasNext(); ) {
String name = it.next();
map.put(
Expand All @@ -1826,14 +1832,14 @@ private <T> Map<String, Object> prepareMapWithValues(Object object, Class<T> typ
: SnowflakeUtil.getJsonNodeStringValue(jsonNode.get(name)));
}
return map;
} else if (object instanceof Map) {
return (Map<String, Object>) object;
} else if (structObjectWrapper.getObject() != null) {
return (Map<String, Object>) structObjectWrapper.getObject();
} else {
throw new SFException(ErrorCode.INVALID_STRUCT_DATA, "Object couldn't be converted to map");
}
}

private Object createJsonSqlInput(int columnIndex, StructObjectWrapper obj) throws SFException {
private SQLInput createJsonSqlInput(int columnIndex, StructObjectWrapper obj) throws SFException {
try {
if (obj == null) {
return null;
Expand Down
Loading