Skip to content

Commit e793e61

Browse files
sfc-gh-pmotackisfc-gh-mkubik
sfc-gh-pmotacki
and
sfc-gh-mkubik
authored
SNOW-1848734: Fix for structured types as strings performance problems (#1991)
Co-authored-by: Mikołaj Kubik <[email protected]>
1 parent 3525adf commit e793e61

File tree

10 files changed

+66
-42
lines changed

10 files changed

+66
-42
lines changed

‎src/main/java/net/snowflake/client/core/SFArrowResultSet.java

Lines changed: 32 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.util.stream.Stream;
2727
import net.snowflake.client.core.arrow.ArrayConverter;
2828
import net.snowflake.client.core.arrow.ArrowVectorConverter;
29+
import net.snowflake.client.core.arrow.StructConverter;
2930
import net.snowflake.client.core.arrow.StructObjectWrapper;
3031
import net.snowflake.client.core.arrow.VarCharConverter;
3132
import net.snowflake.client.core.arrow.VectorTypeConverter;
@@ -564,6 +565,16 @@ public Timestamp getTimestamp(int columnIndex, TimeZone tz) throws SFException {
564565

565566
@Override
566567
public Object getObject(int columnIndex) throws SFException {
568+
return getObjectRepresentation(columnIndex, true);
569+
}
570+
571+
@SnowflakeJdbcInternalApi
572+
@Override
573+
public Object getObjectWithoutString(int columnIndex) throws SFException {
574+
return getObjectRepresentation(columnIndex, false);
575+
}
576+
577+
private Object getObjectRepresentation(int columnIndex, boolean withString) throws SFException {
567578
int type = resultSetMetaData.getColumnType(columnIndex);
568579
if (type == SnowflakeUtil.EXTRA_TYPES_VECTOR) {
569580
return getString(columnIndex);
@@ -575,13 +586,28 @@ public Object getObject(int columnIndex) throws SFException {
575586
converter.setUseSessionTimezone(useSessionTimezone);
576587
converter.setSessionTimeZone(sessionTimeZone);
577588
Object obj = converter.toObject(index);
589+
if (obj == null) {
590+
return null;
591+
}
592+
String jsonString = withString ? converter.toString(index) : null;
578593
boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
579594
if (isVarcharConvertedStruct(type, isStructuredType, converter)) {
580-
if (obj != null) {
581-
return new StructObjectWrapper((String) obj, createJsonSqlInput(columnIndex, obj));
582-
}
595+
return new StructObjectWrapper(jsonString, createJsonSqlInput(columnIndex, obj));
596+
} else if (converter instanceof StructConverter) {
597+
return new StructObjectWrapper(
598+
jsonString, createArrowSqlInput(columnIndex, (Map<String, Object>) obj));
599+
} else {
600+
return new StructObjectWrapper(jsonString, obj);
601+
}
602+
}
603+
604+
private SQLInput createArrowSqlInput(int columnIndex, Map<String, Object> input)
605+
throws SFException {
606+
if (input == null) {
607+
return null;
583608
}
584-
return obj;
609+
return new ArrowSqlInput(
610+
input, session, converters, resultSetMetaData.getColumnFields(columnIndex));
585611
}
586612

587613
private boolean isVarcharConvertedStruct(
@@ -620,11 +646,8 @@ public Array getArray(int columnIndex) throws SFException {
620646
if (converter instanceof VarCharConverter) {
621647
return getJsonArray((String) obj, columnIndex);
622648
} else if (converter instanceof ArrayConverter || converter instanceof VectorTypeConverter) {
623-
StructObjectWrapper structObjectWrapper = (StructObjectWrapper) obj;
624-
return getArrowArray(
625-
structObjectWrapper.getJsonString(),
626-
(List<Object>) structObjectWrapper.getObject(),
627-
columnIndex);
649+
String jsonString = converter.toString(index);
650+
return getArrowArray(jsonString, (List<Object>) obj, columnIndex);
628651
} else {
629652
throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
630653
}

‎src/main/java/net/snowflake/client/core/SFBaseResultSet.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,9 @@ public abstract class SFBaseResultSet {
113113

114114
public abstract Object getObject(int columnIndex) throws SFException;
115115

116+
@SnowflakeJdbcInternalApi
117+
public abstract Object getObjectWithoutString(int columnIndex) throws SFException;
118+
116119
public Array getArray(int columnIndex) throws SFException {
117120
throw new UnsupportedOperationException();
118121
}

‎src/main/java/net/snowflake/client/core/SFJsonResultSet.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,11 @@ public Object getObject(int columnIndex) throws SFException {
103103
}
104104
}
105105

106+
@SnowflakeJdbcInternalApi
107+
@Override
108+
public Object getObjectWithoutString(int columnIndex) throws SFException {
109+
return getObject(columnIndex);
110+
}
106111
/**
107112
* Sometimes large BIGINTS overflow the java Long type. In these cases, return a BigDecimal type
108113
* instead.

‎src/main/java/net/snowflake/client/core/arrow/ArrayConverter.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ public ArrayConverter(ListVector valueVector, int vectorIndex, DataConversionCon
2525

2626
@Override
2727
public Object toObject(int index) throws SFException {
28-
return isNull(index) ? null : new StructObjectWrapper(toString(index), vector.getObject(index));
28+
return isNull(index) ? null : vector.getObject(index);
2929
}
3030

3131
@Override

‎src/main/java/net/snowflake/client/core/arrow/MapConverter.java

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import java.nio.charset.StandardCharsets;
44
import java.util.List;
5-
import java.util.Map;
65
import java.util.stream.Collectors;
76
import net.snowflake.client.core.DataConversionContext;
87
import net.snowflake.client.core.SFException;
@@ -33,15 +32,11 @@ public Object toObject(int index) throws SFException {
3332
if (isNull(index)) {
3433
return null;
3534
}
36-
3735
List<JsonStringHashMap<String, Object>> entriesList =
3836
(List<JsonStringHashMap<String, Object>>) vector.getObject(index);
39-
Map<String, Object> map =
40-
entriesList.stream()
41-
.collect(
42-
Collectors.toMap(
43-
entry -> entry.get("key").toString(), entry -> entry.get("value")));
44-
return new StructObjectWrapper(toString(index), map);
37+
return entriesList.stream()
38+
.collect(
39+
Collectors.toMap(entry -> entry.get("key").toString(), entry -> entry.get("value")));
4540
}
4641

4742
@Override

‎src/main/java/net/snowflake/client/core/arrow/StructConverter.java

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,7 @@ public StructConverter(StructVector vector, int columnIndex, DataConversionConte
2121

2222
@Override
2323
public Object toObject(int index) throws SFException {
24-
return isNull(index)
25-
? null
26-
: new StructObjectWrapper(toString(index), structVector.getObject(index));
24+
return isNull(index) ? null : structVector.getObject(index);
2725
}
2826

2927
@Override

‎src/main/java/net/snowflake/client/core/arrow/VectorTypeConverter.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,7 @@ public Object toObject(int index) throws SFException {
2727
if (isNull(index)) {
2828
return null;
2929
}
30-
Object object = vector.getObject(index);
31-
return new StructObjectWrapper(object.toString(), object);
30+
return vector.getObject(index);
3231
}
3332

3433
@Override

‎src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1398,8 +1398,11 @@ public <T> T getObject(int columnIndex, Class<T> type) throws SQLException {
13981398
SnowflakeUtil.mapSFExceptionToSQLException(
13991399
() -> {
14001400
StructObjectWrapper structObjectWrapper =
1401-
(StructObjectWrapper) sfBaseResultSet.getObject(columnIndex);
1402-
return (SQLInput) createJsonSqlInput(columnIndex, structObjectWrapper);
1401+
(StructObjectWrapper) sfBaseResultSet.getObjectWithoutString(columnIndex);
1402+
if (structObjectWrapper == null) {
1403+
return null;
1404+
}
1405+
return (SQLInput) structObjectWrapper.getObject();
14031406
});
14041407
if (sqlInput == null) {
14051408
return null;
@@ -1639,7 +1642,7 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
16391642
StructObjectWrapper structObjectWrapper =
16401643
(StructObjectWrapper)
16411644
SnowflakeUtil.mapSFExceptionToSQLException(
1642-
() -> sfBaseResultSet.getObject(columnIndex));
1645+
() -> sfBaseResultSet.getObjectWithoutString(columnIndex));
16431646
if (structObjectWrapper == null) {
16441647
return null;
16451648
}

‎src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,6 @@
2828
import java.util.List;
2929
import java.util.Map;
3030
import java.util.TimeZone;
31-
import net.snowflake.client.core.ArrowSqlInput;
32-
import net.snowflake.client.core.JsonSqlInput;
3331
import net.snowflake.client.core.QueryStatus;
3432
import net.snowflake.client.core.SFBaseResultSet;
3533
import net.snowflake.client.core.SFException;
@@ -274,18 +272,21 @@ public Object getObject(int columnIndex) throws SQLException {
274272
SnowflakeUtil.mapSFExceptionToSQLException(() -> sfBaseResultSet.getObject(columnIndex));
275273
if (object == null) {
276274
return null;
277-
} else if (object instanceof JsonSqlInput) {
278-
return ((JsonSqlInput) object).getText();
279-
} else if (object instanceof StructObjectWrapper) {
280-
return ((StructObjectWrapper) object).getJsonString();
281-
} else if (object instanceof SfSqlArray) {
275+
}
276+
if (object instanceof SfSqlArray) {
282277
return ((SfSqlArray) object).getText();
283-
} else if (object instanceof ArrowSqlInput) {
284-
throw new SQLException(
285-
"Arrow native struct couldn't be converted to String. To map to SqlData the method getObject(int columnIndex, Class type) should be used");
286-
} else {
287-
return object;
288278
}
279+
if (object instanceof StructObjectWrapper) {
280+
StructObjectWrapper structObjectWrapper = (StructObjectWrapper) object;
281+
if (resultSetMetaData.isStructuredTypeColumn(columnIndex)
282+
&& structObjectWrapper.getJsonString() != null) {
283+
return structObjectWrapper.getJsonString();
284+
}
285+
if (structObjectWrapper.getObject() != null) {
286+
return structObjectWrapper.getObject();
287+
}
288+
}
289+
return object;
289290
}
290291

291292
public Array getArray(int columnIndex) throws SQLException {

‎src/test/java/net/snowflake/client/jdbc/structuredtypes/ResultSetStructuredTypesLatestIT.java

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -263,10 +263,7 @@ public void testReturnAsArrayOfNullableFieldsInSqlData(ResultSetFormatType forma
263263
+ "'date', null, 'bd', null, 'bytes', null, 'longValue', null)"
264264
+ "::OBJECT(string VARCHAR, nullableIntValue INTEGER, nullableLongValue INTEGER, date DATE, bd DOUBLE, bytes BINARY, longValue INTEGER)",
265265
(resultSet) -> {
266-
NullableFieldsSqlData result =
267-
resultSet
268-
.unwrap(SnowflakeBaseResultSet.class)
269-
.getObject(1, NullableFieldsSqlData.class);
266+
NullableFieldsSqlData result = resultSet.getObject(1, NullableFieldsSqlData.class);
270267
assertNull(result.getString());
271268
assertNull(result.getNullableIntValue());
272269
assertNull(result.getNullableLongValue());

0 commit comments

Comments
 (0)