Skip to content

Commit e8e3bac

Browse files
sfc-gh-mkubiksfc-gh-pmotacki
sfc-gh-mkubik
and
sfc-gh-pmotacki
authored
SNOW-1856886 Narrow toString calculation to structured types only (#2007)
Co-authored-by: Przemyslaw Motacki <[email protected]>
1 parent 1d58f57 commit e8e3bac

File tree

5 files changed

+52
-40
lines changed

5 files changed

+52
-40
lines changed

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

+32-13
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.MapConverter;
2930
import net.snowflake.client.core.arrow.StructConverter;
3031
import net.snowflake.client.core.arrow.StructObjectWrapper;
3132
import net.snowflake.client.core.arrow.VarCharConverter;
@@ -574,10 +575,11 @@ public Object getObjectWithoutString(int columnIndex) throws SFException {
574575
return getObjectRepresentation(columnIndex, false);
575576
}
576577

577-
private Object getObjectRepresentation(int columnIndex, boolean withString) throws SFException {
578+
private StructObjectWrapper getObjectRepresentation(int columnIndex, boolean withString)
579+
throws SFException {
578580
int type = resultSetMetaData.getColumnType(columnIndex);
579581
if (type == SnowflakeUtil.EXTRA_TYPES_VECTOR) {
580-
return getString(columnIndex);
582+
return new StructObjectWrapper(getString(columnIndex), null);
581583
}
582584
ArrowVectorConverter converter = currentChunkIterator.getCurrentConverter(columnIndex - 1);
583585
int index = currentChunkIterator.getCurrentRowInRecordBatch();
@@ -589,15 +591,33 @@ private Object getObjectRepresentation(int columnIndex, boolean withString) thro
589591
if (obj == null) {
590592
return null;
591593
}
592-
String jsonString = withString ? converter.toString(index) : null;
593594
boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
594-
if (isVarcharConvertedStruct(type, isStructuredType, converter)) {
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));
595+
if (isStructuredType) {
596+
if (converter instanceof VarCharConverter) {
597+
if (type == Types.STRUCT) {
598+
JsonSqlInput jsonSqlInput = createJsonSqlInput(columnIndex, obj);
599+
return new StructObjectWrapper(jsonSqlInput.getText(), jsonSqlInput);
600+
} else if (type == Types.ARRAY) {
601+
SfSqlArray sfArray = getJsonArray((String) obj, columnIndex);
602+
return new StructObjectWrapper(sfArray.getText(), sfArray);
603+
} else {
604+
throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
605+
}
606+
} else if (converter instanceof StructConverter) {
607+
String jsonString = withString ? converter.toString(index) : null;
608+
return new StructObjectWrapper(
609+
jsonString, createArrowSqlInput(columnIndex, (Map<String, Object>) obj));
610+
} else if (converter instanceof MapConverter) {
611+
String jsonString = withString ? converter.toString(index) : null;
612+
return new StructObjectWrapper(jsonString, obj);
613+
} else if (converter instanceof ArrayConverter || converter instanceof VectorTypeConverter) {
614+
String jsonString = converter.toString(index);
615+
return new StructObjectWrapper(jsonString, obj);
616+
} else {
617+
throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
618+
}
599619
} else {
600-
return new StructObjectWrapper(jsonString, obj);
620+
return new StructObjectWrapper(null, obj);
601621
}
602622
}
603623

@@ -610,12 +630,11 @@ private SQLInput createArrowSqlInput(int columnIndex, Map<String, Object> input)
610630
input, session, converters, resultSetMetaData.getColumnFields(columnIndex));
611631
}
612632

