SNOW-1234216 Native Arrow structured types array support

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

@@ -187,11 +187,7 @@ public Timestamp readTimestamp(TimeZone tz) throws SQLException {
                   converters
                       .getStructuredTypeDateTimeConverter()
                       .getTimestamp(
-                          (JsonStringHashMap<String, Object>) value,
-                          columnType,
-                          columnSubType,
-                          tz,
-                          scale));
+                          (Map<String, Object>) value, columnType, columnSubType, tz, scale));
         });
   }
 

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

@@ -16,11 +16,15 @@
 import java.sql.Array;
 import java.sql.Date;
 import java.sql.SQLException;
+import java.sql.SQLInput;
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.sql.Types;
+import java.util.List;
 import java.util.Map;
 import java.util.TimeZone;
+import java.util.stream.Stream;
+import net.snowflake.client.core.arrow.ArrayConverter;
 import net.snowflake.client.core.arrow.ArrowVectorConverter;
 import net.snowflake.client.core.arrow.StructConverter;
 import net.snowflake.client.core.arrow.VarCharConverter;
@@ -29,6 +33,8 @@
 import net.snowflake.client.jdbc.ArrowResultChunk;
 import net.snowflake.client.jdbc.ArrowResultChunk.ArrowChunkIterator;
 import net.snowflake.client.jdbc.ErrorCode;
+import net.snowflake.client.jdbc.FieldMetadata;
+import net.snowflake.client.jdbc.SnowflakeColumnMetadata;
 import net.snowflake.client.jdbc.SnowflakeResultSetSerializableV1;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
 import net.snowflake.client.jdbc.SnowflakeSQLLoggedException;
@@ -38,6 +44,7 @@
 import net.snowflake.client.jdbc.telemetry.TelemetryUtil;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
+import net.snowflake.client.util.Converter;
 import net.snowflake.common.core.SFBinaryFormat;
 import net.snowflake.common.core.SnowflakeDateTimeFormat;
 import net.snowflake.common.core.SqlState;
@@ -364,18 +371,45 @@ public Converters getConverters() {
   }
 
   @Override
+  @SnowflakeJdbcInternalApi
+  public SQLInput createSqlInputForColumn(
+      Object input, Class<?> parentObjectClass, int columnIndex, SFBaseSession session) {
+    if (parentObjectClass.equals(JsonSqlInput.class)) {
+      return createJsonSqlInputForColumn(input, columnIndex, session);
+    } else {
+      return new ArrowSqlInput(
+          (Map<String, Object>) input,
+          session,
+          converters,
+          resultSetMetaData.getColumnMetadata().get(columnIndex - 1).getFields());
+    }
+  }
+
+  @Override
+  @SnowflakeJdbcInternalApi
   public Date convertToDate(Object object, TimeZone tz) throws SFException {
+    if (object instanceof String) {
+      return convertStringToDate(object, tz);
+    }
     return converters.getStructuredTypeDateTimeConverter().getDate((int) object, tz);
   }
 
   @Override
+  @SnowflakeJdbcInternalApi
   public Time convertToTime(Object object, int scale) throws SFException {
+    if (object instanceof String) {
+      return convertStringToTime(object, scale);
+    }
     return converters.getStructuredTypeDateTimeConverter().getTime((long) object, scale);
   }
 
   @Override
