SNOW-1170182: Return decimal as int in ARROW results configured by fl…

…ag (#1666)
snowflakedb · Mar 15, 2024 · f279079 · f279079
1 parent 75c57f2
commit f279079
Show file tree

Hide file tree

Showing 9 changed files with 144 additions and 5 deletions.
diff --git 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<String, Object> 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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);
+      }
+    }
+  }
 }