From f279079016c57c6c06e0732bde096808eafbebde Mon Sep 17 00:00:00 2001 From: Jelena Furundzic <141762304+sfc-gh-ext-simba-jf@users.noreply.github.com> Date: Fri, 15 Mar 2024 03:35:44 -0700 Subject: [PATCH] SNOW-1170182: Return decimal as int in ARROW results configured by flag (#1666) --- .../snowflake/client/core/SFBaseSession.java | 10 ++ .../net/snowflake/client/core/SFSession.java | 6 ++ .../client/core/SFSessionProperty.java | 4 +- .../arrow/AbstractArrowVectorConverter.java | 10 ++ .../core/arrow/BigIntToFixedConverter.java | 5 +- .../core/arrow/IntToFixedConverter.java | 7 +- .../core/arrow/SmallIntToFixedConverter.java | 2 + .../core/arrow/TinyIntToFixedConverter.java | 7 +- .../client/jdbc/ResultSetLatestIT.java | 98 +++++++++++++++++++ 9 files changed, 144 insertions(+), 5 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java index 5c0e20fcc..b0da3e9b4 100644 --- a/src/main/java/net/snowflake/client/core/SFBaseSession.java +++ b/src/main/java/net/snowflake/client/core/SFBaseSession.java @@ -142,6 +142,8 @@ public abstract class SFBaseSession { private Map commonParameters; + private boolean isJdbcArrowTreatDecimalAsInt = true; + protected SFBaseSession(SFConnectionHandler sfConnectionHandler) { this.sfConnectionHandler = sfConnectionHandler; } @@ -270,6 +272,14 @@ public void setJdbcTreatDecimalAsInt(boolean jdbcTreatDecimalAsInt) { isJdbcTreatDecimalAsInt = jdbcTreatDecimalAsInt; } + public boolean isJdbcArrowTreatDecimalAsInt() { + return isJdbcArrowTreatDecimalAsInt; + } + + public void setJdbcArrowTreatDecimalAsInt(boolean jdbcArrowTreatDecimalAsInt) { + isJdbcArrowTreatDecimalAsInt = jdbcArrowTreatDecimalAsInt; + } + public String getServerUrl() { if (connectionPropertiesMap.containsKey(SFSessionProperty.SERVER_URL)) { return (String) connectionPropertiesMap.get(SFSessionProperty.SERVER_URL); diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java index 75d5e9600..d7ee69a07 100644 --- a/src/main/java/net/snowflake/client/core/SFSession.java +++ b/src/main/java/net/snowflake/client/core/SFSession.java @@ -475,6 +475,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro } break; + case JDBC_ARROW_TREAT_DECIMAL_AS_INT: + if (propertyValue != null) { + setJdbcArrowTreatDecimalAsInt(getBooleanValue(propertyValue)); + } + break; + default: break; } diff --git a/src/main/java/net/snowflake/client/core/SFSessionProperty.java b/src/main/java/net/snowflake/client/core/SFSessionProperty.java index d99ae7f22..0ca91809c 100644 --- a/src/main/java/net/snowflake/client/core/SFSessionProperty.java +++ b/src/main/java/net/snowflake/client/core/SFSessionProperty.java @@ -80,7 +80,9 @@ public enum SFSessionProperty { ENABLE_PATTERN_SEARCH("enablePatternSearch", false, Boolean.class), - DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class); + DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class), + + JDBC_ARROW_TREAT_DECIMAL_AS_INT("JDBC_ARROW_TREAT_DECIMAL_AS_INT", false, Boolean.class); // property key in string private String propertyKey; diff --git a/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java index 03b2cace6..c460e9e4c 100644 --- a/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java @@ -9,6 +9,7 @@ import java.sql.Timestamp; import java.util.TimeZone; import net.snowflake.client.core.DataConversionContext; +import net.snowflake.client.core.SFBaseSession; import net.snowflake.client.core.SFException; import net.snowflake.client.jdbc.ErrorCode; import net.snowflake.client.jdbc.SnowflakeUtil; @@ -146,6 +147,15 @@ public BigDecimal toBigDecimal(int index) throws SFException { ErrorCode.INVALID_VALUE_CONVERT, logicalTypeStr, SnowflakeUtil.BIG_DECIMAL_STR, ""); } + public boolean shouldTreatDecimalAsInt(SFBaseSession session) { + if (session != null) { + if (!session.isJdbcArrowTreatDecimalAsInt() && !session.isJdbcTreatDecimalAsInt()) { + return false; + } + } + return true; + } + @Override public void setTreatNTZAsUTC(boolean isUTC) { this.treatNTZasUTC = isUTC; diff --git a/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java index f05a9960d..c38ef2d68 100644 --- a/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java @@ -137,9 +137,10 @@ public BigDecimal toBigDecimal(int index) { public Object toObject(int index) throws SFException { if (bigIntVector.isNull(index)) { return null; - } else { - return getLong(index); + } else if (!shouldTreatDecimalAsInt(context.getSession())) { + return BigDecimal.valueOf(getLong(index), sfScale); } + return getLong(index); } @Override diff --git a/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java index a5faa2215..b7464eb6a 100644 --- a/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java @@ -105,7 +105,12 @@ public BigDecimal toBigDecimal(int index) throws SFException { @Override public Object toObject(int index) throws SFException { - return isNull(index) ? null : (long) getInt(index); + if (isNull(index)) { + return null; + } else if (!shouldTreatDecimalAsInt(context.getSession())) { + return BigDecimal.valueOf((long) getInt(index), sfScale); + } + return (long) getInt(index); } @Override diff --git a/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java index ad44cd224..25122c81d 100644 --- a/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java @@ -99,6 +99,8 @@ public double toDouble(int index) throws SFException { public Object toObject(int index) throws SFException { if (isNull(index)) { return null; + } else if (!shouldTreatDecimalAsInt(context.getSession())) { + return BigDecimal.valueOf((long) getShort(index), sfScale); } return (long) getShort(index); } diff --git a/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java b/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java index 5cfbd3e5d..46b2cdea8 100644 --- a/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java @@ -90,7 +90,12 @@ public BigDecimal toBigDecimal(int index) throws SFException { @Override public Object toObject(int index) throws SFException { - return isNull(index) ? null : (long) toByte(index); + if (isNull(index)) { + return null; + } else if (!shouldTreatDecimalAsInt(context.getSession())) { + return BigDecimal.valueOf((long) getByte(index), sfScale); + } + return (long) toByte(index); } @Override diff --git a/src/test/java/net/snowflake/client/jdbc/ResultSetLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ResultSetLatestIT.java index f9e5f0293..da91f2c28 100644 --- a/src/test/java/net/snowflake/client/jdbc/ResultSetLatestIT.java +++ b/src/test/java/net/snowflake/client/jdbc/ResultSetLatestIT.java @@ -73,6 +73,14 @@ public ResultSetLatestIT() { super(queryResultFormat); } + private String createTableSql = + "Create or replace table get_object_for_numeric_types (c1 INT, c2 BIGINT, c3 SMALLINT, c4 TINYINT) "; + private String insertStmt = + "Insert into get_object_for_numeric_types (c1, c2, c3, c4) values (1000000000, 2000000000000000000000000, 3, 4)"; + private String selectQuery = "Select * from get_object_for_numeric_types"; + private String setJdbcTreatDecimalAsIntFalse = + "alter session set JDBC_TREAT_DECIMAL_AS_INT = false"; + /** * Test that when closing of results is interrupted by Thread.Interrupt(), the memory is released * safely before driver execution ends. @@ -971,4 +979,94 @@ public void testLargeStringRetrieval() throws SQLException { fail("executeQuery should not fail"); } } + + private static void assertAllColumnsAreLongButBigIntIsBigDecimal(ResultSet rs) + throws SQLException { + while (rs.next()) { + assertEquals(java.lang.Long.class, rs.getObject(1).getClass()); + assertEquals(java.math.BigDecimal.class, rs.getObject(2).getClass()); + assertEquals(java.lang.Long.class, rs.getObject(3).getClass()); + assertEquals(java.lang.Long.class, rs.getObject(4).getClass()); + } + } + + private static void assertAllColumnsAreBigDecimal(ResultSet rs) throws SQLException { + while (rs.next()) { + assertEquals(java.math.BigDecimal.class, rs.getObject(1).getClass()); + assertEquals(java.math.BigDecimal.class, rs.getObject(2).getClass()); + assertEquals(java.math.BigDecimal.class, rs.getObject(3).getClass()); + assertEquals(java.math.BigDecimal.class, rs.getObject(4).getClass()); + } + } + + // Test setting new connection property JDBC_ARROW_TREAT_DECIMAL_AS_INT=false. Connection property + // introduced after version 3.15.0. + @Test + public void testGetObjectForArrowResultFormatJDBCArrowDecimalAsIntFalse() throws SQLException { + Properties properties = new Properties(); + properties.put("JDBC_ARROW_TREAT_DECIMAL_AS_INT", false); + try (Connection con = getConnection(properties); + Statement stmt = con.createStatement()) { + stmt.execute("alter session set jdbc_query_result_format = 'ARROW'"); + stmt.execute(createTableSql); + stmt.execute(insertStmt); + + // Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=false and JDBC_TREAT_DECIMAL_AS_INT=true + try (ResultSet rs = stmt.executeQuery(selectQuery)) { + assertAllColumnsAreLongButBigIntIsBigDecimal(rs); + } + + // Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=false and JDBC_TREAT_DECIMAL_AS_INT=false + stmt.execute(setJdbcTreatDecimalAsIntFalse); + try (ResultSet rs = stmt.executeQuery(selectQuery)) { + assertAllColumnsAreBigDecimal(rs); + } + } + } + + // Test default setting of new connection property JDBC_ARROW_TREAT_DECIMAL_AS_INT=true. + // Connection property introduced after version 3.15.0. + @Test + public void testGetObjectForArrowResultFormatJDBCArrowDecimalAsIntTrue() throws SQLException { + try (Connection con = BaseJDBCTest.getConnection(); + Statement stmt = con.createStatement()) { + stmt.execute("alter session set jdbc_query_result_format = 'ARROW'"); + stmt.execute(createTableSql); + stmt.execute(insertStmt); + + // Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=true and JDBC_TREAT_DECIMAL_AS_INT=true + try (ResultSet rs = stmt.executeQuery(selectQuery)) { + assertAllColumnsAreLongButBigIntIsBigDecimal(rs); + } + + // Test with JDBC_ARROW_TREAT_DECIMAL_AS_INT=true and JDBC_TREAT_DECIMAL_AS_INT=false + stmt.execute(setJdbcTreatDecimalAsIntFalse); + try (ResultSet rs = stmt.executeQuery(selectQuery)) { + assertAllColumnsAreLongButBigIntIsBigDecimal(rs); + } + } + } + + // Test getObject for numeric types when JDBC_TREAT_DECIMAL_AS_INT is set and using JSON result + // format. + @Test + public void testGetObjectForJSONResultFormatUsingJDBCDecimalAsInt() throws SQLException { + try (Connection con = BaseJDBCTest.getConnection(); + Statement stmt = con.createStatement()) { + stmt.execute("alter session set jdbc_query_result_format = 'JSON'"); + stmt.execute(createTableSql); + stmt.execute(insertStmt); + + // Test with JDBC_TREAT_DECIMAL_AS_INT=true (default value) + try (ResultSet rs = stmt.executeQuery(selectQuery)) { + assertAllColumnsAreLongButBigIntIsBigDecimal(rs); + } + + // Test with JDBC_TREAT_DECIMAL_AS_INT=false + stmt.execute(setJdbcTreatDecimalAsIntFalse); + try (ResultSet rs = stmt.executeQuery(selectQuery)) { + assertAllColumnsAreBigDecimal(rs); + } + } + } }