Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 2cfa1e9

Browse files
sfc-gh-dheymansfc-gh-dheyman
committedMar 7, 2024·
Add native arrow structured types support
1 parent 63ba15c commit 2cfa1e9

14 files changed

+510
-71
lines changed
 

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

+293
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,293 @@
1+
package net.snowflake.client.core;
2+
3+
import net.snowflake.client.core.json.Converters;
4+
import net.snowflake.client.core.structs.SQLDataCreationHelper;
5+
import net.snowflake.client.jdbc.FieldMetadata;
6+
import net.snowflake.client.jdbc.SnowflakeLoggedFeatureNotSupportedException;
7+
import net.snowflake.client.util.ThrowingBiFunction;
8+
import net.snowflake.common.core.SFTimestamp;
9+
import net.snowflake.common.core.SnowflakeDateTimeFormat;
10+
import org.apache.arrow.vector.util.JsonStringHashMap;
11+
12+
import java.io.InputStream;
13+
import java.io.Reader;
14+
import java.math.BigDecimal;
15+
import java.net.URL;
16+
import java.sql.*;
17+
import java.time.Instant;
18+
import java.time.ZoneOffset;
19+
import java.util.Arrays;
20+
import java.util.Iterator;
21+
import java.util.List;
22+
import java.util.TimeZone;
23+
24+
import static net.snowflake.client.jdbc.SnowflakeUtil.mapExceptions;
25+
26+
@SnowflakeJdbcInternalApi
27+
public class ArrowSqlInput implements SFSqlInput {
28+
29+
private final JsonStringHashMap<String, Object> input;
30+
private final SFBaseSession session;
31+
private final Iterator<Object> elements;
32+
private final Converters converters;
33+
private final List<FieldMetadata> fields;
34+
35+
private int currentIndex = 0;
36+
37+
public ArrowSqlInput(JsonStringHashMap<String, Object> input, SFBaseSession session, Converters converters, List<FieldMetadata> fields) {
38+
this.input = input;
39+
this.elements = input.values().iterator();
40+
this.session = session;
41+
this.converters = converters;
42+
this.fields = fields;
43+
}
44+
45+
@Override
46+
public String readString() throws SQLException {
47+
return withNextValue(
48+
((value, fieldMetadata) -> {
49+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
50+
int columnSubType = fieldMetadata.getType();
51+
int scale = fieldMetadata.getScale();
52+
return mapExceptions(
53+
() ->
54+
converters
55+
.getStringConverter()
56+
.getString(value, columnType, columnSubType, scale));
57+
}));
58+
}
59+
60+
@Override
61+
public boolean readBoolean() throws SQLException {
62+
return withNextValue((value, fieldMetadata) -> {
63+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
64+
return mapExceptions(
65+
() -> converters.getBooleanConverter().getBoolean(value, columnType));
66+
});
67+
}
68+
69+
@Override
70+
public byte readByte() throws SQLException {
71+
return withNextValue(
72+
(value, fieldMetadata) ->
73+
mapExceptions(() -> converters.getNumberConverter().getByte(value)));
74+
}
75+
76+
@Override
77+
public short readShort() throws SQLException {
78+
return withNextValue(
79+
(value, fieldMetadata) -> {
80+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
81+
return mapExceptions(() -> converters.getNumberConverter().getShort(value, columnType));
82+
});
83+
}
84+
85+
@Override
86+
public int readInt() throws SQLException {
87+
return withNextValue(
88+
(value, fieldMetadata) -> {
89+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
90+
return mapExceptions(() -> converters.getNumberConverter().getInt(value, columnType));
91+
});
92+
}
93+
94+
@Override
95+
public long readLong() throws SQLException {
96+
return withNextValue(
97+
(value, fieldMetadata) -> {
98+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
99+
return mapExceptions(() -> converters.getNumberConverter().getLong(value, columnType));
100+
});
101+
}
102+
103+
@Override
104+
public float readFloat() throws SQLException {
105+
return withNextValue(
106+
(value, fieldMetadata) -> {
107+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
108+
return mapExceptions(() -> converters.getNumberConverter().getFloat(value, columnType));
109+
});
110+
}
111+
112+
@Override
113+
public double readDouble() throws SQLException {
114+
return withNextValue(
115+
(value, fieldMetadata) -> {
116+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
117+
return mapExceptions(() -> converters.getNumberConverter().getDouble(value, columnType));
118+
});
119+
}
120+
121+
@Override
122+
public BigDecimal readBigDecimal() throws SQLException {
123+
return withNextValue(
124+
(value, fieldMetadata) -> {
125+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
126+
return mapExceptions(
127+
() -> converters.getNumberConverter().getBigDecimal(value, columnType));
128+
});
129+
}
130+
131+
@Override
132+
public byte[] readBytes() throws SQLException {
133+
return withNextValue(
134+
(value, fieldMetadata) -> {
135+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
136+
int columnSubType = fieldMetadata.getType();
137+
int scale = fieldMetadata.getScale();
138+
return mapExceptions(
139+
() ->
140+
converters.getBytesConverter().getBytes(value, columnType, columnSubType, scale));
141+
});
142+
}
143+
144+
@Override
145+
public Date readDate() throws SQLException {
146+
return withNextValue(
147+
(value, fieldMetadata) -> {
148+
SnowflakeDateTimeFormat formatter = SqlInputTimestampUtil.extractDateTimeFormat(session, "DATE_OUTPUT_FORMAT");
149+
SFTimestamp timestamp = formatter.parse((String) value);
150+
return Date.valueOf(
151+
Instant.ofEpochMilli(timestamp.getTime()).atZone(ZoneOffset.UTC).toLocalDate());
152+
});
153+
}
154+
155+
@Override
156+
public Time readTime() throws SQLException {
157+
return withNextValue(
158+
(value, fieldMetadata) -> {
159+
SnowflakeDateTimeFormat formatter = SqlInputTimestampUtil.extractDateTimeFormat(session, "TIME_OUTPUT_FORMAT");
160+
SFTimestamp timestamp = formatter.parse((String) value);
161+
return Time.valueOf(
162+
Instant.ofEpochMilli(timestamp.getTime()).atZone(ZoneOffset.UTC).toLocalTime());
163+
});
164+
}
165+
166+
@Override
167+
public Timestamp readTimestamp() throws SQLException {
168+
return readTimestamp(null);
169+
}
170+
171+
@Override
172+
public Timestamp readTimestamp(TimeZone tz) throws SQLException {
173+
return withNextValue(
174+
(value, fieldMetadata) -> {
175+
if (value == null) {
176+
return null;
177+
}
178+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
179+
int columnSubType = fieldMetadata.getType();
180+
int scale = fieldMetadata.getScale();
181+
// TODO structuredType what if not a string value?
182+
Timestamp result = SqlInputTimestampUtil.getTimestampFromType(columnSubType, (String) value, session);
183+
if (result != null) {
184+
return result;
185+
}
186+
return mapExceptions(
187+
() ->
188+
converters
189+
.getDateTimeConverter()
190+
.getTimestamp(value, columnType, columnSubType, tz, scale));
191+
});
192+
}
193+
194+
195+
@Override
196+
public Reader readCharacterStream() throws SQLException {
197+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readCharacterStream");
198+
}
199+
200+
@Override
201+
public InputStream readAsciiStream() throws SQLException {
202+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readAsciiStream");
203+
}
204+
205+
@Override
206+
public InputStream readBinaryStream() throws SQLException {
207+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readBinaryStream");
208+
}
209+
210+
@Override
211+
public Object readObject() throws SQLException {
212+
return withNextValue((value, fieldMetadata) -> {
213+
if (!(value instanceof JsonStringHashMap)) {
214+
throw new SQLException("Invalid value passed to 'readObject()', expected Map; got: " + value.getClass());
215+
}
216+
return value;
217+
});
218+
}
219+
220+
@Override
221+
public <T> T readObject(Class<T> type) throws SQLException {
222+
return withNextValue(
223+
(value, fieldMetadata) -> {
224+
SQLData instance = (SQLData) SQLDataCreationHelper.create(type);
225+
instance.readSQL(
226+
new ArrowSqlInput(
227+
(JsonStringHashMap<String, Object>) value,
228+
session,
229+
converters,
230+
Arrays.asList(fieldMetadata.getFields())
231+
),
232+
null
233+
);
234+
return (T) instance;
235+
});
236+
}
237+
238+
@Override
239+
public Ref readRef() throws SQLException {
240+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readRef");
241+
}
242+
243+
@Override
244+
public Blob readBlob() throws SQLException {
245+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readBlob");
246+
}
247+
248+
@Override
249+
public Clob readClob() throws SQLException {
250+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readClob");
251+
}
252+
253+
@Override
254+
public Array readArray() throws SQLException {
255+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readArray");
256+
}
257+
258+
@Override
259+
public boolean wasNull() throws SQLException {
260+
return false; // nulls are not allowed in structure types
261+
}
262+
263+
@Override
264+
public URL readURL() throws SQLException {
265+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readCharacterStream");
266+
}
267+
268+
@Override
269+
public NClob readNClob() throws SQLException {
270+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readNClob");
271+
}
272+
273+
@Override
274+
public String readNString() throws SQLException {
275+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readNString");
276+
}
277+
278+
@Override
279+
public SQLXML readSQLXML() throws SQLException {
280+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readSQLXML");
281+
}
282+
283+
@Override
284+
public RowId readRowId() throws SQLException {
285+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readRowId");
286+
}
287+
288+
private <T> T withNextValue(
289+
ThrowingBiFunction<Object, FieldMetadata, T, SQLException> action)
290+
throws SQLException {
291+
return action.apply(elements.next(), fields.get(currentIndex++));
292+
}
293+
}

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

