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

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

@@ -139,6 +139,8 @@ public abstract class SFBaseSession {
 
   private Map<String, Object> commonParameters;
 
+  private boolean isJdbcArrowTreatDecimalAsInt = true;
+
   protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
     this.sfConnectionHandler = sfConnectionHandler;
   }
@@ -267,6 +269,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);

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

@@ -470,6 +470,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;
       }

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

@@ -78,7 +78,9 @@ public enum SFSessionProperty {
 
   RETRY_TIMEOUT("retryTimeout", false, Integer.class),
 
-  ENABLE_PATTERN_SEARCH("enablePatternSearch", false, Boolean.class);
+  ENABLE_PATTERN_SEARCH("enablePatternSearch", 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;

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

@@ -137,9 +137,11 @@ public BigDecimal toBigDecimal(int index) {
   public Object toObject(int index) throws SFException {
     if (bigIntVector.isNull(index)) {
       return null;
-    } else {
-      return getLong(index);
+    } else if (!context.getSession().isJdbcArrowTreatDecimalAsInt()
+        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+      return BigDecimal.valueOf(getLong(index), sfScale);
     }
+    return getLong(index);
   }
 
   @Override

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

@@ -105,7 +105,13 @@ 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 (!context.getSession().isJdbcArrowTreatDecimalAsInt()
+        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getInt(index), sfScale);
+    }
+    return (long) getInt(index);
   }
 
   @Override

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

@@ -99,6 +99,9 @@ public double toDouble(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
+    } else if (!context.getSession().isJdbcArrowTreatDecimalAsInt()
+        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getShort(index), sfScale);
     }
     return (long) getShort(index);
   }

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

@@ -90,7 +90,13 @@ 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 (!context.getSession().isJdbcArrowTreatDecimalAsInt()
+        && !context.getSession().isJdbcTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getByte(index), sfScale);
+    }
+    return (long) toByte(index);
   }
 
   @Override

src/test/java/net/snowflake/client/jdbc/ResultSetArrowLatestIT.java

@@ -4,8 +4,17 @@
 package net.snowflake.client.jdbc;
 
 import net.snowflake.client.category.TestCategoryArrow;
+import org.junit.Test;
 import org.junit.experimental.categories.Category;
 
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Properties;
+
+import static org.junit.Assert.assertEquals;
+
 /**
  * ResultSet integration tests for the latest JDBC driver. This doesn't work for the oldest
  * supported driver. Drop this file when ResultSetLatestIT is dropped.

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,138 @@ public void testLargeStringRetrieval() throws SQLException {
       fail("executeQuery should not fail");
     }
   }
+
+  // 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)) {
+        while(rs.next()) {
+          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
+        }
+      }
+
+      // 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)) {
+        while (rs.next()) {
+          assertEquals("class java.math.BigDecimal", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(3).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(4).getClass().toString());
+        }
+      }
+    }
+  }
+
+  // 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)) {
+        while(rs.next()) {
+          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
+        }
+      }
+
+      // 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)) {
+        while (rs.next()) {
+          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
+        }
+      }
+    }
+  }
+
+  // Test setting new connection property jdbc_arrow_treat_decimal_as_int=false. Connection property
+  // introduced after version 3.15.0.
+  @Test
+  public void testGetObjectForJSONResultFormatJDBCArrowDecimalAsIntFalse() 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 = 'JSON'");
+      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)) {
+        while (rs.next()) {
+          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
+        }
+      }
+
+      // 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)) {
+        while (rs.next()) {
+          assertEquals("class java.math.BigDecimal", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(3).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(4).getClass().toString());
+        }
+      }
+    }
+  }
+
+  // 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 testGetObjectForJSONResultFormatJDBCArrowDecimalAsIntTrue() 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_arrow_treat_decimal_as_int=true and JDBC_TREAT_DECIMAL_AS_INT=true
+      try (ResultSet rs = stmt.executeQuery(selectQuery)) {
+        while (rs.next()) {
+          assertEquals("class java.lang.Long", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(3).getClass().toString());
+          assertEquals("class java.lang.Long", rs.getObject(4).getClass().toString());
+        }
+      }
+
+      // 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)) {
+        while (rs.next()) {
+          assertEquals("class java.math.BigDecimal", rs.getObject(1).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(2).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(3).getClass().toString());
+          assertEquals("class java.math.BigDecimal", rs.getObject(4).getClass().toString());
+        }
+      }
+    }
+  }
 }