Arrow batches complete (for testing)

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

@@ -30,6 +30,8 @@
 import net.snowflake.client.core.arrow.VarCharConverter;
 import net.snowflake.client.core.arrow.VectorTypeConverter;
 import net.snowflake.client.core.json.Converters;
+import net.snowflake.client.jdbc.ArrowBatch;
+import net.snowflake.client.jdbc.ArrowBatches;
 import net.snowflake.client.jdbc.ArrowResultChunk;
 import net.snowflake.client.jdbc.ArrowResultChunk.ArrowChunkIterator;
 import net.snowflake.client.jdbc.ErrorCode;
@@ -112,6 +114,15 @@ public class SFArrowResultSet extends SFBaseResultSet implements DataConversionC
    */
   private boolean formatDateWithTimezone;
 
+  /** The result set should be read either only as rows or only as batches */
+  private enum ReadingMode {
+    UNSPECIFIED,
+    ROW_MODE,
+    BATCH_MODE
+  }
+
+  private ReadingMode readingMode = ReadingMode.UNSPECIFIED;
+
   @SnowflakeJdbcInternalApi protected Converters converters;
 
   /**
@@ -238,6 +249,11 @@ public SFArrowResultSet(
     }
   }
 
+  @SnowflakeJdbcInternalApi
+  public long getAllocatedMemory() {
+    return rootAllocator.getAllocatedMemory();
+  }
+
   private boolean fetchNextRow() throws SnowflakeSQLException {
     if (sortResult) {
       return fetchNextRowSorted();
@@ -246,6 +262,31 @@ private boolean fetchNextRow() throws SnowflakeSQLException {
     }
   }
 
+  private ArrowResultChunk fetchNextChunk() throws SnowflakeSQLException {
+    try {
+      logger.debug("Fetching chunk number " + nextChunkIndex);
+      eventHandler.triggerStateTransition(
+          BasicEvent.QueryState.CONSUMING_RESULT,
+          String.format(
+              BasicEvent.QueryState.CONSUMING_RESULT.getArgString(), queryId, nextChunkIndex));
+      ArrowResultChunk nextChunk = (ArrowResultChunk) chunkDownloader.getNextChunkToConsume();
+
+      if (nextChunk == null) {
+        throw new SnowflakeSQLLoggedException(
+            queryId,
+            session,
+            ErrorCode.INTERNAL_ERROR.getMessageCode(),
+            SqlState.INTERNAL_ERROR,
+            "Expect chunk but got null for chunk index " + nextChunkIndex);
+      }
+      logger.debug("Chunk number " + nextChunkIndex + " fetched successfully.");
+      return nextChunk;
+    } catch (InterruptedException ex) {
+      throw new SnowflakeSQLLoggedException(
+          queryId, session, ErrorCode.INTERRUPTED.getMessageCode(), SqlState.QUERY_CANCELED);
+    }
+  }
+
   /**
    * Goto next row. If end of current chunk, update currentChunkIterator to the beginning of next
    * chunk, if any chunk not being consumed yet.
@@ -259,40 +300,19 @@ private boolean fetchNextRowUnsorted() throws SnowflakeSQLException {
       return true;
     } else {
       if (nextChunkIndex < chunkCount) {
-        try {
-          eventHandler.triggerStateTransition(
-              BasicEvent.QueryState.CONSUMING_RESULT,
-              String.format(
-                  BasicEvent.QueryState.CONSUMING_RESULT.getArgString(), queryId, nextChunkIndex));
-
-          ArrowResultChunk nextChunk = (ArrowResultChunk) chunkDownloader.getNextChunkToConsume();
-
-          if (nextChunk == null) {
-            throw new SnowflakeSQLLoggedException(
-                queryId,
-                session,
-                ErrorCode.INTERNAL_ERROR.getMessageCode(),
-                SqlState.INTERNAL_ERROR,
-                "Expect chunk but got null for chunk index " + nextChunkIndex);
-          }
+        ArrowResultChunk nextChunk = fetchNextChunk();
 
-          currentChunkIterator.getChunk().freeData();
-          currentChunkIterator = nextChunk.getIterator(this);
-          if (currentChunkIterator.next()) {
-
-            logger.debug(
-                "Moving to chunk index: {}, row count: {}",
-                nextChunkIndex,
-                nextChunk.getRowCount());
-
-            nextChunkIndex++;
-            return true;
-          } else {
-            return false;
-          }
-        } catch (InterruptedException ex) {
-          throw new SnowflakeSQLLoggedException(
-              queryId, session, ErrorCode.INTERRUPTED.getMessageCode(), SqlState.QUERY_CANCELED);
+        currentChunkIterator.getChunk().freeData();
+        currentChunkIterator = nextChunk.getIterator(this);
+        if (currentChunkIterator.next()) {
+
+          logger.debug(
+              "Moving to chunk index: {}, row count: {}", nextChunkIndex, nextChunk.getRowCount());
+
+          nextChunkIndex++;
+          return true;
+        } else {
+          return false;
         }
       } else {
         // always free current chunk
@@ -431,6 +451,11 @@ public boolean next() throws SFException, SnowflakeSQLException {
     if (isClosed()) {
       return false;
     }
+    if (readingMode == ReadingMode.BATCH_MODE) {
+      logger.warn("Cannot read rows after getArrowBatches() was called.");
+      return false;
+    }
+    readingMode = ReadingMode.ROW_MODE;
 
     // otherwise try to fetch again
     if (fetchNextRow()) {
@@ -763,6 +788,43 @@ public BigDecimal getBigDecimal(int columnIndex, int scale) throws SFException {
     return bigDec == null ? null : bigDec.setScale(scale, RoundingMode.HALF_UP);
   }
 
+  public ArrowBatches getArrowBatches() {
+    if (readingMode == ReadingMode.ROW_MODE) {
+      logger.warn("Cannot read arrow batches after next() was called.");
+      return null;
+    }
+    readingMode = ReadingMode.BATCH_MODE;
+    return new SFArrowBatchesIterator();
+  }
+
+  private class SFArrowBatchesIterator implements ArrowBatches {
+    private boolean firstFetched = false;
+
+    @Override
+    public long getRowCount() throws SQLException {
+      return resultSetSerializable.getRowCount();
+    }
+
+    @Override
+    public boolean hasNext() {
+      return nextChunkIndex < chunkCount || !firstFetched;
+    }
+
+    @Override
+    public ArrowBatch next() throws SQLException {
+      if (!firstFetched) {
+        firstFetched = true;
+        return currentChunkIterator
+            .getChunk()
+            .getArrowBatch(SFArrowResultSet.this, useSessionTimezone ? sessionTimeZone : null);
+      } else {
+        nextChunkIndex++;
+        return fetchNextChunk()
+            .getArrowBatch(SFArrowResultSet.this, useSessionTimezone ? sessionTimeZone : null);
+      }
+    }
+  }
+
   @Override
   public boolean isLast() {
     return nextChunkIndex == chunkCount && currentChunkIterator.isLast();

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

@@ -30,8 +30,10 @@
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 import net.snowflake.client.core.json.Converters;
+import net.snowflake.client.jdbc.ArrowBatches;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.FieldMetadata;
+import net.snowflake.client.jdbc.SnowflakeLoggedFeatureNotSupportedException;
 import net.snowflake.client.jdbc.SnowflakeResultSetSerializable;
 import net.snowflake.client.jdbc.SnowflakeResultSetSerializableV1;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
@@ -137,6 +139,10 @@ public SFBaseSession getSession() {
     return this.session;
   }
 
+  public ArrowBatches getArrowBatches() throws SnowflakeLoggedFeatureNotSupportedException {
+    throw new SnowflakeLoggedFeatureNotSupportedException(session);
+  }
+
   // default implementation
   public boolean next() throws SFException, SnowflakeSQLException {
     logger.trace("boolean next()", false);

src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/ArrowFullVectorConverter.java

@@ -0,0 +1,142 @@
+package net.snowflake.client.core.arrow.fullvectorconverters;
+
+import java.util.Map;
+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.core.SnowflakeJdbcInternalApi;
+import net.snowflake.client.jdbc.SnowflakeSQLException;
+import net.snowflake.client.jdbc.SnowflakeType;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.types.Types;
+
+@SnowflakeJdbcInternalApi
+public interface ArrowFullVectorConverter {
+  static Types.MinorType deduceType(ValueVector vector) {
+    Types.MinorType type = Types.getMinorTypeForArrowType(vector.getField().getType());
+    // each column's metadata
+    Map<String, String> customMeta = vector.getField().getMetadata();
+    if (type == Types.MinorType.DECIMAL) {
+      // Note: Decimal vector is different from others
+      return Types.MinorType.DECIMAL;
+    } else if (!customMeta.isEmpty()) {
+      SnowflakeType st = SnowflakeType.valueOf(customMeta.get("logicalType"));
+      switch (st) {
+        case FIXED:
+          {
+            String scaleStr = vector.getField().getMetadata().get("scale");
+            int sfScale = Integer.parseInt(scaleStr);
+            if (sfScale != 0) {
+              return Types.MinorType.DECIMAL;
+            }
+            break;
+          }
+        case VECTOR:
+          return Types.MinorType.FIXED_SIZE_LIST;
+        case TIME:
+          {
+            String scaleStr = vector.getField().getMetadata().get("scale");
+            int sfScale = Integer.parseInt(scaleStr);
+            if (sfScale == 0) {
+              return Types.MinorType.TIMESEC;
+            }
+            if (sfScale <= 3) {
+              return Types.MinorType.TIMEMILLI;
+            }
+            if (sfScale <= 6) {
+              return Types.MinorType.TIMEMICRO;
+            }
+            if (sfScale <= 9) {
+              return Types.MinorType.TIMENANO;
+            }
+          }
+        case TIMESTAMP_NTZ:
+          return Types.MinorType.TIMESTAMPNANO;
+        case TIMESTAMP_LTZ:
+        case TIMESTAMP_TZ:
+          return Types.MinorType.TIMESTAMPNANOTZ;
+      }
+    }
+    return type;
+  }
+
+  static FieldVector convert(
+      RootAllocator allocator,
+      ValueVector vector,
+      DataConversionContext context,
+      SFBaseSession session,
+      TimeZone timeZoneToUse,
+      int idx,
+      Object targetType)
+      throws SnowflakeSQLException {
+    try {
+      if (targetType == null) {
+        targetType = deduceType(vector);
+      }
+      if (targetType instanceof Types.MinorType) {
+        switch ((Types.MinorType) targetType) {
+          case TINYINT:
+            return new TinyIntVectorConverter(allocator, vector, context, session, idx).convert();
+          case SMALLINT:
+            return new SmallIntVectorConverter(allocator, vector, context, session, idx).convert();
+          case INT:
+            return new IntVectorConverter(allocator, vector, context, session, idx).convert();
+          case BIGINT:
+            return new BigIntVectorConverter(allocator, vector, context, session, idx).convert();
+          case DECIMAL:
+            return new DecimalVectorConverter(allocator, vector, context, session, idx).convert();
+          case FLOAT8:
+            return new FloatVectorConverter(allocator, vector, context, session, idx).convert();
+          case BIT:
+            return new BitVectorConverter(allocator, vector, context, session, idx).convert();
+          case VARBINARY:
+            return new BinaryVectorConverter(allocator, vector, context, session, idx).convert();
+          case DATEDAY:
+            return new DateVectorConverter(allocator, vector, context, session, idx, timeZoneToUse)
+                .convert();
+          case TIMESEC:
+            return new TimeSecVectorConverter(allocator, vector).convert();
+          case TIMEMILLI:
+            return new TimeMilliVectorConverter(allocator, vector).convert();
+          case TIMEMICRO:
+            return new TimeMicroVectorConverter(allocator, vector).convert();
+          case TIMENANO:
+            return new TimeNanoVectorConverter(allocator, vector).convert();
+          case TIMESTAMPNANOTZ:
+            return new TimestampVectorConverter(allocator, vector, context, timeZoneToUse, false)
+                .convert();
+          case TIMESTAMPNANO:
+            return new TimestampVectorConverter(allocator, vector, context, timeZoneToUse, true)
+                .convert();
+          case STRUCT:
+            return new StructVectorConverter(
+                    allocator, vector, context, session, timeZoneToUse, idx, null)
+                .convert();
+          case LIST:
+            return new ListVectorConverter(
+                    allocator, vector, context, session, timeZoneToUse, idx, null)
+                .convert();
+          case VARCHAR:
+            return new VarCharVectorConverter(allocator, vector, context, session, idx).convert();
+          case MAP:
+            return new MapVectorConverter(
+                    allocator, vector, context, session, timeZoneToUse, idx, null)
+                .convert();
+          case FIXED_SIZE_LIST:
+            return new FixedSizeListVectorConverter(
+                    allocator, vector, context, session, timeZoneToUse, idx, null)
+                .convert();
+        }
+      }
+    } catch (SFException ex) {
+      throw new SnowflakeSQLException(
+          ex.getCause(), ex.getSqlState(), ex.getVendorCode(), ex.getParams());
+    }
+    return null;
+  }
+
+  FieldVector convert() throws SFException, SnowflakeSQLException;
+}

src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BigIntVectorConverter.java

@@ -0,0 +1,41 @@
+package net.snowflake.client.core.arrow.fullvectorconverters;
+
+import net.snowflake.client.core.DataConversionContext;
+import net.snowflake.client.core.SFBaseSession;
+import net.snowflake.client.core.SFException;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+import net.snowflake.client.core.arrow.ArrowVectorConverter;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.ValueVector;
+
+@SnowflakeJdbcInternalApi
+public class BigIntVectorConverter extends SimpleArrowFullVectorConverter<BigIntVector> {
+
+  public BigIntVectorConverter(
+      RootAllocator allocator,
+      ValueVector vector,
+      DataConversionContext context,
+      SFBaseSession session,
+      int idx) {
+    super(allocator, vector, context, session, idx);
+  }
+
+  @Override
+  protected boolean matchingType() {
+    return (vector instanceof BigIntVector);
+  }
+
+  @Override
+  protected BigIntVector initVector() {
+    BigIntVector resultVector = new BigIntVector(vector.getName(), allocator);
+    resultVector.allocateNew(vector.getValueCount());
+    return resultVector;
+  }
+
+  @Override
+  protected void convertValue(ArrowVectorConverter from, BigIntVector to, int idx)
+      throws SFException {
+    to.set(idx, from.toLong(idx));
+  }
+}