From 5cd55a7514f1295558f79d8ee7b7a571daf9d6e9 Mon Sep 17 00:00:00 2001
From: sfc-gh-astachowski <antoni.stachowski@snowflake.com>
Date: Wed, 25 Sep 2024 11:37:29 +0200
Subject: [PATCH] Added extra tests and related fixes

---
 .../TimestampVectorConverter.java             | 20 ++++++++++---------
 .../snowflake/client/jdbc/ArrowBatchesIT.java |  8 +++++++-
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java
index 6e00410af..2d9f5b121 100644
--- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java
+++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java
@@ -74,10 +74,12 @@ private SFPair<BigIntVector, IntVector> normalizeTimeSinceEpoch(BigIntVector vec
     int scale = Integer.parseInt(vector.getField().getMetadata().get("scale"));
     if (scale == 0) {
       IntVector fractions = makeVectorOfZeroes(length);
+      BigIntVector epoch = new BigIntVector(FIELD_NAME_EPOCH, allocator);
       fractions
           .getValidityBuffer()
-          .setBytes(0L, vector.getValidityBuffer(), 0L, fractions.getValidityBuffer().capacity());
-      return SFPair.of(vector, fractions);
+          .setBytes(0L, vector.getValidityBuffer(), 0L, vector.getValidityBuffer().capacity());
+      vector.makeTransferPair(epoch).transfer();
+      return SFPair.of(epoch, fractions);
     }
     long scaleFactor = ArrowResultUtil.powerOfTen(scale);
     long fractionScaleFactor = ArrowResultUtil.powerOfTen(9 - scale);
@@ -97,9 +99,9 @@ private SFPair<BigIntVector, IntVector> normalizeTimeSinceEpoch(BigIntVector vec
   private IntVector makeTimeZoneOffsets(
       BigIntVector seconds, IntVector fractions, TimeZone timeZone) {
     IntVector offsets = new IntVector(FIELD_NAME_TIME_ZONE_INDEX, allocator);
-    offsets.allocateNew(vector.getValueCount());
-    offsets.setValueCount(vector.getValueCount());
-    for (int i = 0; i < vector.getValueCount(); i++) {
+    offsets.allocateNew(seconds.getValueCount());
+    offsets.setValueCount(seconds.getValueCount());
+    for (int i = 0; i < seconds.getValueCount(); i++) {
       offsets.set(
           i,
           UTC_OFFSET
@@ -125,10 +127,10 @@ private StructVector pack(BigIntVector seconds, IntVector fractions, IntVector o
     seconds.makeTransferPair(result.getChild(FIELD_NAME_EPOCH)).transfer();
     fractions.makeTransferPair(result.getChild(FIELD_NAME_FRACTION)).transfer();
     offsets.makeTransferPair(result.getChild(FIELD_NAME_TIME_ZONE_INDEX)).transfer();
-    result.setValueCount(vector.getValueCount());
+    result.setValueCount(seconds.getValueCount());
     result
         .getValidityBuffer()
-        .setBytes(0L, vector.getValidityBuffer(), 0L, vector.getValidityBuffer().capacity());
+        .setBytes(0L, seconds.getValidityBuffer(), 0L, seconds.getValidityBuffer().capacity());
     return result;
   }
 
@@ -163,13 +165,13 @@ public FieldVector convert() throws SFException, SnowflakeSQLException {
       if (timeZoneIndices == null) {
         if (isNTZ && context.getHonorClientTZForTimestampNTZ()) {
           timeZoneIndices = makeTimeZoneOffsets(seconds, fractions, TimeZone.getDefault());
-          for (int i = 0; i < vector.getValueCount(); i++) {
+          for (int i = 0; i < seconds.getValueCount(); i++) {
             seconds.set(
                 i,
                 seconds.get(i) - (long) (timeZoneIndices.get(i) - UTC_OFFSET) * SECONDS_PER_MINUTE);
           }
         } else if (isNTZ || timeZoneToUse == null) {
-          timeZoneIndices = makeVectorOfUTCOffsets(vector.getValueCount());
+          timeZoneIndices = makeVectorOfUTCOffsets(seconds.getValueCount());
         } else {
           timeZoneIndices = makeTimeZoneOffsets(seconds, fractions, timeZoneToUse);
         }
diff --git a/src/test/java/net/snowflake/client/jdbc/ArrowBatchesIT.java b/src/test/java/net/snowflake/client/jdbc/ArrowBatchesIT.java
index b1030b6b5..eb663aea5 100644
--- a/src/test/java/net/snowflake/client/jdbc/ArrowBatchesIT.java
+++ b/src/test/java/net/snowflake/client/jdbc/ArrowBatchesIT.java
@@ -551,7 +551,7 @@ public void testTimeNanoBatch() throws Exception, SFException {
     assertTrue(values.containsAll(expected));
   }
 
-  private void testTimestampBase(String query) throws Exception, SFException {
+  private void testTimestampCase(String query) throws Exception, SFException {
     Timestamp tsFromBatch;
     Timestamp tsFromRow;
 
@@ -575,6 +575,12 @@ private void testTimestampBase(String query) throws Exception, SFException {
     assertTrue(tsFromBatch.equals(tsFromRow));
   }
 
+  private void testTimestampBase(String query) throws Exception, SFException {
+    testTimestampCase(query);
+    testTimestampCase(query + "(0)");
+    testTimestampCase(query + "(1)");
+  }
+
   @Test
   public void testTimestampTZBatch() throws Exception, SFException {
     testTimestampBase("select '2020-04-05 12:22:12+0700'::TIMESTAMP_TZ");