SNOW-1234216 Native Arrow structured types array support

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

@@ -16,7 +16,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.TimeZone;
-import net.snowflake.client.core.json.Converters;
 import net.snowflake.client.core.structs.SQLDataCreationHelper;
 import net.snowflake.client.jdbc.FieldMetadata;
 import net.snowflake.client.util.ThrowingBiFunction;
@@ -187,11 +186,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/BaseSqlInput.java

@@ -17,7 +17,6 @@
 import java.sql.SQLXML;
 import java.sql.Timestamp;
 import java.util.List;
-import net.snowflake.client.core.json.Converters;
 import net.snowflake.client.jdbc.FieldMetadata;
 import net.snowflake.client.jdbc.SnowflakeLoggedFeatureNotSupportedException;
 

src/main/java/net/snowflake/client/core/json/Converters.java → src/main/java/net/snowflake/client/core/Converters.java

@@ -1,4 +1,4 @@
-package net.snowflake.client.core.json;
+package net.snowflake.client.core;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -11,18 +11,20 @@
 import java.util.Arrays;
 import java.util.Map;
 import java.util.TimeZone;
-import net.snowflake.client.core.SFBaseSession;
-import net.snowflake.client.core.SFException;
-import net.snowflake.client.core.SnowflakeJdbcInternalApi;
-import net.snowflake.client.core.SqlInputTimestampUtil;
 import net.snowflake.client.core.arrow.StructuredTypeDateTimeConverter;
+import net.snowflake.client.core.json.BooleanConverter;
+import net.snowflake.client.core.json.BytesConverter;
+import net.snowflake.client.core.json.DateTimeConverter;
+import net.snowflake.client.core.json.NumberConverter;
+import net.snowflake.client.core.json.StringConverter;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.SnowflakeResultSetSerializableV1;
-import net.snowflake.client.util.JsonStringToTypeConverter;
+import net.snowflake.client.util.Converter;
 import net.snowflake.common.core.SFBinaryFormat;
 import net.snowflake.common.core.SFTimestamp;
 import net.snowflake.common.core.SnowflakeDateTimeFormat;
 
