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: Fix for structured types as strings performance problems #1991

Merged
merged 9 commits into from
Dec 10, 2024
41 changes: 32 additions & 9 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 @@ -564,6 +565,16 @@ public Timestamp getTimestamp(int columnIndex, TimeZone tz) throws SFException {

@Override
public Object getObject(int columnIndex) throws SFException {
return getObjectRepresentation(columnIndex, true);
}

@SnowflakeJdbcInternalApi
@Override
public Object getObject(int columnIndex, Class<?> type) throws SFException {
return getObjectRepresentation(columnIndex, false);
sfc-gh-dprzybysz marked this conversation as resolved.
Show resolved Hide resolved
}

private Object getObjectRepresentation(int columnIndex, boolean withString) throws SFException {
int type = resultSetMetaData.getColumnType(columnIndex);
if (type == SnowflakeUtil.EXTRA_TYPES_VECTOR) {
return getString(columnIndex);
Expand All @@ -575,13 +586,28 @@ public Object getObject(int columnIndex) throws SFException {
converter.setUseSessionTimezone(useSessionTimezone);
converter.setSessionTimeZone(sessionTimeZone);
Object obj = converter.toObject(index);
if (obj == null) {
return null;
}
String jsonString = withString ? converter.toString(index) : null;
boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
if (isVarcharConvertedStruct(type, isStructuredType, converter)) {
if (obj != null) {
return new StructObjectWrapper((String) obj, createJsonSqlInput(columnIndex, obj));
}
return new StructObjectWrapper(jsonString, createJsonSqlInput(columnIndex, obj));
sfc-gh-dprzybysz marked this conversation as resolved.
Show resolved Hide resolved
} else if (converter instanceof StructConverter) {
return new StructObjectWrapper(
jsonString, createArrowSqlInput(columnIndex, (Map<String, Object>) obj));
} else {
return new StructObjectWrapper(jsonString, obj);
}
}

private SQLInput createArrowSqlInput(int columnIndex, Map<String, Object> input)
throws SFException {
if (input == null) {
return null;
}
return obj;
return new ArrowSqlInput(
input, session, converters, resultSetMetaData.getColumnFields(columnIndex));
}

private boolean isVarcharConvertedStruct(
Expand Down Expand Up @@ -620,11 +646,8 @@ 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);
String jsonString = converter.toString(index);
return getArrowArray(jsonString, (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,9 @@ public abstract class SFBaseResultSet {

public abstract Object getObject(int columnIndex) throws SFException;

@SnowflakeJdbcInternalApi
public abstract Object getObject(int columnIndex, Class<?> type) throws SFException;

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 {
}
}

@SnowflakeJdbcInternalApi
@Override
public Object getObject(int columnIndex, Class<?> type) throws SFException {
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 @@ -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 @@ -33,15 +32,11 @@ public Object toObject(int index) throws SFException {
if (isNull(index)) {
return null;
}

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 @@ -28,7 +28,7 @@ public Object toObject(int index) throws SFException {
return null;
}
Object object = vector.getObject(index);
return new StructObjectWrapper(object.toString(), object);
return object;
sfc-gh-dprzybysz marked this conversation as resolved.
Show resolved Hide resolved
}

@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 (SQLInput) structObjectWrapper.getObject();
});
if (sqlInput == null) {
return null;
Expand Down Expand Up @@ -1639,7 +1642,7 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
StructObjectWrapper structObjectWrapper =
(StructObjectWrapper)
SnowflakeUtil.mapSFExceptionToSQLException(
() -> sfBaseResultSet.getObject(columnIndex));
() -> sfBaseResultSet.getObject(columnIndex, type));
if (structObjectWrapper == null) {
return null;
}
Expand Down
25 changes: 13 additions & 12 deletions src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,6 @@
import java.util.List;
import java.util.Map;
import java.util.TimeZone;
import net.snowflake.client.core.ArrowSqlInput;
import net.snowflake.client.core.JsonSqlInput;
import net.snowflake.client.core.QueryStatus;
import net.snowflake.client.core.SFBaseResultSet;
import net.snowflake.client.core.SFException;
Expand Down Expand Up @@ -274,18 +272,21 @@ public Object getObject(int columnIndex) throws SQLException {
SnowflakeUtil.mapSFExceptionToSQLException(() -> sfBaseResultSet.getObject(columnIndex));
if (object == null) {
return null;
} else if (object instanceof JsonSqlInput) {
return ((JsonSqlInput) object).getText();
} else if (object instanceof StructObjectWrapper) {
return ((StructObjectWrapper) object).getJsonString();
} else if (object instanceof SfSqlArray) {
}
if (object instanceof SfSqlArray) {
return ((SfSqlArray) object).getText();
} else if (object instanceof ArrowSqlInput) {
throw new SQLException(
"Arrow native struct couldn't be converted to String. To map to SqlData the method getObject(int columnIndex, Class type) should be used");
} else {
return object;
}
if (object instanceof StructObjectWrapper) {
StructObjectWrapper structObjectWrapper = (StructObjectWrapper) object;
if (resultSetMetaData.isStructuredTypeColumn(columnIndex)
&& structObjectWrapper.getJsonString() != null) {
return structObjectWrapper.getJsonString();
}
if (structObjectWrapper.getObject() != null) {
return structObjectWrapper.getObject();
}
}
return object;
}

public Array getArray(int columnIndex) throws SQLException {
Expand Down
Loading