Added extra tests and related fixes

snowflakedb · Sep 25, 2024 · 5cd55a7 · 5cd55a7
1 parent 82a68b7
commit 5cd55a7
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/...n/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java b/...n/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
@@ -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");