Skip to content

Commit f279079

Browse files
sfc-gh-ext-simba-jfsfc-gh-ext-simba-jf
authored
SNOW-1170182: Return decimal as int in ARROW results configured by flag (#1666)
1 parent 75c57f2 commit f279079

9 files changed

+144
-5
lines changed

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

+10
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,8 @@ public abstract class SFBaseSession {
142142

143143
private Map<String, Object> commonParameters;
144144

145+
private boolean isJdbcArrowTreatDecimalAsInt = true;
146+
145147
protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
146148
this.sfConnectionHandler = sfConnectionHandler;
147149
}
@@ -270,6 +272,14 @@ public void setJdbcTreatDecimalAsInt(boolean jdbcTreatDecimalAsInt) {
270272
isJdbcTreatDecimalAsInt = jdbcTreatDecimalAsInt;
271273
}
272274

275+
public boolean isJdbcArrowTreatDecimalAsInt() {
276+
return isJdbcArrowTreatDecimalAsInt;
277+
}
278+
279+
public void setJdbcArrowTreatDecimalAsInt(boolean jdbcArrowTreatDecimalAsInt) {
280+
isJdbcArrowTreatDecimalAsInt = jdbcArrowTreatDecimalAsInt;
281+
}
282+
273283
public String getServerUrl() {
274284
if (connectionPropertiesMap.containsKey(SFSessionProperty.SERVER_URL)) {
275285
return (String) connectionPropertiesMap.get(SFSessionProperty.SERVER_URL);

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

+6
Original file line numberDiff line numberDiff line change
@@ -475,6 +475,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
475475
}
476476
break;
477477

478+
case JDBC_ARROW_TREAT_DECIMAL_AS_INT:
479+
if (propertyValue != null) {
480+
setJdbcArrowTreatDecimalAsInt(getBooleanValue(propertyValue));
481+
}
482+
break;
483+
478484
default:
479485
break;
480486
}

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,9 @@ public enum SFSessionProperty {
8080

8181
ENABLE_PATTERN_SEARCH("enablePatternSearch", false, Boolean.class),
8282

83-
DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class);
83+
DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class),
84+
85+
JDBC_ARROW_TREAT_DECIMAL_AS_INT("JDBC_ARROW_TREAT_DECIMAL_AS_INT", false, Boolean.class);
8486

8587
// property key in string
8688
private String propertyKey;

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

+10
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import java.sql.Timestamp;
1010
import java.util.TimeZone;
1111
import net.snowflake.client.core.DataConversionContext;
12+
import net.snowflake.client.core.SFBaseSession;
1213
import net.snowflake.client.core.SFException;
1314
import net.snowflake.client.jdbc.ErrorCode;
1415
import net.snowflake.client.jdbc.SnowflakeUtil;
@@ -146,6 +147,15 @@ public BigDecimal toBigDecimal(int index) throws SFException {
146147
ErrorCode.INVALID_VALUE_CONVERT, logicalTypeStr, SnowflakeUtil.BIG_DECIMAL_STR, "");
147148
}
148149

