SNOW-1812949 add get object and get bytes support for native arrow structured types

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

@@ -26,7 +26,7 @@
 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.StructObjectWrapper;
 import net.snowflake.client.core.arrow.VarCharConverter;
 import net.snowflake.client.core.arrow.VectorTypeConverter;
 import net.snowflake.client.core.json.Converters;
@@ -576,12 +576,11 @@ public Object getObject(int columnIndex) throws SFException {
     converter.setSessionTimeZone(sessionTimeZone);
     Object obj = converter.toObject(index);
     boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
-    if (type == Types.STRUCT && isStructuredType) {
-      if (converter instanceof VarCharConverter) {
-        return createJsonSqlInput(columnIndex, obj);
-      } else if (converter instanceof StructConverter) {
-        return createArrowSqlInput(columnIndex, (Map<String, Object>) obj);
+    if (type == Types.STRUCT && isStructuredType && converter instanceof VarCharConverter) {
+      if (obj == null) {
+        return null;
       }
+      return new StructObjectWrapper((String) obj, createJsonSqlInput(columnIndex, obj));
     }
     return obj;
   }
@@ -605,15 +604,6 @@ private Object createJsonSqlInput(int columnIndex, Object obj) throws SFExceptio
     }
   }
 
-  private Object createArrowSqlInput(int columnIndex, Map<String, Object> input)
-      throws SFException {
-    if (input == null) {
-      return null;
-    }
-    return new ArrowSqlInput(
-        input, session, converters, resultSetMetaData.getColumnFields(columnIndex));
-  }
-
   @Override
   public Array getArray(int columnIndex) throws SFException {
     ArrowVectorConverter converter = currentChunkIterator.getCurrentConverter(columnIndex - 1);
@@ -625,16 +615,19 @@ public Array getArray(int columnIndex) throws SFException {
     }
     if (converter instanceof VarCharConverter) {
       return getJsonArray((String) obj, columnIndex);
-    } else if (converter instanceof ArrayConverter) {
-      return getArrowArray((List<Object>) obj, columnIndex);
-    } else if (converter instanceof VectorTypeConverter) {
-      return getArrowArray((List<Object>) obj, columnIndex);
+    } else if (converter instanceof ArrayConverter || converter instanceof VectorTypeConverter) {
+      StructObjectWrapper structObjectWrapper = (StructObjectWrapper) obj;
+      return getArrowArray(
+          structObjectWrapper.getJsonString(),
+          (List<Object>) structObjectWrapper.getObject(),
+          columnIndex);
     } else {
       throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
     }
   }
 
-  private SfSqlArray getArrowArray(List<Object> elements, int columnIndex) throws SFException {
+  private SfSqlArray getArrowArray(String text, List<Object> elements, int columnIndex)
+      throws SFException {
     try {
       List<FieldMetadata> fieldMetadataList = resultSetMetaData.getColumnFields(columnIndex);
       if (fieldMetadataList.size() != 1) {
@@ -651,62 +644,74 @@ private SfSqlArray getArrowArray(List<Object> elements, int columnIndex) throws
       switch (columnType) {
         case Types.INTEGER:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.integerConverter(columnType))
                   .toArray(Integer[]::new));
         case Types.SMALLINT:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.smallIntConverter(columnType))
                   .toArray(Short[]::new));
         case Types.TINYINT:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.tinyIntConverter(columnType))
                   .toArray(Byte[]::new));
         case Types.BIGINT:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.bigIntConverter(columnType)).toArray(Long[]::new));
         case Types.DECIMAL:
         case Types.NUMERIC:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.bigDecimalConverter(columnType))
                   .toArray(BigDecimal[]::new));
         case Types.CHAR:
         case Types.VARCHAR:
         case Types.LONGNVARCHAR:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.varcharConverter(columnType, columnSubType, scale))
                   .toArray(String[]::new));
         case Types.BINARY:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.bytesConverter(columnType, scale))
                   .toArray(Byte[][]::new));
         case Types.FLOAT:
         case Types.REAL:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.floatConverter(columnType)).toArray(Float[]::new));
         case Types.DOUBLE:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.doubleConverter(columnType))
                   .toArray(Double[]::new));
         case Types.DATE:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.dateFromIntConverter(sessionTimeZone))
                   .toArray(Date[]::new));
         case Types.TIME:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.timeFromIntConverter(scale)).toArray(Time[]::new));
         case Types.TIMESTAMP:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(
                       elements,
@@ -715,13 +720,16 @@ private SfSqlArray getArrowArray(List<Object> elements, int columnIndex) throws
                   .toArray(Timestamp[]::new));
         case Types.BOOLEAN:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, converters.booleanConverter(columnType))
                   .toArray(Boolean[]::new));
         case Types.STRUCT:
-          return new SfSqlArray(columnSubType, mapAndConvert(elements, e -> e).toArray(Map[]::new));
+          return new SfSqlArray(
+              text, columnSubType, mapAndConvert(elements, e -> e).toArray(Map[]::new));
         case Types.ARRAY:
           return new SfSqlArray(
+              text,
               columnSubType,
               mapAndConvert(elements, e -> ((List) e).stream().toArray(Map[]::new))
                   .toArray(Map[][]::new));

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

@@ -294,67 +294,79 @@ protected SfSqlArray getJsonArray(String obj, int columnIndex) throws SFExceptio
       switch (columnType) {
         case Types.INTEGER:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().integerConverter(columnType))
                   .toArray(Integer[]::new));
         case Types.SMALLINT:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().smallIntConverter(columnType))
                   .toArray(Short[]::new));
         case Types.TINYINT:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().tinyIntConverter(columnType))
                   .toArray(Byte[]::new));
         case Types.BIGINT:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().bigIntConverter(columnType))
                   .toArray(Long[]::new));
         case Types.DECIMAL:
         case Types.NUMERIC:
           return new SfSqlArray(
+              obj,
               columnSubType,
               convertToFixedArray(
                   getStream(nodeElements, getConverters().bigDecimalConverter(columnType))));
         case Types.CHAR:
         case Types.VARCHAR:
         case Types.LONGNVARCHAR:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(
                       nodeElements,
                       getConverters().varcharConverter(columnType, columnSubType, scale))
                   .toArray(String[]::new));
         case Types.BINARY:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().bytesConverter(columnType, scale))
                   .toArray(Byte[][]::new));
         case Types.FLOAT:
         case Types.REAL:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().floatConverter(columnType))
                   .toArray(Float[]::new));
         case Types.DOUBLE:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().doubleConverter(columnType))
                   .toArray(Double[]::new));
         case Types.DATE:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().dateStringConverter(session))
                   .toArray(Date[]::new));
         case Types.TIME:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().timeFromStringConverter(session))
                   .toArray(Time[]::new));
         case Types.TIMESTAMP:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(
                       nodeElements,
@@ -364,16 +376,19 @@ protected SfSqlArray getJsonArray(String obj, int columnIndex) throws SFExceptio
                   .toArray(Timestamp[]::new));
         case Types.BOOLEAN:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().booleanConverter(columnType))
                   .toArray(Boolean[]::new));
         case Types.STRUCT:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().structConverter(OBJECT_MAPPER))
                   .toArray(Map[]::new));
         case Types.ARRAY:
           return new SfSqlArray(
+              obj,
               columnSubType,
               getStream(nodeElements, getConverters().arrayConverter(OBJECT_MAPPER))
                   .toArray(Map[][]::new));

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

@@ -15,6 +15,7 @@
 import java.sql.Types;
 import java.util.List;
 import java.util.TimeZone;
+import net.snowflake.client.core.arrow.StructObjectWrapper;
 import net.snowflake.client.core.json.Converters;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.FieldMetadata;
@@ -87,7 +88,7 @@ public Object getObject(int columnIndex) throws SFException {
 
       case Types.STRUCT:
         if (resultSetMetaData.isStructuredTypeColumn(columnIndex)) {
-          return getSqlInput((String) obj, columnIndex);
+          return new StructObjectWrapper((String) obj, getSqlInput((String) obj, columnIndex));
         } else {
           throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "data type: " + type);
         }

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