613-
private boolean isVarcharConvertedStruct(
614-
int type, boolean isStructuredType, ArrowVectorConverter converter) {
615-
return type == Types.STRUCT && isStructuredType && converter instanceof VarCharConverter;
633+
private boolean isVarcharConvertedStruct(int type, ArrowVectorConverter converter) {
634+
return (type == Types.STRUCT || type == Types.ARRAY) && converter instanceof VarCharConverter;
616635
}
617636

618-
private Object createJsonSqlInput(int columnIndex, Object obj) throws SFException {
637+
private JsonSqlInput createJsonSqlInput(int columnIndex, Object obj) throws SFException {
619638
try {
620639
if (obj == null) {
621640
return null;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ public Object getObject(int columnIndex) throws SFException {
9494
}
9595
case Types.ARRAY:
9696
if (resultSetMetaData.isStructuredTypeColumn(columnIndex)) {
97-
return getArray(columnIndex);
97+
return new StructObjectWrapper((String) obj, getArray(columnIndex));
9898
} else {
9999
throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "data type: " + type);
100100
}

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

+1-19
Original file line numberDiff line numberDiff line change
@@ -1643,7 +1643,7 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
16431643
(StructObjectWrapper)
16441644
SnowflakeUtil.mapSFExceptionToSQLException(
16451645
() -> sfBaseResultSet.getObjectWithoutString(columnIndex));
1646-
if (structObjectWrapper == null) {
1646+
if (structObjectWrapper == null || structObjectWrapper.getObject() == null) {
16471647
return null;
16481648
}
16491649
Map<String, Object> map =
@@ -1835,22 +1835,4 @@ private <T> Map<String, Object> prepareMapWithValues(Object object, Class<T> typ
18351835
throw new SFException(ErrorCode.INVALID_STRUCT_DATA, "Object couldn't be converted to map");
18361836
}
18371837
}
1838-
1839-
private Object createJsonSqlInput(int columnIndex, StructObjectWrapper obj) throws SFException {
1840-
try {
1841-
if (obj == null) {
1842-
return null;
1843-
}
1844-
JsonNode jsonNode = OBJECT_MAPPER.readTree(obj.getJsonString());
1845-
return new JsonSqlInput(
1846-
obj.getJsonString(),
1847-
jsonNode,
1848-
session,
1849-
sfBaseResultSet.getConverters(),
1850-
sfBaseResultSet.getMetaData().getColumnFields(columnIndex),
1851-
sfBaseResultSet.getSessionTimeZone());
1852-
} catch (JsonProcessingException e) {
1853-
throw new SFException(sfBaseResultSet.getQueryId(), e, ErrorCode.INVALID_STRUCT_DATA);
1854-
}
1855-
}
18561838
}

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

+5-7
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
import net.snowflake.client.core.QueryStatus;
3232
import net.snowflake.client.core.SFBaseResultSet;
3333
import net.snowflake.client.core.SFException;
34-
import net.snowflake.client.core.SfSqlArray;
3534
import net.snowflake.client.core.arrow.StructObjectWrapper;
3635
import net.snowflake.client.log.SFLogger;
3736
import net.snowflake.client.log.SFLoggerFactory;
@@ -270,22 +269,21 @@ public Object getObject(int columnIndex) throws SQLException {
270269
raiseSQLExceptionIfResultSetIsClosed();
271270
Object object =
272271
SnowflakeUtil.mapSFExceptionToSQLException(() -> sfBaseResultSet.getObject(columnIndex));
272+
273273
if (object == null) {
274274
return null;
275275
}
276-
if (object instanceof SfSqlArray) {
277-
return ((SfSqlArray) object).getText();
278-
}
276+
279277
if (object instanceof StructObjectWrapper) {
280278
StructObjectWrapper structObjectWrapper = (StructObjectWrapper) object;
281279
if (resultSetMetaData.isStructuredTypeColumn(columnIndex)
282280
&& structObjectWrapper.getJsonString() != null) {
283281
return structObjectWrapper.getJsonString();
284282
}
285-
if (structObjectWrapper.getObject() != null) {
286-
return structObjectWrapper.getObject();
287-
}
283+
284+
return structObjectWrapper.getObject();
288285
}
286+
289287
return object;
290288
}
291289

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

+13
Original file line numberDiff line numberDiff line change
@@ -635,6 +635,19 @@ public void testMapIntegerArray(ResultSetFormatType format) throws SQLException
635635
format);
636636
}
637637

638+
@ParameterizedTest
639+
@ArgumentsSource(ResultFormatProvider.class)
640+
@DontRunOnGithubActions
641+
public void testMapIntegerArrayGetObject(ResultSetFormatType format) throws SQLException {
642+
withFirstRow(
643+
"SELECT ARRAY_CONSTRUCT(10, 20, 30)::ARRAY(INTEGER)",
644+
(resultSet) -> {
645+
Object resultArray = resultSet.getObject(1);
646+
TestUtil.assertEqualsIgnoringWhitespace("[10,20,30]", (String) resultArray);
647+
},
648+
format);
649+
}
650+
638651
@ParameterizedTest
639652
@ArgumentsSource(ResultFormatProvider.class)
640653
@DontRunOnGithubActions

0 commit comments

Comments
 (0)