150+
public boolean shouldTreatDecimalAsInt(SFBaseSession session) {
151+
if (session != null) {
152+
if (!session.isJdbcArrowTreatDecimalAsInt() && !session.isJdbcTreatDecimalAsInt()) {
153+
return false;
154+
}
155+
}
156+
return true;
157+
}
158+
149159
@Override
150160
public void setTreatNTZAsUTC(boolean isUTC) {
151161
this.treatNTZasUTC = isUTC;

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -137,9 +137,10 @@ public BigDecimal toBigDecimal(int index) {
137137
public Object toObject(int index) throws SFException {
138138
if (bigIntVector.isNull(index)) {
139139
return null;
140-
} else {
141-
return getLong(index);
140+
} else if (!shouldTreatDecimalAsInt(context.getSession())) {
141+
return BigDecimal.valueOf(getLong(index), sfScale);
142142
}
143+
return getLong(index);
143144
}
144145

145146
@Override

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

+6-1
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,12 @@ public BigDecimal toBigDecimal(int index) throws SFException {
105105

106106
@Override
107107
public Object toObject(int index) throws SFException {
108-
return isNull(index) ? null : (long) getInt(index);
108+
if (isNull(index)) {
109+
return null;
110+
} else if (!shouldTreatDecimalAsInt(context.getSession())) {
111+
return BigDecimal.valueOf((long) getInt(index), sfScale);
112+
}
113+
return (long) getInt(index);
109114
}
110115

111116
@Override

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

+2
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,8 @@ public double toDouble(int index) throws SFException {
9999
public Object toObject(int index) throws SFException {
100100
if (isNull(index)) {
101101
return null;
102+
} else if (!shouldTreatDecimalAsInt(context.getSession())) {
103+
return BigDecimal.valueOf((long) getShort(index), sfScale);
102104
}
103105
return (long) getShort(index);
104106
}

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

+6-1
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,12 @@ public BigDecimal toBigDecimal(int index) throws SFException {
9090

9191
@Override
9292
public Object toObject(int index) throws SFException {
93-
return isNull(index) ? null : (long) toByte(index);
93+
if (isNull(index)) {
94+
return null;
95+
} else if (!shouldTreatDecimalAsInt(context.getSession())) {
96+
return BigDecimal.valueOf((long) getByte(index), sfScale);
97+
}
98+
return (long) toByte(index);
9499
}
95100

96101
@Override

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

+98
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,14 @@ public ResultSetLatestIT() {
7373
super(queryResultFormat);
7474
}
7575

76+
private String createTableSql =
77+
"Create or replace table get_object_for_numeric_types (c1 INT, c2 BIGINT, c3 SMALLINT, c4 TINYINT) ";
78+
private String insertStmt =
79+
"Insert into get_object_for_numeric_types (c1, c2, c3, c4) values (1000000000, 2000000000000000000000000, 3, 4)";
80+
private String selectQuery = "Select * from get_object_for_numeric_types";
81+
private String setJdbcTreatDecimalAsIntFalse =
82+
"alter session set JDBC_TREAT_DECIMAL_AS_INT = false";
83+
7684
/**
7785
* Test that when closing of results is interrupted by Thread.Interrupt(), the memory is released
7886
* safely before driver execution ends.
@@ -971,4 +979,94 @@ public void testLargeStringRetrieval() throws SQLException {
971979
fail("executeQuery should not fail");
972980
}
973981
}
982+
983+
private static void assertAllColumnsAreLongButBigIntIsBigDecimal(ResultSet rs)
984+
throws SQLException {
985+
while (rs.next()) {
986+
assertEquals(java.lang.Long.class, rs.getObject(1).getClass());
987+
assertEquals(java.math.BigDecimal.class, rs.getObject(2).getClass());
988+
assertEquals(java.lang.Long.class, rs.getObject(3).getClass());
989+
assertEquals(java.lang.Long.class, rs.getObject(4).getClass());
990+
}
991+
}
992+
993+
private static void assertAllColumnsAreBigDecimal(ResultSet rs) throws SQLException {
994+
while (rs.next()) {
995+
assertEquals(java.math.BigDecimal.class, rs.getObject(1).getClass());
996+
assertEquals(java.math.BigDecimal.class, rs.getObject(2).getClass());
997+
assertEquals(java.math.BigDecimal.class, rs.getObject(3).getClass());
998+
assertEquals(java.math.BigDecimal.class, rs.getObject(4).getClass());
999+
}
1000+
}
1001+
1002+
// Test setting new connection property JDBC_ARROW_TREAT_DECIMAL_AS_INT=false. Connection property
1003+
// introduced after version 3.15.0.
1004+
@Test
1005+
public void testGetObjectForArrowResultFormatJDBCArrowDecimalAsIntFalse() throws SQLException {
1006+
Properties properties = new Properties();
1007+
properties.put("JDBC_ARROW_TREAT_DECIMAL_AS_INT", false);
1008+
try (Connection con = getConnection(properties);
1009+
Statement stmt = con.createStatement()) {
1010+
stmt.execute("alter session set jdbc_query_result_format = 'ARROW'");
1011+
stmt.execute(createTableSql);
1012+
stmt.execute(insertStmt);
1013+
1014+
// Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=false and JDBC_TREAT_DECIMAL_AS_INT=true
1015+
try (ResultSet rs = stmt.executeQuery(selectQuery)) {
1016+
assertAllColumnsAreLongButBigIntIsBigDecimal(rs);
1017+
}
1018+
1019+
// Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=false and JDBC_TREAT_DECIMAL_AS_INT=false
1020+
stmt.execute(setJdbcTreatDecimalAsIntFalse);
1021+
try (ResultSet rs = stmt.executeQuery(selectQuery)) {
1022+
assertAllColumnsAreBigDecimal(rs);
1023+
}
1024+
}
1025+
}
1026+
1027+
// Test default setting of new connection property JDBC_ARROW_TREAT_DECIMAL_AS_INT=true.
1028+
// Connection property introduced after version 3.15.0.
1029+
@Test
1030+
public void testGetObjectForArrowResultFormatJDBCArrowDecimalAsIntTrue() throws SQLException {
1031+
try (Connection con = BaseJDBCTest.getConnection();
1032+
Statement stmt = con.createStatement()) {
1033+
stmt.execute("alter session set jdbc_query_result_format = 'ARROW'");
1034+
stmt.execute(createTableSql);
1035+
stmt.execute(insertStmt);
1036+
1037+
// Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=true and JDBC_TREAT_DECIMAL_AS_INT=true
1038+
try (ResultSet rs = stmt.executeQuery(selectQuery)) {
1039+
assertAllColumnsAreLongButBigIntIsBigDecimal(rs);
1040+
}
1041+
1042+
// Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=true and JDBC_TREAT_DECIMAL_AS_INT=false
1043+
stmt.execute(setJdbcTreatDecimalAsIntFalse);
1044+
try (ResultSet rs = stmt.executeQuery(selectQuery)) {
1045+
assertAllColumnsAreLongButBigIntIsBigDecimal(rs);
1046+
}
1047+
}
1048+
}
1049+
1050+
// Test getObject for numeric types when JDBC_TREAT_DECIMAL_AS_INT is set and using JSON result
1051+
// format.
1052+
@Test
1053+
public void testGetObjectForJSONResultFormatUsingJDBCDecimalAsInt() throws SQLException {
1054+
try (Connection con = BaseJDBCTest.getConnection();
1055+
Statement stmt = con.createStatement()) {
1056+
stmt.execute("alter session set jdbc_query_result_format = 'JSON'");
1057+
stmt.execute(createTableSql);
1058+
stmt.execute(insertStmt);
1059+
1060+
// Test with JDBC_TREAT_DECIMAL_AS_INT=true (default value)
1061+
try (ResultSet rs = stmt.executeQuery(selectQuery)) {
1062+
assertAllColumnsAreLongButBigIntIsBigDecimal(rs);
1063+
}
1064+
1065+
// Test with JDBC_TREAT_DECIMAL_AS_INT=false
1066+
stmt.execute(setJdbcTreatDecimalAsIntFalse);
1067+
try (ResultSet rs = stmt.executeQuery(selectQuery)) {
1068+
assertAllColumnsAreBigDecimal(rs);
1069+
}
1070+
}
1071+
}
9741072
}

0 commit comments

Comments
 (0)