+@SnowflakeJdbcInternalApi
 public class Converters {
   private final BooleanConverter booleanConverter;
   private final NumberConverter numberConverter;
@@ -122,47 +124,47 @@ public StructuredTypeDateTimeConverter getStructuredTypeDateTimeConverter() {
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter integerConverter(int columnType) {
+  public Converter integerConverter(int columnType) {
     return value -> getNumberConverter().getInt(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter smallIntConverter(int columnType) {
+  public Converter smallIntConverter(int columnType) {
     return value -> getNumberConverter().getShort(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter tinyIntConverter(int columnType) {
+  public Converter tinyIntConverter(int columnType) {
     return value -> getNumberConverter().getByte(value);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter bigIntConverter(int columnType) {
+  public Converter bigIntConverter(int columnType) {
     return value -> getNumberConverter().getBigInt(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter longConverter(int columnType) {
+  public Converter longConverter(int columnType) {
     return value -> getNumberConverter().getLong(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter bigDecimalConverter(int columnType) {
+  public Converter bigDecimalConverter(int columnType) {
     return value -> getNumberConverter().getBigDecimal(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter floatConverter(int columnType) {
+  public Converter floatConverter(int columnType) {
     return value -> getNumberConverter().getBigDecimal(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter doubleConverter(int columnType) {
+  public Converter doubleConverter(int columnType) {
     return value -> getNumberConverter().getBigDecimal(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter bytesConverter(int columnType, int scale) {
+  public Converter bytesConverter(int columnType, int scale) {
     return value -> {
       byte[] primitiveArray = getBytesConverter().getBytes(value, columnType, Types.BINARY, scale);
       Byte[] newByteArray = new Byte[primitiveArray.length];
@@ -172,29 +174,34 @@ public JsonStringToTypeConverter bytesConverter(int columnType, int scale) {
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter varcharConverter(int columnType, int columnSubType, int scale) {
+  public Converter varcharConverter(int columnType, int columnSubType, int scale) {
     return value -> getStringConverter().getString(value, columnType, columnSubType, scale);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter booleanConverter(int columnType) {
+  public Converter booleanConverter(int columnType) {
     return value -> getBooleanConverter().getBoolean(value, columnType);
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter dateConverter(SFBaseSession session) {
+  public Converter dateStringConverter(SFBaseSession session) {
     return value -> {
       SnowflakeDateTimeFormat formatter =
           SnowflakeDateTimeFormat.fromSqlFormat(
               (String) session.getCommonParameters().get("DATE_OUTPUT_FORMAT"));
-      SFTimestamp timestamp = formatter.parse(value);
+      SFTimestamp timestamp = formatter.parse((String) value);
       return Date.valueOf(
           Instant.ofEpochMilli(timestamp.getTime()).atZone(ZoneOffset.UTC).toLocalDate());
     };
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter timeConverter(SFBaseSession session) {
+  public Converter dateFromIntConverter(TimeZone tz) {
+    return value -> structuredTypeDateTimeConverter.getDate((Integer) value, tz);
+  }
+
+  @SnowflakeJdbcInternalApi
+  public Converter timeFromStringConverter(SFBaseSession session) {
     return value -> {
       SnowflakeDateTimeFormat formatter =
           SnowflakeDateTimeFormat.fromSqlFormat(
@@ -205,8 +212,12 @@ public JsonStringToTypeConverter timeConverter(SFBaseSession session) {
     };
   }
 
+  public Converter timeFromIntConverter(int scale) {
+    return value -> structuredTypeDateTimeConverter.getTime((Long) value, scale);
+  }
+
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter timestampConverter(
+  public Converter timestampFromStringConverter(
       int columnSubType,
       int columnType,
       int scale,
@@ -225,22 +236,29 @@ public JsonStringToTypeConverter timestampConverter(
     };
   }
 
+  public Converter timestampFromStructConverter(
+      int columnType, int columnSubType, TimeZone tz, int scale) {
+    return value ->
+        structuredTypeDateTimeConverter.getTimestamp(
+            (Map<String, Object>) value, columnType, columnSubType, tz, scale);
+  }
+
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter structConverter(ObjectMapper objectMapper) {
+  public Converter structConverter(ObjectMapper objectMapper) {
     return value -> {
       try {
-        return objectMapper.readValue(value, Map.class);
+        return objectMapper.readValue((String) value, Map.class);
       } catch (JsonProcessingException e) {
         throw new SFException(e, ErrorCode.INVALID_STRUCT_DATA);
       }
     };
   }
 
   @SnowflakeJdbcInternalApi
-  public JsonStringToTypeConverter arrayConverter(ObjectMapper objectMapper) {
+  public Converter arrayConverter(ObjectMapper objectMapper) {
     return value -> {
       try {
-        return objectMapper.readValue(value, Map[].class);
+        return objectMapper.readValue((String) value, Map[].class);
       } catch (JsonProcessingException e) {
         throw new SFException(e, ErrorCode.INVALID_STRUCT_DATA);
       }

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

@@ -17,7 +17,6 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.TimeZone;
-import net.snowflake.client.core.json.Converters;
 import net.snowflake.client.core.structs.SQLDataCreationHelper;
 import net.snowflake.client.jdbc.FieldMetadata;
 import net.snowflake.client.jdbc.SnowflakeLoggedFeatureNotSupportedException;

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

@@ -16,19 +16,23 @@
 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.ArrowVectorConverter;
 import net.snowflake.client.core.arrow.StructConverter;
 import net.snowflake.client.core.arrow.VarCharConverter;
-import net.snowflake.client.core.json.Converters;
 import net.snowflake.client.core.structs.StructureTypeHelper;
 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 +42,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,16 +369,29 @@ public Converters getConverters() {
   }
 
   @Override
+  @SnowflakeJdbcInternalApi
+  public SQLInput createSqlInputForColumn(Object input, int columnIndex, SFBaseSession session) {
+    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 {
     return converters.getStructuredTypeDateTimeConverter().getDate((int) object, tz);
   }
 
   @Override
+  @SnowflakeJdbcInternalApi
   public Time convertToTime(Object object, int scale) throws SFException {
     return converters.getStructuredTypeDateTimeConverter().getTime((long) object, scale);
   }
 
   @Override
+  @SnowflakeJdbcInternalApi
   public Timestamp convertToTimestamp(
       Object object, int columnType, int columnSubType, TimeZone tz, int scale) throws SFException {
     return converters
@@ -566,8 +584,120 @@ 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);
+    return getArrayInternal((List<Object>) obj, columnIndex);
+  }
+
+  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