@@ -16,9 +16,16 @@
 @SnowflakeJdbcInternalApi
 public class SfSqlArray implements Array {
 
+  private String text;
   private int baseType;
   private Object elements;
 
+  public SfSqlArray(String text, int baseType, Object elements) {
+    this.text = text;
+    this.baseType = baseType;
+    this.elements = elements;
+  }
+
   public SfSqlArray(int baseType, Object elements) {
     this.baseType = baseType;
     this.elements = elements;
@@ -82,6 +89,13 @@ public ResultSet getResultSet(long index, int count, Map<String, Class<?>> map)
   public void free() throws SQLException {}
 
   public String getJsonString() throws SQLException {
+    if (text == null) {
+      text = buildJsonStringFromElements(elements);
+    }
+    return text;
+  }
+
+  private static String buildJsonStringFromElements(Object elements) throws SQLException {
     try {
       return SnowflakeUtil.mapJson(elements);
     } catch (JsonProcessingException e) {

src/main/java/net/snowflake/client/core/arrow/ArrayConverter.java

@@ -25,7 +25,12 @@ public ArrayConverter(ListVector valueVector, int vectorIndex, DataConversionCon
 
   @Override
   public Object toObject(int index) throws SFException {
-    return vector.getObject(index);
+    return isNull(index) ? null : new StructObjectWrapper(toString(index), vector.getObject(index));
+  }
+
+  @Override
+  public byte[] toBytes(int index) throws SFException {
+    return isNull(index) ? null : toString(index).getBytes();
   }
 
   @Override

src/main/java/net/snowflake/client/core/arrow/MapConverter.java

@@ -1,6 +1,8 @@
 package net.snowflake.client.core.arrow;
 
+import java.nio.charset.StandardCharsets;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 import net.snowflake.client.core.DataConversionContext;
 import net.snowflake.client.core.SFException;
@@ -28,11 +30,23 @@ public MapConverter(MapVector valueVector, int columnIndex, DataConversionContex
 
   @Override
   public Object toObject(int index) throws SFException {
+    if (isNull(index)) {
+      return null;
+    }
+
     List<JsonStringHashMap<String, Object>> entriesList =
         (List<JsonStringHashMap<String, Object>>) vector.getObject(index);
-    return entriesList.stream()
-        .collect(
-            Collectors.toMap(entry -> entry.get("key").toString(), entry -> entry.get("value")));
+    Map<String, Object> map =
+        entriesList.stream()
+            .collect(
+                Collectors.toMap(
+                    entry -> entry.get("key").toString(), entry -> entry.get("value")));
+    return new StructObjectWrapper(toString(index), map);
+  }
+
+  @Override
+  public byte[] toBytes(int index) throws SFException {
+    return isNull(index) ? null : toString(index).getBytes(StandardCharsets.UTF_8);
   }
 
   @Override

src/main/java/net/snowflake/client/core/arrow/StructConverter.java

@@ -21,7 +21,14 @@ public StructConverter(StructVector vector, int columnIndex, DataConversionConte
 
   @Override
   public Object toObject(int index) throws SFException {
-    return structVector.getObject(index);
+    return isNull(index)
+        ? null
+        : new StructObjectWrapper(toString(index), structVector.getObject(index));
+  }
+
+  @Override
+  public byte[] toBytes(int index) throws SFException {
+    return isNull(index) ? null : toString(index).getBytes();
   }
 
   @Override
@@ -32,9 +39,13 @@ public String toString(int index) throws SFException {
       SnowflakeType logicalType =
           ArrowVectorConverterUtil.getSnowflakeTypeFromFieldMetadata(fieldVector.getField());
       try {
-        ArrowVectorConverter converter =
-            ArrowVectorConverterUtil.initConverter(fieldVector, context, columnIndex);
-        builder.appendKeyValue(childName, converter.toString(index), logicalType);
+        if (fieldVector.isNull(index)) {
+          builder.appendKeyValue(childName, null, logicalType);
+        } else {
+          ArrowVectorConverter converter =
+              ArrowVectorConverterUtil.initConverter(fieldVector, context, columnIndex);
+          builder.appendKeyValue(childName, converter.toString(index), logicalType);
+        }
       } catch (SnowflakeSQLException e) {
         return structVector.getObject(index).toString();
       }