+  @SnowflakeJdbcInternalApi
   public Timestamp convertToTimestamp(
       Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException {
+    if (object instanceof String) {
+      return convertStringToTimestamp(object, columnType, columnSubType, tz, scale);
+    }
     return converters
         .getStructuredTypeDateTimeConverter()
         .getTimestamp(
@@ -566,8 +600,126 @@ private Object createArrowSqlInput(int columnIndex, Map<String, Object> input) {
 
   @Override
   public Array getArray(int columnIndex) throws SFException {
-    //    TODO: handleArray SNOW-969794
-    throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "data type ARRAY");
+    ArrowVectorConverter converter = currentChunkIterator.getCurrentConverter(columnIndex - 1);
+    int index = currentChunkIterator.getCurrentRowInRecordBatch();
+    wasNull = converter.isNull(index);
+    Object obj = converter.toObject(index);
+    if (converter instanceof VarCharConverter) {
+      return getJsonArrayInternal((String) obj, columnIndex);
+    } else if (converter instanceof ArrayConverter) {
+      return getArrayInternal((List<Object>) obj, columnIndex);
+    } else {
+      throw new SFException(ErrorCode.INVALID_STRUCT_DATA);
+    }
+  }
+
+  private SfSqlArray getArrayInternal(List<Object> elements, int columnIndex) throws SFException {
+    try {
+      SnowflakeColumnMetadata arrayMetadata =
+          resultSetMetaData.getColumnMetadata().get(columnIndex - 1);
+      FieldMetadata fieldMetadata = arrayMetadata.getFields().get(0);
+
+      int columnSubType = fieldMetadata.getType();
+      int columnType = ColumnTypeHelper.getColumnType(columnSubType, session);
+      int scale = fieldMetadata.getScale();
+
+      switch (columnSubType) {
+        case Types.INTEGER:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.integerConverter(columnType))
+                  .toArray(Integer[]::new));
+        case Types.SMALLINT:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.smallIntConverter(columnType))
+                  .toArray(Short[]::new));
+        case Types.TINYINT:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.tinyIntConverter(columnType))
+                  .toArray(Byte[]::new));
+        case Types.BIGINT:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.bigIntConverter(columnType)).toArray(Long[]::new));
+        case Types.DECIMAL:
+        case Types.NUMERIC:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.bigDecimalConverter(columnType))
+                  .toArray(BigDecimal[]::new));
+        case Types.CHAR:
+        case Types.VARCHAR:
+        case Types.LONGNVARCHAR:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.varcharConverter(columnType, columnSubType, scale))
+                  .toArray(String[]::new));
+        case Types.BINARY:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.bytesConverter(columnType, scale))
+                  .toArray(Byte[][]::new));
+        case Types.FLOAT:
+        case Types.REAL:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.floatConverter(columnType)).toArray(Float[]::new));
+        case Types.DOUBLE:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.doubleConverter(columnType))
+                  .toArray(Double[]::new));
+        case Types.DATE:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.dateFromIntConverter(sessionTimezone))
+                  .toArray(Date[]::new));
+        case Types.TIME:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.timeFromIntConverter(scale)).toArray(Time[]::new));
+        case Types.TIMESTAMP:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(
+                      elements,
+                      converters.timestampFromStructConverter(
+                          columnType, columnSubType, sessionTimezone, scale))
+                  .toArray(Timestamp[]::new));
+        case Types.BOOLEAN:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, converters.booleanConverter(columnType))
+                  .toArray(Boolean[]::new));
+        case Types.STRUCT:
+          return new SfSqlArray(columnSubType, mapAndConvert(elements, e -> e).toArray(Map[]::new));
+        case Types.ARRAY:
+          return new SfSqlArray(
+              columnSubType,
+              mapAndConvert(elements, e -> ((List) e).stream().toArray(Map[]::new))
+                  .toArray(Map[][]::new));
+        default:
+          throw new SFException(
+              ErrorCode.FEATURE_UNSUPPORTED,
+              "Can't construct array for data type: " + columnSubType);
+      }
+    } catch (RuntimeException e) {
+      throw new SFException(e, ErrorCode.INVALID_STRUCT_DATA);
+    }
+  }
+
+  private <T> Stream<T> mapAndConvert(List<Object> elements, Converter<T> converter) {
+    return elements.stream()
+        .map(
+            obj -> {
+              try {
+                return converter.convert(obj);
+              } catch (SFException e) {
+                throw new RuntimeException(e);
+              }
+            });
   }
 
   @Override