SNOW-873466 arrow batches

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,45 @@ 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, nextChunkIndex);
+      } else {
+        nextChunkIndex++;
+        return fetchNextChunk()
+            .getArrowBatch(
+                SFArrowResultSet.this, useSessionTimezone ? sessionTimeZone : null, nextChunkIndex);
+      }
+    }
+  }
+
   @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/ArrowVectorConverterUtil.java

@@ -20,6 +20,20 @@
 public final class ArrowVectorConverterUtil {
   private ArrowVectorConverterUtil() {}
 
+  public static int getScale(ValueVector vector, SFBaseSession session)
+      throws SnowflakeSQLLoggedException {
+    try {
+      String scaleStr = vector.getField().getMetadata().get("scale");
+      return Integer.parseInt(scaleStr);
+    } catch (NullPointerException | NumberFormatException e) {
+      throw new SnowflakeSQLLoggedException(
+          session,
+          ErrorCode.INTERNAL_ERROR.getMessageCode(),
+          SqlState.INTERNAL_ERROR,
+          "Invalid scale metadata");
+    }
+  }
+
   /**
    * Given an arrow vector (a single column in a single record batch), return an arrow vector
    * converter. Note, converter is built on top of arrow vector, so that arrow data can be converted
@@ -102,8 +116,7 @@ public static ArrowVectorConverter initConverter(
           return new DateConverter(vector, idx, context, getFormatDateWithTimeZone);
 
         case FIXED:
-          String scaleStr = vector.getField().getMetadata().get("scale");
-          int sfScale = Integer.parseInt(scaleStr);
+          int sfScale = getScale(vector, session);
           switch (type) {
             case TINYINT:
               if (sfScale == 0) {

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

@@ -0,0 +1,23 @@
+package net.snowflake.client.core.arrow.fullvectorconverters;
+
+import net.snowflake.client.core.SFException;
+import net.snowflake.client.jdbc.SnowflakeSQLException;
+import org.apache.arrow.vector.FieldVector;
+
+public abstract class AbstractFullVectorConverter implements ArrowFullVectorConverter {
+  private boolean converted;
+
+  protected abstract FieldVector convertVector()
+      throws SFException, SnowflakeSQLException, SFArrowException;
+
+  @Override
+  public FieldVector convert() throws SFException, SnowflakeSQLException, SFArrowException {
+    if (converted) {
+      throw new SFArrowException(
+          ArrowErrorCode.VECTOR_ALREADY_CONVERTED, "Convert has already been called");
+    } else {
+      converted = true;
+      return convertVector();
+    }
+  }
+}

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

@@ -0,0 +1,7 @@
+package net.snowflake.client.core.arrow.fullvectorconverters;
+
+public enum ArrowErrorCode {
+  VECTOR_ALREADY_CONVERTED,
+  CONVERT_FAILED,
+  CHUNK_FETCH_FAILED,
+}

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

@@ -0,0 +1,11 @@
+package net.snowflake.client.core.arrow.fullvectorconverters;
+
+import net.snowflake.client.core.SFException;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+import net.snowflake.client.jdbc.SnowflakeSQLException;
+import org.apache.arrow.vector.FieldVector;
+
+@SnowflakeJdbcInternalApi
+public interface ArrowFullVectorConverter {
+  FieldVector convert() throws SFException, SnowflakeSQLException, SFArrowException;
+}