+5-41
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import java.sql.SQLXML;
2121
import java.sql.Time;
2222
import java.sql.Timestamp;
23-
import java.sql.Types;
2423
import java.time.Instant;
2524
import java.time.ZoneOffset;
2625
import java.util.Arrays;
@@ -31,12 +30,12 @@
3130
import net.snowflake.client.core.structs.SQLDataCreationHelper;
3231
import net.snowflake.client.jdbc.FieldMetadata;
3332
import net.snowflake.client.jdbc.SnowflakeLoggedFeatureNotSupportedException;
34-
import net.snowflake.client.jdbc.SnowflakeUtil;
35-
import net.snowflake.client.util.ThrowingCallable;
3633
import net.snowflake.client.util.ThrowingTriFunction;
3734
import net.snowflake.common.core.SFTimestamp;
3835
import net.snowflake.common.core.SnowflakeDateTimeFormat;
3936

37+
import static net.snowflake.client.jdbc.SnowflakeUtil.mapExceptions;
38+
4039
@SnowflakeJdbcInternalApi
4140
public class JsonSqlInput implements SFSqlInput {
4241
private final JsonNode input;
@@ -163,7 +162,7 @@ public byte[] readBytes() throws SQLException {
163162
public Date readDate() throws SQLException {
164163
return withNextValue(
165164
(value, jsonNode, fieldMetadata) -> {
166-
SnowflakeDateTimeFormat formatter = getFormat(session, "DATE_OUTPUT_FORMAT");
165+
SnowflakeDateTimeFormat formatter = SqlInputTimestampUtil.extractDateTimeFormat(session, "DATE_OUTPUT_FORMAT");
167166
SFTimestamp timestamp = formatter.parse((String) value);
168167
return Date.valueOf(
169168
Instant.ofEpochMilli(timestamp.getTime()).atZone(ZoneOffset.UTC).toLocalDate());
@@ -174,7 +173,7 @@ public Date readDate() throws SQLException {
174173
public Time readTime() throws SQLException {
175174
return withNextValue(
176175
(value, jsonNode, fieldMetadata) -> {
177-
SnowflakeDateTimeFormat formatter = getFormat(session, "TIME_OUTPUT_FORMAT");
176+
SnowflakeDateTimeFormat formatter = SqlInputTimestampUtil.extractDateTimeFormat(session, "TIME_OUTPUT_FORMAT");
178177
SFTimestamp timestamp = formatter.parse((String) value);
179178
return Time.valueOf(
180179
Instant.ofEpochMilli(timestamp.getTime()).atZone(ZoneOffset.UTC).toLocalTime());
@@ -197,7 +196,7 @@ public Timestamp readTimestamp(TimeZone tz) throws SQLException {
197196
int columnSubType = fieldMetadata.getType();
198197
int scale = fieldMetadata.getScale();
199198
// TODO structuredType what if not a string value?
200-
Timestamp result = getTimestampFromType(columnSubType, (String) value);
199+
Timestamp result = SqlInputTimestampUtil.getTimestampFromType(columnSubType, (String) value, session);
201200
if (result != null) {
202201
return result;
203202
}
@@ -209,28 +208,6 @@ public Timestamp readTimestamp(TimeZone tz) throws SQLException {
209208
});
210209
}
211210

212-
private Timestamp getTimestampFromType(int columnSubType, String value) {
213-
if (columnSubType == SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_LTZ) {
214-
return getTimestampFromFormat("TIMESTAMP_LTZ_OUTPUT_FORMAT", value);
215-
} else if (columnSubType == SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_NTZ
216-
|| columnSubType == Types.TIMESTAMP) {
217-
return getTimestampFromFormat("TIMESTAMP_NTZ_OUTPUT_FORMAT", value);
218-
} else if (columnSubType == SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_TZ) {
219-
return getTimestampFromFormat("TIMESTAMP_TZ_OUTPUT_FORMAT", value);
220-
} else {
221-
return null;
222-
}
223-
}
224-
225-
private Timestamp getTimestampFromFormat(String format, String value) {
226-
String rawFormat = (String) session.getCommonParameters().get(format);
227-
if (rawFormat == null || rawFormat.isEmpty()) {
228-
rawFormat = (String) session.getCommonParameters().get("TIMESTAMP_OUTPUT_FORMAT");
229-
}
230-
SnowflakeDateTimeFormat formatter = SnowflakeDateTimeFormat.fromSqlFormat(rawFormat);
231-
return formatter.parse(value).getTimestamp();
232-
}
233-
234211
@Override
235212
public Reader readCharacterStream() throws SQLException {
236213
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readCharacterStream");
@@ -333,17 +310,4 @@ private Object getValue(JsonNode jsonNode) {
333310
}
334311
return null;
335312
}
336-
337-
private <T> T mapExceptions(ThrowingCallable<T, SFException> action) throws SQLException {
338-
try {
339-
return action.call();
340-
} catch (SFException e) {
341-
throw new SQLException(e);
342-
}
343-
}
344-
345-
private static SnowflakeDateTimeFormat getFormat(SFBaseSession session, String format) {
346-
return SnowflakeDateTimeFormat.fromSqlFormat(
347-
(String) session.getCommonParameters().get(format));
348-
}
349313
}

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

+32-15
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import java.util.Arrays;
2222
import java.util.TimeZone;
2323
import net.snowflake.client.core.arrow.ArrowVectorConverter;
24+
import net.snowflake.client.core.arrow.StructConverter;
25+
import net.snowflake.client.core.arrow.VarCharConverter;
2426
import net.snowflake.client.core.json.Converters;
2527
import net.snowflake.client.jdbc.ArrowResultChunk;
2628
import net.snowflake.client.jdbc.ArrowResultChunk.ArrowChunkIterator;
@@ -38,6 +40,7 @@
3840
import net.snowflake.common.core.SnowflakeDateTimeFormat;
3941
import net.snowflake.common.core.SqlState;
4042
import org.apache.arrow.memory.RootAllocator;
43+
import org.apache.arrow.vector.util.JsonStringHashMap;
4144

4245
/** Arrow result set implementation */
4346
public class SFArrowResultSet extends SFBaseResultSet implements DataConversionContext {
@@ -506,25 +509,39 @@ public Object getObject(int columnIndex) throws SFException {
506509
converter.setUseSessionTimezone(useSessionTimezone);
507510
converter.setSessionTimeZone(timeZone);
508511
Object obj = converter.toObject(index);
509-
return handleObjectType(columnIndex, obj);
512+
int type = resultSetMetaData.getColumnType(columnIndex);
513+
if (type == Types.STRUCT &&
514+
Boolean.parseBoolean(System.getProperty(STRUCTURED_TYPE_ENABLED_PROPERTY_NAME))) {
515+
if (converter instanceof VarCharConverter) {
516+
return createJsonSqlInput(columnIndex, obj);
517+
} else if (converter instanceof StructConverter) {
518+
return createArrowSqlInput(columnIndex, (JsonStringHashMap<String, Object>) obj);
519+
}
520+
}
521+
return obj;
510522
}
511523

512-
private Object handleObjectType(int columnIndex, Object obj) throws SFException {
513-
int columnType = resultSetMetaData.getColumnType(columnIndex);
514-
if (columnType == Types.STRUCT
515-
&& Boolean.valueOf(System.getProperty(STRUCTURED_TYPE_ENABLED_PROPERTY_NAME))) {
516-
try {
517-
JsonNode jsonNode = OBJECT_MAPPER.readTree((String) obj);
518-
return new JsonSqlInput(
519-
jsonNode,
524+
private Object createJsonSqlInput(int columnIndex, Object obj) throws SFException {
525+
try {
526+
JsonNode jsonNode = OBJECT_MAPPER.readTree((String) obj);
527+
return new JsonSqlInput(
528+
jsonNode,
529+
session,
530+
jsonConverters,
531+
Arrays.asList(resultSetMetaData.getColumnMetadata().get(columnIndex - 1).getFields())
532+
);
533+
} catch (JsonProcessingException e) {
534+
throw new SFException(e, ErrorCode.INVALID_STRUCT_DATA);
535+
}
536+
}
537+
538+
private Object createArrowSqlInput(int columnIndex, JsonStringHashMap<String, Object> input) {
539+
return new ArrowSqlInput(
540+
input,
520541
session,
521542
jsonConverters,
522-
Arrays.asList(resultSetMetaData.getColumnMetadata().get(columnIndex - 1).getFields()));
523-
} catch (JsonProcessingException e) {
524-
throw new SFException(e, ErrorCode.INVALID_STRUCT_DATA);
525-
}
526-
}
527-
return obj;
543+
Arrays.asList(resultSetMetaData.getColumnMetadata().get(columnIndex - 1).getFields())
544+
);
528545
}
529546

530547
@Override

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

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package net.snowflake.client.core;
2+
3+
import net.snowflake.client.jdbc.SnowflakeUtil;
4+
import net.snowflake.common.core.SnowflakeDateTimeFormat;
5+
6+
import java.sql.Timestamp;
7+
import java.sql.Types;
8+
9+
public class SqlInputTimestampUtil {
10+
11+
public static SnowflakeDateTimeFormat extractDateTimeFormat(SFBaseSession session, String format) {
12+
return SnowflakeDateTimeFormat.fromSqlFormat(
13+
(String) session.getCommonParameters().get(format));
14+
}
15+
16+
public static Timestamp getTimestampFromType(int columnSubType, String value, SFBaseSession session) {
17+
if (columnSubType == SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_LTZ) {
18+
return getTimestampFromFormat("TIMESTAMP_LTZ_OUTPUT_FORMAT", value, session);
19+
} else if (columnSubType == SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_NTZ
20+
|| columnSubType == Types.TIMESTAMP) {
21+
return getTimestampFromFormat("TIMESTAMP_NTZ_OUTPUT_FORMAT", value, session);
22+
} else if (columnSubType == SnowflakeUtil.EXTRA_TYPES_TIMESTAMP_TZ) {
23+
return getTimestampFromFormat("TIMESTAMP_TZ_OUTPUT_FORMAT", value, session);
24+
} else {
25+
return null;
26+
}
27+
}
28+
29+
private static Timestamp getTimestampFromFormat(String format, String value, SFBaseSession session) {
30+
String rawFormat = (String) session.getCommonParameters().get(format);
31+
if (rawFormat == null || rawFormat.isEmpty()) {
32+
rawFormat = (String) session.getCommonParameters().get("TIMESTAMP_OUTPUT_FORMAT");
33+
}
34+
SnowflakeDateTimeFormat formatter = SnowflakeDateTimeFormat.fromSqlFormat(rawFormat);
35+
return formatter.parse(value).getTimestamp();
36+
}
37+
38+
}

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

+29
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package net.snowflake.client.core.arrow;
2+
3+
import net.snowflake.client.core.DataConversionContext;
4+
import net.snowflake.client.core.SFException;
5+
import net.snowflake.client.core.SnowflakeJdbcInternalApi;
6+
import net.snowflake.client.jdbc.SnowflakeType;
7+
import org.apache.arrow.vector.ValueVector;
8+
import org.apache.arrow.vector.complex.StructVector;
9+
10+
@SnowflakeJdbcInternalApi
11+
public class StructConverter extends AbstractArrowVectorConverter {
12+
13+
private final StructVector structVector;
14+
15+
public StructConverter(ValueVector vector, int columnIndex, DataConversionContext context) {
16+
super(SnowflakeType.OBJECT.name(), vector, columnIndex, context);
17+
structVector = (StructVector) vector;
18+
}
19+
20+
@Override
21+
public Object toObject(int index) throws SFException {
22+
return structVector.getObject(index);
23+
}
24+
25+
@Override
26+
public String toString(int index) throws SFException {
27+
return structVector.getObject(index).toString();
28+
}
29+
}

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

+43
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package net.snowflake.client.core.arrow;
2+
3+
import net.snowflake.client.core.SFBaseSession;
4+
import net.snowflake.client.core.SFException;
5+
6+
import java.sql.Timestamp;
7+
import java.util.TimeZone;
8+
9+
public class StructuredTypeDateTimeConverter {
10+
11+
private final TimeZone sessionTimeZone;
12+
private final SFBaseSession session;
13+
private final long resultVersion;
14+
private final boolean honorClientTZForTimestampNTZ;
15+
private final boolean treatNTZAsUTC;
16+
private final boolean useSessionTimezone;
17+
private final boolean formatDateWithTimeZone;
18+
19+
public StructuredTypeDateTimeConverter(
20+
TimeZone sessionTimeZone,
21+
SFBaseSession session,
22+
long resultVersion,
23+
boolean honorClientTZForTimestampNTZ,
24+
boolean treatNTZAsUTC,
25+
boolean useSessionTimezone,
26+
boolean formatDateWithTimeZone) {
27+
28+
this.sessionTimeZone = sessionTimeZone;
29+
this.session = session;
30+
this.resultVersion = resultVersion;
31+
this.honorClientTZForTimestampNTZ = honorClientTZForTimestampNTZ;
32+
this.treatNTZAsUTC = treatNTZAsUTC;
33+
this.useSessionTimezone = useSessionTimezone;
34+
this.formatDateWithTimeZone = formatDateWithTimeZone;
35+
}
36+
37+
public Timestamp getTimestamp(Object obj, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException {
38+
if (obj == null) {
39+
return null;
40+
}
41+
return null;
42+
}
43+
}

‎src/main/java/net/snowflake/client/core/json/Converters.java

+11
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import java.util.TimeZone;
44
import net.snowflake.client.core.SFBaseSession;
5+
import net.snowflake.client.core.arrow.StructuredTypeDateTimeConverter;
56
import net.snowflake.common.core.SFBinaryFormat;
67
import net.snowflake.common.core.SnowflakeDateTimeFormat;
78

@@ -11,6 +12,7 @@ public class Converters {
1112
private final DateTimeConverter dateTimeConverter;
1213
private final BytesConverter bytesConverter;
1314
private final StringConverter stringConverter;
15+
private final StructuredTypeDateTimeConverter structuredTypeDateTimeConverter;
1416

1517
public Converters(
1618
TimeZone sessionTimeZone,
@@ -50,6 +52,15 @@ public Converters(
5052
resultVersion,
5153
session,
5254
this);
55+
structuredTypeDateTimeConverter =
56+
new StructuredTypeDateTimeConverter(
57+
sessionTimeZone,
58+
session,
59+
resultVersion,
60+
honorClientTZForTimestampNTZ,
61+
treatNTZAsUTC,
62+
useSessionTimezone,
63+
formatDateWithTimeZone);
5364
}
5465

5566
public BooleanConverter getBooleanConverter() {

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

+5-1
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import net.snowflake.client.core.arrow.IntToTimeConverter;
3030
import net.snowflake.client.core.arrow.SmallIntToFixedConverter;
3131
import net.snowflake.client.core.arrow.SmallIntToScaledFixedConverter;
32+
import net.snowflake.client.core.arrow.StructConverter;
3233
import net.snowflake.client.core.arrow.ThreeFieldStructToTimestampTZConverter;
3334
import net.snowflake.client.core.arrow.TinyIntToFixedConverter;
3435
import net.snowflake.client.core.arrow.TinyIntToScaledFixedConverter;
@@ -201,11 +202,14 @@ private static List<ArrowVectorConverter> initConverters(
201202
case ARRAY:
202203
case CHAR:
203204
case TEXT:
204-
case OBJECT:
205205
case VARIANT:
206206
converters.add(new VarCharConverter(vector, i, context));
207207
break;
208208

209+
case OBJECT:
210+
converters.add(new StructConverter(vector, i , context));
211+
break;
212+
209213
case BINARY:
210214
converters.add(new VarBinaryToBinaryConverter(vector, i, context));
211215
break;

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

+12-5
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.io.PrintWriter;
1616
import java.io.StringWriter;
1717
import java.lang.reflect.Field;
18+
import java.sql.SQLException;
1819
import java.sql.Time;
1920
import java.sql.Types;
2021
import java.time.Instant;
@@ -32,13 +33,11 @@
3233
import java.util.concurrent.ThreadFactory;
3334
import java.util.concurrent.ThreadPoolExecutor;
3435
import java.util.concurrent.TimeUnit;
35-
import net.snowflake.client.core.HttpClientSettingsKey;
36-
import net.snowflake.client.core.OCSPMode;
37-
import net.snowflake.client.core.SFBaseSession;
38-
import net.snowflake.client.core.SFSession;
39-
import net.snowflake.client.core.SFSessionProperty;
36+
37+
import net.snowflake.client.core.*;
4038
import net.snowflake.client.log.SFLogger;
4139
import net.snowflake.client.log.SFLoggerFactory;
40+
import net.snowflake.client.util.ThrowingCallable;
4241
import net.snowflake.common.core.SqlState;
4342
import net.snowflake.common.util.ClassUtil;
4443
import net.snowflake.common.util.FixedViewColumn;
@@ -755,4 +754,12 @@ public static Time getTimeInSessionTimezone(Long time, int nanos) {
755754
ts.setTime(c.getTimeInMillis());
756755
return ts;
757756
}
757+
758+
public static <T> T mapExceptions(ThrowingCallable<T, SFException> action) throws SQLException {
759+
try {
760+
return action.call();
761+
} catch (SFException e) {
762+
throw new SQLException(e);
763+
}
764+
}
758765
}

‎src/main/java/net/snowflake/client/util/ThrowingBiFunction.java

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package net.snowflake.client.util;
2+
3+
import net.snowflake.client.core.SnowflakeJdbcInternalApi;
4+
5+
@SnowflakeJdbcInternalApi
6+
@FunctionalInterface
7+
public interface ThrowingBiFunction<A, B, R, T extends Throwable> {
8+
R apply(A a, B b) throws T;
9+
}

‎src/test/java/net/snowflake/client/AbstractDriverIT.java

-2
Original file line numberDiff line numberDiff line change
@@ -324,8 +324,6 @@ public static Connection getConnection(
324324

325325
properties.put("internal", Boolean.TRUE.toString()); // TODO: do we need this?
326326

327-
properties.put("insecureMode", false); // use OCSP for all tests.
328-
329327
if (injectSocketTimeout > 0) {
330328
properties.put("injectSocketTimeout", String.valueOf(injectSocketTimeout));
331329
}

‎src/test/java/net/snowflake/client/jdbc/ArrowResultSetStructuredTypesLatestIT.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,6 @@
55

66
public class ArrowResultSetStructuredTypesLatestIT extends ResultSetStructuredTypesLatestIT {
77
public ArrowResultSetStructuredTypesLatestIT() {
8-
super("ARROW");
8+
super(ResultSetFormatType.ARROW_WITH_JSON_STRUCTURED_TYPES);
99
}
1010
}

‎src/test/java/net/snowflake/client/jdbc/NativeArrowResultSetStructuredTypesLatestIT.java

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
/*
2+
* Copyright (c) 2012-2024 Snowflake Computing Inc. All right reserved.
3+
*/
4+
package net.snowflake.client.jdbc;
5+
6+
public class NativeArrowResultSetStructuredTypesLatestIT extends ResultSetStructuredTypesLatestIT {
7+
public NativeArrowResultSetStructuredTypesLatestIT() {
8+
super(ResultSetFormatType.NATIVE_ARROW);
9+
}
10+
}

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

+22-6
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,13 @@
2727

2828
@Category(TestCategoryStructuredType.class)
2929
public class ResultSetStructuredTypesLatestIT extends BaseJDBCTest {
30-
private final String queryResultFormat;
30+
private final ResultSetFormatType queryResultFormat;
3131

3232
public ResultSetStructuredTypesLatestIT() {
33-
this("JSON");
33+
this(ResultSetFormatType.JSON);
3434
}
3535

36-
protected ResultSetStructuredTypesLatestIT(String queryResultFormat) {
36+
protected ResultSetStructuredTypesLatestIT(ResultSetFormatType queryResultFormat) {
3737
this.queryResultFormat = queryResultFormat;
3838
}
3939

@@ -42,7 +42,11 @@ public Connection init() throws SQLException {
4242
try (Statement stmt = conn.createStatement()) {
4343
stmt.execute("alter session set ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE = true");
4444
stmt.execute("alter session set IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE = true");
45-
stmt.execute("alter session set jdbc_query_result_format = '" + queryResultFormat + "'");
45+
stmt.execute("alter session set jdbc_query_result_format = '" + queryResultFormat.sessionParameterTypeValue + "'");
46+
if (queryResultFormat == ResultSetFormatType.NATIVE_ARROW) {
47+
stmt.execute("alter session set ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT = true");
48+
stmt.execute("alter session set FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT = true");
49+
}
4650
}
4751
return conn;
4852
}
@@ -70,8 +74,8 @@ private void testMapJson(boolean registerFactory) throws SQLException {
7074
}
7175
try (Connection connection = init();
7276
Statement statement = connection.createStatement();
73-
ResultSet resultSet =
74-
statement.executeQuery("select {'string':'a'}::OBJECT(string VARCHAR)"); ) {
77+
ResultSet resultSet =
78+
statement.executeQuery("select {'string':'a'}::OBJECT(string VARCHAR)"); ) {
7579
resultSet.next();
7680
SimpleClass object = resultSet.getObject(1, SimpleClass.class);
7781
assertEquals("a", object.getString());
@@ -164,4 +168,16 @@ public void testMapStructsFromChunks() throws SQLException {
164168
}
165169
}
166170
}
171+
172+
enum ResultSetFormatType {
173+
JSON("JSON"),
174+
ARROW_WITH_JSON_STRUCTURED_TYPES("ARROW"),
175+
NATIVE_ARROW("ARROW");
176+
177+
public final String sessionParameterTypeValue;
178+
179+
ResultSetFormatType(String sessionParameterTypeValue) {
180+
this.sessionParameterTypeValue = sessionParameterTypeValue;
181+
}
182+
}
167183
}

0 commit comments

Comments
 (0)
Please sign in to comment.