From 9851856e4946335ddae54724bb66069fbd633004 Mon Sep 17 00:00:00 2001
From: Dominik Przybysz <dominik.przybysz@snowflake.com>
Date: Mon, 18 Mar 2024 08:48:28 +0100
Subject: [PATCH] SNOW-1170182: Simplify condition returning decimal in ARROW

---
 .../arrow/AbstractArrowVectorConverter.java     | 17 +++++++++--------
 .../core/arrow/BigIntToFixedConverter.java      |  2 +-
 .../client/core/arrow/IntToFixedConverter.java  |  2 +-
 .../core/arrow/SmallIntToFixedConverter.java    |  2 +-
 .../core/arrow/TinyIntToFixedConverter.java     |  2 +-
 5 files changed, 13 insertions(+), 12 deletions(-)

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 c460e9e4c..c7054442c 100644
--- a/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java
+++ b/src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java
@@ -9,7 +9,6 @@
 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;
@@ -39,6 +38,8 @@ abstract class AbstractArrowVectorConverter implements ArrowVectorConverter {
 
   protected TimeZone sessionTimeZone;
 
+  private boolean shouldTreatDecimalAsInt;
+
   /** Field names of the struct vectors used by timestamp */
   public static final String FIELD_NAME_EPOCH = "epoch"; // seconds since epoch
 
@@ -54,6 +55,11 @@ abstract class AbstractArrowVectorConverter implements ArrowVectorConverter {
     this.valueVector = valueVector;
     this.columnIndex = vectorIndex + 1;
     this.context = context;
+    this.shouldTreatDecimalAsInt =
+        context == null
+            || context.getSession() == null
+            || context.getSession().isJdbcArrowTreatDecimalAsInt()
+            || context.getSession().isJdbcTreatDecimalAsInt();
   }
 
   @Override
@@ -147,13 +153,8 @@ 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;
+  boolean shouldTreatDecimalAsInt() {
+    return shouldTreatDecimalAsInt;
   }
 
   @Override
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 c38ef2d68..71bd123a0 100644
--- a/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java
+++ b/src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java
@@ -137,7 +137,7 @@ public BigDecimal toBigDecimal(int index) {
   public Object toObject(int index) throws SFException {
     if (bigIntVector.isNull(index)) {
       return null;
-    } else if (!shouldTreatDecimalAsInt(context.getSession())) {
+    } else if (!shouldTreatDecimalAsInt()) {
       return BigDecimal.valueOf(getLong(index), sfScale);
     }
     return getLong(index);
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 b7464eb6a..8055081ef 100644
--- a/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java
+++ b/src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java
@@ -107,7 +107,7 @@ public BigDecimal toBigDecimal(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
-    } else if (!shouldTreatDecimalAsInt(context.getSession())) {
+    } else if (!shouldTreatDecimalAsInt()) {
       return BigDecimal.valueOf((long) getInt(index), sfScale);
     }
     return (long) getInt(index);
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 25122c81d..bfa398d88 100644
--- a/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java
+++ b/src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java
@@ -99,7 +99,7 @@ public double toDouble(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
-    } else if (!shouldTreatDecimalAsInt(context.getSession())) {
+    } else if (!shouldTreatDecimalAsInt()) {
       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 46b2cdea8..26c90c228 100644
--- a/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java
+++ b/src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java
@@ -92,7 +92,7 @@ public BigDecimal toBigDecimal(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
-    } else if (!shouldTreatDecimalAsInt(context.getSession())) {
+    } else if (!shouldTreatDecimalAsInt()) {
       return BigDecimal.valueOf((long) getByte(index), sfScale);
     }
     return (long) toByte(index);