diff --git a/src/main/java/net/snowflake/client/core/ArrowSqlInput.java b/src/main/java/net/snowflake/client/core/ArrowSqlInput.java index 8e5d9a787..2ea95de0b 100644 --- a/src/main/java/net/snowflake/client/core/ArrowSqlInput.java +++ b/src/main/java/net/snowflake/client/core/ArrowSqlInput.java @@ -14,6 +14,7 @@ import java.sql.Timestamp; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.TimeZone; import net.snowflake.client.core.json.Converters; import net.snowflake.client.core.structs.SQLDataCreationHelper; @@ -24,16 +25,22 @@ @SnowflakeJdbcInternalApi public class ArrowSqlInput extends BaseSqlInput { + private final Map input; private final Iterator structuredTypeFields; private int currentIndex = 0; public ArrowSqlInput( - JsonStringHashMap input, + Map input, SFBaseSession session, Converters converters, List fields) { super(session, converters, fields); this.structuredTypeFields = input.values().iterator(); + this.input = input; + } + + public Map getInput() { + return input; } @Override @@ -172,6 +179,8 @@ public Timestamp readTimestamp(TimeZone tz) throws SQLException { if (value == null) { return null; } + int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session); + int columnSubType = fieldMetadata.getType(); int scale = fieldMetadata.getScale(); return mapSFExceptionToSQLException( () -> @@ -179,7 +188,8 @@ public Timestamp readTimestamp(TimeZone tz) throws SQLException { .getStructuredTypeDateTimeConverter() .getTimestamp( (JsonStringHashMap) value, - fieldMetadata.getBase(), + columnType, + columnSubType, tz, scale)); }); @@ -204,10 +214,7 @@ public T readObject(Class type) throws SQLException { SQLData instance = (SQLData) SQLDataCreationHelper.create(type); instance.readSQL( new ArrowSqlInput( - (JsonStringHashMap) value, - session, - converters, - fieldMetadata.getFields()), + (Map) value, session, converters, fieldMetadata.getFields()), null); return (T) instance; }); diff --git a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java index daf60b804..58aa3a709 100644 --- a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java +++ b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java @@ -19,6 +19,7 @@ import java.sql.Time; import java.sql.Timestamp; import java.sql.Types; +import java.util.Map; import java.util.TimeZone; import net.snowflake.client.core.arrow.ArrowVectorConverter; import net.snowflake.client.core.arrow.StructConverter; @@ -104,7 +105,7 @@ public class SFArrowResultSet extends SFBaseResultSet implements DataConversionC */ private boolean formatDateWithTimezone; - @SnowflakeJdbcInternalApi protected Converters jsonConverters; + @SnowflakeJdbcInternalApi protected Converters converters; /** * Constructor takes a result from the API response that we get from executing a SQL statement. @@ -124,7 +125,7 @@ public SFArrowResultSet( boolean sortResult) throws SQLException { this(resultSetSerializable, session.getTelemetryClient(), sortResult); - this.jsonConverters = + this.converters = new Converters( resultSetSerializable.getTimeZone(), session, @@ -356,6 +357,31 @@ private boolean fetchNextRowSorted() throws SnowflakeSQLException { } } + @Override + @SnowflakeJdbcInternalApi + public Converters getConverters() { + return converters; + } + + @Override + public Date convertToDate(Object object, TimeZone tz) throws SFException { + return converters.getStructuredTypeDateTimeConverter().getDate((int) object, tz); + } + + @Override + public Time convertToTime(Object object, int scale) throws SFException { + return converters.getStructuredTypeDateTimeConverter().getTime((long) object, scale); + } + + @Override + public Timestamp convertToTimestamp( + Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException { + return converters + .getStructuredTypeDateTimeConverter() + .getTimestamp( + (JsonStringHashMap) object, columnType, columnSubType, tz, scale); + } + /** * Advance to next row * @@ -510,7 +536,7 @@ public Object getObject(int columnIndex) throws SFException { if (converter instanceof VarCharConverter) { return createJsonSqlInput(columnIndex, obj); } else if (converter instanceof StructConverter) { - return createArrowSqlInput(columnIndex, (JsonStringHashMap) obj); + return createArrowSqlInput(columnIndex, (Map) obj); } } return obj; @@ -522,7 +548,7 @@ private Object createJsonSqlInput(int columnIndex, Object obj) throws SFExceptio return new JsonSqlInput( jsonNode, session, - jsonConverters, + converters, resultSetMetaData.getColumnMetadata().get(columnIndex - 1).getFields(), sessionTimezone); } catch (JsonProcessingException e) { @@ -530,11 +556,11 @@ private Object createJsonSqlInput(int columnIndex, Object obj) throws SFExceptio } } - private Object createArrowSqlInput(int columnIndex, JsonStringHashMap input) { + private Object createArrowSqlInput(int columnIndex, Map input) { return new ArrowSqlInput( input, session, - jsonConverters, + converters, resultSetMetaData.getColumnMetadata().get(columnIndex - 1).getFields()); } diff --git a/src/main/java/net/snowflake/client/core/SFBaseResultSet.java b/src/main/java/net/snowflake/client/core/SFBaseResultSet.java index 97509d2a7..a68bf82b6 100644 --- a/src/main/java/net/snowflake/client/core/SFBaseResultSet.java +++ b/src/main/java/net/snowflake/client/core/SFBaseResultSet.java @@ -213,4 +213,14 @@ public Converters getConverters() { public TimeZone getSessionTimeZone() { return resultSetSerializable.getTimeZone(); } + + @SnowflakeJdbcInternalApi + public abstract Date convertToDate(Object object, TimeZone tz) throws SFException; + + @SnowflakeJdbcInternalApi + public abstract Time convertToTime(Object object, int scale) throws SFException; + + @SnowflakeJdbcInternalApi + public abstract Timestamp convertToTimestamp( + Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException; } diff --git a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java index d8955bdeb..5a91b7591 100644 --- a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java +++ b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java @@ -426,4 +426,23 @@ private static Object convert(JsonStringToTypeConverter converter, JsonNode node return converter.convert(node.toString()); } } + + @Override + public Date convertToDate(Object object, TimeZone tz) throws SFException { + return (Date) converters.dateConverter(session).convert((String) object); + } + + @Override + public Time convertToTime(Object object, int scale) throws SFException { + return (Time) converters.timeConverter(session).convert((String) object); + } + + @Override + public Timestamp convertToTimestamp( + Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException { + return (Timestamp) + converters + .timestampConverter(columnSubType, columnType, scale, session, null, tz) + .convert((String) object); + } } diff --git a/src/main/java/net/snowflake/client/core/arrow/MapConverter.java b/src/main/java/net/snowflake/client/core/arrow/MapConverter.java new file mode 100644 index 000000000..433792294 --- /dev/null +++ b/src/main/java/net/snowflake/client/core/arrow/MapConverter.java @@ -0,0 +1,33 @@ +package net.snowflake.client.core.arrow; + +import java.util.List; +import java.util.stream.Collectors; +import net.snowflake.client.core.DataConversionContext; +import net.snowflake.client.core.SFException; +import net.snowflake.client.jdbc.SnowflakeType; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.util.JsonStringHashMap; + +public class MapConverter extends AbstractArrowVectorConverter { + + private final MapVector vector; + + public MapConverter(MapVector valueVector, int columnIndex, DataConversionContext context) { + super(SnowflakeType.MAP.name(), valueVector, columnIndex, context); + this.vector = valueVector; + } + + @Override + public Object toObject(int index) throws SFException { + List> entriesList = + (List>) vector.getObject(index); + return entriesList.stream() + .collect( + Collectors.toMap(entry -> entry.get("key").toString(), entry -> entry.get("value"))); + } + + @Override + public String toString(int index) throws SFException { + return vector.getObject(index).toString(); + } +} diff --git a/src/main/java/net/snowflake/client/core/arrow/StructuredTypeDateTimeConverter.java b/src/main/java/net/snowflake/client/core/arrow/StructuredTypeDateTimeConverter.java index cd30c4bf5..c5ff2d19b 100644 --- a/src/main/java/net/snowflake/client/core/arrow/StructuredTypeDateTimeConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/StructuredTypeDateTimeConverter.java @@ -11,11 +11,12 @@ import java.sql.Date; import java.sql.Time; import java.sql.Timestamp; +import java.sql.Types; import java.util.TimeZone; import net.snowflake.client.core.SFException; import net.snowflake.client.core.SnowflakeJdbcInternalApi; import net.snowflake.client.jdbc.ErrorCode; -import net.snowflake.client.jdbc.SnowflakeType; +import net.snowflake.client.jdbc.SnowflakeUtil; import org.apache.arrow.vector.util.JsonStringHashMap; @SnowflakeJdbcInternalApi @@ -45,22 +46,33 @@ public StructuredTypeDateTimeConverter( } public Timestamp getTimestamp( - JsonStringHashMap obj, SnowflakeType type, TimeZone tz, int scale) + JsonStringHashMap obj, + int columnType, + int columnSubType, + TimeZone tz, + int scale) throws SFException { if (tz == null) { tz = TimeZone.getDefault(); } - switch (type) { - case TIMESTAMP_LTZ: + if (Types.TIMESTAMP == columnType) { + if (SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_LTZ == columnSubType) { return convertTimestampLtz(obj, scale); - case TIMESTAMP_NTZ: + } else { return convertTimestampNtz(obj, tz, scale); - case TIMESTAMP_TZ: - return convertTimestampTz(obj, scale); + } + } else if (Types.TIMESTAMP_WITH_TIMEZONE == columnType + && SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_TZ == columnSubType) { + return convertTimestampTz(obj, scale); } throw new SFException( ErrorCode.INVALID_VALUE_CONVERT, - "Unexpected Arrow Field for " + type.name() + " and object type " + obj.getClass()); + "Unexpected Arrow Field for columnType " + + columnType + + " , column subtype " + + columnSubType + + " , and object type " + + obj.getClass()); } public Date getDate(int value, TimeZone tz) throws SFException { diff --git a/src/main/java/net/snowflake/client/jdbc/ArrowResultChunk.java b/src/main/java/net/snowflake/client/jdbc/ArrowResultChunk.java index c273a8817..123eb0139 100644 --- a/src/main/java/net/snowflake/client/jdbc/ArrowResultChunk.java +++ b/src/main/java/net/snowflake/client/jdbc/ArrowResultChunk.java @@ -27,6 +27,7 @@ import net.snowflake.client.core.arrow.IntToFixedConverter; import net.snowflake.client.core.arrow.IntToScaledFixedConverter; import net.snowflake.client.core.arrow.IntToTimeConverter; +import net.snowflake.client.core.arrow.MapConverter; import net.snowflake.client.core.arrow.SmallIntToFixedConverter; import net.snowflake.client.core.arrow.SmallIntToScaledFixedConverter; import net.snowflake.client.core.arrow.StructConverter; @@ -55,6 +56,7 @@ import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.ipc.ArrowStreamReader; import org.apache.arrow.vector.types.Types; @@ -206,6 +208,10 @@ private static List initConverters( converters.add(new VarCharConverter(vector, i, context)); break; + case MAP: + converters.add(new MapConverter((MapVector) vector, i, context)); + break; + case OBJECT: if (vector instanceof StructVector) { converters.add(new StructConverter((StructVector) vector, i, context)); diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java index 6f2804b3a..2bad0a4e5 100644 --- a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java +++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java @@ -37,6 +37,7 @@ import java.util.List; import java.util.Map; import java.util.TimeZone; +import net.snowflake.client.core.ArrowSqlInput; import net.snowflake.client.core.ColumnTypeHelper; import net.snowflake.client.core.JsonSqlInput; import net.snowflake.client.core.ObjectMapperFactory; @@ -1355,8 +1356,14 @@ public T getObject(int columnIndex, Class type) throws SQLException { instance.readSQL(sqlInput, null); return (T) instance; } else if (Map.class.isAssignableFrom(type)) { - JsonNode jsonNode = ((JsonSqlInput) getObject(columnIndex)).getInput(); - return (T) OBJECT_MAPPER.convertValue(jsonNode, new TypeReference>() {}); + Object object = getObject(columnIndex); + if (object instanceof JsonSqlInput) { + JsonNode jsonNode = ((JsonSqlInput) object).getInput(); + return (T) + OBJECT_MAPPER.convertValue(jsonNode, new TypeReference>() {}); + } else { + return (T) ((ArrowSqlInput) object).getInput(); + } } else if (String.class.isAssignableFrom(type)) { return (T) getString(columnIndex); } else if (Boolean.class.isAssignableFrom(type)) { @@ -1543,20 +1550,50 @@ public Map getMap(int columnIndex, Class type) throws SQLExcep int scale = valueFieldMetadata.getScale(); TimeZone tz = sfBaseResultSet.getSessionTimeZone(); Object object = getObject(columnIndex); - JsonNode jsonNode = ((JsonSqlInput) object).getInput(); - Map map = - OBJECT_MAPPER.convertValue(jsonNode, new TypeReference>() {}); + Map map; + if (object instanceof JsonSqlInput) { + map = new HashMap<>(); + JsonNode jsonNode = ((JsonSqlInput) object).getInput(); + jsonNode + .fieldNames() + .forEachRemaining(node -> map.put(node.toString(), jsonNode.get(node.toString()))); + } else { + map = (Map) object; + } Map resultMap = new HashMap<>(); for (Map.Entry entry : map.entrySet()) { if (SQLData.class.isAssignableFrom(type)) { SQLData instance = (SQLData) SQLDataCreationHelper.create(type); - SQLInput sqlInput = - new JsonSqlInput( - jsonNode.get(entry.getKey()), - session, - sfBaseResultSet.getConverters(), - sfBaseResultSet.getMetaData().getColumnMetadata().get(columnIndex - 1).getFields(), - sfBaseResultSet.getSessionTimezone()); + SQLInput sqlInput; + if (object instanceof JsonSqlInput) { + sqlInput = + new JsonSqlInput( + (JsonNode) entry.getValue(), + session, + sfBaseResultSet.getConverters(), + sfBaseResultSet + .getMetaData() + .getColumnMetadata() + .get(columnIndex - 1) + .getFields(), + sfBaseResultSet.getSessionTimezone()); + } else if (object instanceof ArrowSqlInput || object instanceof Map) { + sqlInput = + new ArrowSqlInput( + (Map) entry.getValue(), + session, + sfBaseResultSet.getConverters(), + sfBaseResultSet + .getMetaData() + .getColumnMetadata() + .get(columnIndex - 1) + .getFields()); + } else { + throw new SQLException( + "SqlInput type " + + object.getClass() + + " is not supported when mapping to SQLData class"); + } instance.readSQL(sqlInput, null); resultMap.put(entry.getKey(), (T) instance); } else if (String.class.isAssignableFrom(type)) { @@ -1658,32 +1695,22 @@ public Map getMap(int columnIndex, Class type) throws SQLExcep resultMap.put( entry.getKey(), mapSFExceptionToSQLException( - () -> - (T) - sfBaseResultSet - .getConverters() - .dateConverter(session) - .convert((String) entry.getValue()))); + () -> (T) sfBaseResultSet.convertToDate(entry.getValue(), tz))); } else if (Time.class.isAssignableFrom(type)) { resultMap.put( entry.getKey(), mapSFExceptionToSQLException( - () -> - (T) - sfBaseResultSet - .getConverters() - .timeConverter(session) - .convert((String) entry.getValue()))); + () -> (T) sfBaseResultSet.convertToTime(entry.getValue(), scale))); + } else if (Timestamp.class.isAssignableFrom(type)) { resultMap.put( entry.getKey(), mapSFExceptionToSQLException( () -> (T) - sfBaseResultSet - .getConverters() - .timestampConverter(columnSubType, columnType, scale, session, null, tz) - .convert((String) entry.getValue()))); + sfBaseResultSet.convertToTimestamp( + entry.getValue(), columnType, columnSubType, tz, scale))); + } else { logger.debug( "Unsupported type passed to getObject(int columnIndex,Class type): " diff --git a/src/test/java/net/snowflake/client/jdbc/ResultSetStructuredTypesLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ResultSetStructuredTypesLatestIT.java index 492214894..81833c8dc 100644 --- a/src/test/java/net/snowflake/client/jdbc/ResultSetStructuredTypesLatestIT.java +++ b/src/test/java/net/snowflake/client/jdbc/ResultSetStructuredTypesLatestIT.java @@ -173,7 +173,6 @@ private void testMapAllTypes(boolean registerFactory) throws SQLException { @Test public void testMapJsonToMap() throws SQLException { - Assume.assumeTrue(queryResultFormat != ResultSetFormatType.NATIVE_ARROW); withFirstRow( "SELECT OBJECT_CONSTRUCT('string','a','string2',1)", (resultSet) -> { @@ -200,6 +199,7 @@ public void testReturnAsArrayOfSqlData() throws SQLException { @Test public void testReturnAsArrayOfString() throws SQLException { + Assume.assumeTrue(queryResultFormat != ResultSetFormatType.NATIVE_ARROW); withFirstRow( "SELECT ARRAY_CONSTRUCT('one', 'two','three')::ARRAY(VARCHAR)", (resultSet) -> { @@ -213,6 +213,7 @@ public void testReturnAsArrayOfString() throws SQLException { @Test public void testReturnAsListOfIntegers() throws SQLException { + Assume.assumeTrue(queryResultFormat != ResultSetFormatType.NATIVE_ARROW); withFirstRow( "SELECT ARRAY_CONSTRUCT(1,2,3)::ARRAY(INTEGER)", (resultSet) -> { @@ -226,7 +227,6 @@ public void testReturnAsListOfIntegers() throws SQLException { @Test public void testReturnAsMap() throws SQLException { - Assume.assumeTrue(queryResultFormat != ResultSetFormatType.NATIVE_ARROW); SnowflakeObjectTypeFactories.register(SimpleClass.class, SimpleClass::new); withFirstRow( "select {'x':{'string':'one'},'y':{'string':'two'},'z':{'string':'three'}}::MAP(VARCHAR, OBJECT(string VARCHAR));", @@ -239,6 +239,18 @@ public void testReturnAsMap() throws SQLException { }); } + @Test + public void testReturnAsMapOfTimestampsNtz() throws SQLException { + withFirstRow( + "SELECT {'x': TO_TIMESTAMP_NTZ('2021-12-23 09:44:44'), 'y': TO_TIMESTAMP_NTZ('2021-12-24 09:55:55')}::MAP(VARCHAR, TIMESTAMP)", + (resultSet) -> { + Map map = + resultSet.unwrap(SnowflakeBaseResultSet.class).getMap(1, Timestamp.class); + assertEquals(Timestamp.valueOf(LocalDateTime.of(2021, 12, 23, 9, 44, 44)), map.get("x")); + assertEquals(Timestamp.valueOf(LocalDateTime.of(2021, 12, 24, 9, 55, 55)), map.get("y")); + }); + } + @Test public void testReturnAsMapOfLong() throws SQLException { withFirstRow( @@ -271,8 +283,10 @@ public void testReturnAsMapOfDate() throws SQLException { (resultSet) -> { Map map = resultSet.unwrap(SnowflakeBaseResultSet.class).getMap(1, Date.class); - assertEquals(Date.valueOf(LocalDate.of(2023, 12, 24)), map.get("x")); - assertEquals(Date.valueOf(LocalDate.of(2023, 12, 25)), map.get("y")); + assertEquals( + Date.valueOf(LocalDate.of(2023, 12, 24)).toString(), map.get("x").toString()); + assertEquals( + Date.valueOf(LocalDate.of(2023, 12, 25)).toString(), map.get("y").toString()); }); } @@ -316,7 +330,6 @@ public void testReturnAsList() throws SQLException { @Test public void testMapStructsFromChunks() throws SQLException { - Assume.assumeTrue(queryResultFormat != ResultSetFormatType.NATIVE_ARROW); withFirstRow( "select {'string':'a'}::OBJECT(string VARCHAR) FROM TABLE(GENERATOR(ROWCOUNT=>30000))", (resultSet) -> {