Skip to content

Commit 63b654d

Browse files
sfc-gh-dheymansfc-gh-dheyman
authored
SNOW-969794 Arrow native structured type base support (#1649)
1 parent a36390a commit 63b654d

30 files changed

+943
-321
lines changed

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

+215
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
/*
2+
* Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
3+
*/
4+
5+
package net.snowflake.client.core;
6+
7+
import static net.snowflake.client.jdbc.SnowflakeUtil.mapExceptions;
8+
9+
import java.math.BigDecimal;
10+
import java.sql.Date;
11+
import java.sql.SQLData;
12+
import java.sql.SQLException;
13+
import java.sql.Time;
14+
import java.sql.Timestamp;
15+
import java.util.Iterator;
16+
import java.util.List;
17+
import java.util.TimeZone;
18+
import net.snowflake.client.core.json.Converters;
19+
import net.snowflake.client.core.structs.SQLDataCreationHelper;
20+
import net.snowflake.client.jdbc.FieldMetadata;
21+
import net.snowflake.client.util.ThrowingBiFunction;
22+
import org.apache.arrow.vector.util.JsonStringHashMap;
23+
24+
@SnowflakeJdbcInternalApi
25+
public class ArrowSqlInput extends BaseSqlInput {
26+
27+
private final Iterator<Object> structuredTypeFields;
28+
private int currentIndex = 0;
29+
30+
public ArrowSqlInput(
31+
JsonStringHashMap<String, Object> input,
32+
SFBaseSession session,
33+
Converters converters,
34+
List<FieldMetadata> fields) {
35+
super(session, converters, fields);
36+
this.structuredTypeFields = input.values().iterator();
37+
}
38+
39+
@Override
40+
public String readString() throws SQLException {
41+
return withNextValue(
42+
((value, fieldMetadata) -> {
43+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
44+
int columnSubType = fieldMetadata.getType();
45+
int scale = fieldMetadata.getScale();
46+
return mapExceptions(
47+
() ->
48+
converters
49+
.getStringConverter()
50+
.getString(value, columnType, columnSubType, scale));
51+
}));
52+
}
53+
54+
@Override
55+
public boolean readBoolean() throws SQLException {
56+
return withNextValue(
57+
(value, fieldMetadata) -> {
58+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
59+
return mapExceptions(
60+
() -> converters.getBooleanConverter().getBoolean(value, columnType));
61+
});
62+
}
63+
64+
@Override
65+
public byte readByte() throws SQLException {
66+
return withNextValue(
67+
(value, fieldMetadata) ->
68+
mapExceptions(() -> converters.getNumberConverter().getByte(value)));
69+
}
70+
71+
@Override
72+
public short readShort() throws SQLException {
73+
return withNextValue(
74+
(value, fieldMetadata) -> {
75+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
76+
return mapExceptions(() -> converters.getNumberConverter().getShort(value, columnType));
77+
});
78+
}
79+
80+
@Override
81+
public int readInt() throws SQLException {
82+
return withNextValue(
83+
(value, fieldMetadata) -> {
84+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
85+
return mapExceptions(() -> converters.getNumberConverter().getInt(value, columnType));
86+
});
87+
}
88+
89+
@Override
90+
public long readLong() throws SQLException {
91+
return withNextValue(
92+
(value, fieldMetadata) -> {
93+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
94+
return mapExceptions(() -> converters.getNumberConverter().getLong(value, columnType));
95+
});
96+
}
97+
98+
@Override
99+
public float readFloat() throws SQLException {
100+
return withNextValue(
101+
(value, fieldMetadata) -> {
102+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
103+
return mapExceptions(() -> converters.getNumberConverter().getFloat(value, columnType));
104+
});
105+
}
106+
107+
@Override
108+
public double readDouble() throws SQLException {
109+
return withNextValue(
110+
(value, fieldMetadata) -> {
111+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
112+
return mapExceptions(() -> converters.getNumberConverter().getDouble(value, columnType));
113+
});
114+
}
115+
116+
@Override
117+
public BigDecimal readBigDecimal() throws SQLException {
118+
return withNextValue(
119+
(value, fieldMetadata) -> {
120+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
121+
return mapExceptions(
122+
() -> converters.getNumberConverter().getBigDecimal(value, columnType));
123+
});
124+
}
125+
126+
@Override
127+
public byte[] readBytes() throws SQLException {
128+
return withNextValue(
129+
(value, fieldMetadata) -> {
130+
int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
131+
int columnSubType = fieldMetadata.getType();
132+
int scale = fieldMetadata.getScale();
133+
return mapExceptions(
134+
() ->
135+
converters.getBytesConverter().getBytes(value, columnType, columnSubType, scale));
136+
});
137+
}
138+
139+
@Override
140+
public Date readDate() throws SQLException {
141+
return withNextValue(
142+
(value, fieldMetadata) ->
143+
mapExceptions(
144+
() ->
145+
converters
146+
.getStructuredTypeDateTimeConverter()
147+
.getDate((int) value, TimeZone.getDefault())));
148+
}
149+
150+
@Override
151+
public Time readTime() throws SQLException {
152+
return withNextValue(
153+
(value, fieldMetadata) ->
154+
mapExceptions(
155+
() -> {
156+
int scale = fieldMetadata.getScale();
157+
return converters
158+
.getStructuredTypeDateTimeConverter()
159+
.getTime((long) value, scale);
160+
}));
161+
}
162+
163+
@Override
164+
public Timestamp readTimestamp(TimeZone tz) throws SQLException {
165+
return withNextValue(
166+
(value, fieldMetadata) -> {
167+
if (value == null) {
168+
return null;
169+
}
170+
int scale = fieldMetadata.getScale();
171+
return mapExceptions(
172+
() ->
173+
converters
174+
.getStructuredTypeDateTimeConverter()
175+
.getTimestamp(
176+
(JsonStringHashMap<String, Object>) value,
177+
fieldMetadata.getBase(),
178+
tz,
179+
scale));
180+
});
181+
}
182+
183+
@Override
184+
public Object readObject() throws SQLException {
185+
return withNextValue(
186+
(value, fieldMetadata) -> {
187+
if (!(value instanceof JsonStringHashMap)) {
188+
throw new SQLException(
189+
"Invalid value passed to 'readObject()', expected Map; got: " + value.getClass());
190+
}
191+
return value;
192+
});
193+
}
194+
195+
@Override
196+
public <T> T readObject(Class<T> type) throws SQLException {
197+
return withNextValue(
198+
(value, fieldMetadata) -> {
199+
SQLData instance = (SQLData) SQLDataCreationHelper.create(type);
200+
instance.readSQL(
201+
new ArrowSqlInput(
202+
(JsonStringHashMap<String, Object>) value,
203+
session,
204+
converters,
205+
fieldMetadata.getFields()),
206+
null);
207+
return (T) instance;
208+
});
209+
}
210+
211+
private <T> T withNextValue(ThrowingBiFunction<Object, FieldMetadata, T, SQLException> action)
212+
throws SQLException {
213+
return action.apply(structuredTypeFields.next(), fields.get(currentIndex++));
214+
}
215+
}

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

+106
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
/*
2+
* Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
3+
*/
4+
5+
package net.snowflake.client.core;
6+
7+
import java.io.InputStream;
8+
import java.io.Reader;
9+
import java.net.URL;
10+
import java.sql.Array;
11+
import java.sql.Blob;
12+
import java.sql.Clob;
13+
import java.sql.NClob;
14+
import java.sql.Ref;
15+
import java.sql.RowId;
16+
import java.sql.SQLException;
17+
import java.sql.SQLXML;
18+
import java.sql.Timestamp;
19+
import java.util.List;
20+
import net.snowflake.client.core.json.Converters;
21+
import net.snowflake.client.jdbc.FieldMetadata;
22+
import net.snowflake.client.jdbc.SnowflakeLoggedFeatureNotSupportedException;
23+
24+
@SnowflakeJdbcInternalApi
25+
public abstract class BaseSqlInput implements SFSqlInput {
26+
27+
protected final SFBaseSession session;
28+
protected final Converters converters;
29+
protected final List<FieldMetadata> fields;
30+
31+
protected BaseSqlInput(SFBaseSession session, Converters converters, List<FieldMetadata> fields) {
32+
this.session = session;
33+
this.converters = converters;
34+
this.fields = fields;
35+
}
36+
37+
@Override
38+
public Timestamp readTimestamp() throws SQLException {
39+
return readTimestamp(null);
40+
}
41+
42+
@Override
43+
public Reader readCharacterStream() throws SQLException {
44+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readCharacterStream");
45+
}
46+
47+
@Override
48+
public InputStream readAsciiStream() throws SQLException {
49+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readAsciiStream");
50+
}
51+
52+
@Override
53+
public InputStream readBinaryStream() throws SQLException {
54+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readBinaryStream");
55+
}
56+
57+
@Override
58+
public Ref readRef() throws SQLException {
59+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readRef");
60+
}
61+
62+
@Override
63+
public Blob readBlob() throws SQLException {
64+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readBlob");
65+
}
66+
67+
@Override
68+
public Clob readClob() throws SQLException {
69+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readClob");
70+
}
71+
72+
@Override
73+
public Array readArray() throws SQLException {
74+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readArray");
75+
}
76+
77+
@Override
78+
public boolean wasNull() throws SQLException {
79+
return false; // nulls are not allowed in structure types
80+
}
81+
82+
@Override
83+
public URL readURL() throws SQLException {
84+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readCharacterStream");
85+
}
86+
87+
@Override
88+
public NClob readNClob() throws SQLException {
89+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readNClob");
90+
}
91+
92+
@Override
93+
public String readNString() throws SQLException {
94+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readNString");
95+
}
96+
97+
@Override
98+
public SQLXML readSQLXML() throws SQLException {
99+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readSQLXML");
100+
}
101+
102+
@Override
103+
public RowId readRowId() throws SQLException {
104+
throw new SnowflakeLoggedFeatureNotSupportedException(session, "readRowId");
105+
}
106+
}

0 commit comments

Comments
 (0)