SNOW-1856886 Narrow toString calculation to structured types only (#2007

) Co-authored-by: Przemyslaw Motacki <przemyslaw.motacki@snowflake.com>
snowflakedb · Jan 14, 2025 · e8e3bac · e8e3bac
1 parent 1d58f57
commit e8e3bac
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 40 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFArrowResultSet.java b/src/main/java/net/snowflake/client/core/SFArrowResultSet.java
@@ -26,6 +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.MapConverter;
 import net.snowflake.client.core.arrow.StructConverter;
 import net.snowflake.client.core.arrow.StructObjectWrapper;
 import net.snowflake.client.core.arrow.VarCharConverter;
@@ -574,10 +575,11 @@ public Object getObjectWithoutString(int columnIndex) throws SFException {
     return getObjectRepresentation(columnIndex, false);
   }
 
-  private Object getObjectRepresentation(int columnIndex, boolean withString) throws SFException {
+  private StructObjectWrapper getObjectRepresentation(int columnIndex, boolean withString)
+      throws SFException {
     int type = resultSetMetaData.getColumnType(columnIndex);
     if (type == SnowflakeUtil.EXTRA_TYPES_VECTOR) {
-      return getString(columnIndex);
+      return new StructObjectWrapper(getString(columnIndex), null);
     }
     ArrowVectorConverter converter = currentChunkIterator.getCurrentConverter(columnIndex - 1);
     int index = currentChunkIterator.getCurrentRowInRecordBatch();
@@ -589,15 +591,33 @@ private Object getObjectRepresentation(int columnIndex, boolean withString) thro
     if (obj == null) {
       return null;
     }
-    String jsonString = withString ? converter.toString(index) : null;
     boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
-    if (isVarcharConvertedStruct(type, isStructuredType, converter)) {
-      return new StructObjectWrapper(jsonString, createJsonSqlInput(columnIndex, obj));
-    } else if (converter instanceof StructConverter) {
-      return new StructObjectWrapper(
-          jsonString, createArrowSqlInput(columnIndex, (Map<String, Object>) obj));
+    if (isStructuredType) {
+      if (converter instanceof VarCharConverter) {
+        if (type == Types.STRUCT) {
+          JsonSqlInput jsonSqlInput = createJsonSqlInput(columnIndex, obj);
+          return new StructObjectWrapper(jsonSqlInput.getText(), jsonSqlInput);
+        } else if (type == Types.ARRAY) {
+          SfSqlArray sfArray = getJsonArray((String) obj, columnIndex);
+          return new StructObjectWrapper(sfArray.getText(), sfArray);
+        } else {
+          throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
+        }
+      } else if (converter instanceof StructConverter) {
+        String jsonString = withString ? converter.toString(index) : null;
+        return new StructObjectWrapper(
+            jsonString, createArrowSqlInput(columnIndex, (Map<String, Object>) obj));
+      } else if (converter instanceof MapConverter) {
+        String jsonString = withString ? converter.toString(index) : null;
+        return new StructObjectWrapper(jsonString, obj);
+      } else if (converter instanceof ArrayConverter || converter instanceof VectorTypeConverter) {
+        String jsonString = converter.toString(index);
+        return new StructObjectWrapper(jsonString, obj);
+      } else {
+        throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
+      }
     } else {
-      return new StructObjectWrapper(jsonString, obj);
+      return new StructObjectWrapper(null, obj);
     }
   }
 
@@ -610,12 +630,11 @@ private SQLInput createArrowSqlInput(int columnIndex, Map<String, Object> input)
         input, session, converters, resultSetMetaData.getColumnFields(columnIndex));
   }
 
-  private boolean isVarcharConvertedStruct(
-      int type, boolean isStructuredType, ArrowVectorConverter converter) {
-    return type == Types.STRUCT && isStructuredType && converter instanceof VarCharConverter;
+  private boolean isVarcharConvertedStruct(int type, ArrowVectorConverter converter) {
+    return (type == Types.STRUCT || type == Types.ARRAY) && converter instanceof VarCharConverter;
   }
 
-  private Object createJsonSqlInput(int columnIndex, Object obj) throws SFException {
+  private JsonSqlInput createJsonSqlInput(int columnIndex, Object obj) throws SFException {
     try {
       if (obj == null) {
         return null;

diff --git a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
@@ -94,7 +94,7 @@ public Object getObject(int columnIndex) throws SFException {
         }
       case Types.ARRAY:
         if (resultSetMetaData.isStructuredTypeColumn(columnIndex)) {
-          return getArray(columnIndex);
+          return new StructObjectWrapper((String) obj, getArray(columnIndex));
         } else {
           throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "data type: " + type);
         }

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
@@ -1643,7 +1643,7 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
         (StructObjectWrapper)
             SnowflakeUtil.mapSFExceptionToSQLException(
                 () -> sfBaseResultSet.getObjectWithoutString(columnIndex));
-    if (structObjectWrapper == null) {
+    if (structObjectWrapper == null || structObjectWrapper.getObject() == null) {
       return null;
     }
     Map<String, Object> map =
@@ -1835,22 +1835,4 @@ private <T> Map<String, Object> prepareMapWithValues(Object object, Class<T> typ
       throw new SFException(ErrorCode.INVALID_STRUCT_DATA, "Object couldn't be converted to map");
     }
   }
-
-  private Object createJsonSqlInput(int columnIndex, StructObjectWrapper obj) throws SFException {
-    try {
-      if (obj == null) {
-        return null;
-      }
-      JsonNode jsonNode = OBJECT_MAPPER.readTree(obj.getJsonString());
-      return new JsonSqlInput(
-          obj.getJsonString(),
-          jsonNode,
-          session,
-          sfBaseResultSet.getConverters(),
-          sfBaseResultSet.getMetaData().getColumnFields(columnIndex),
-          sfBaseResultSet.getSessionTimeZone());
-    } catch (JsonProcessingException e) {
-      throw new SFException(sfBaseResultSet.getQueryId(), e, ErrorCode.INVALID_STRUCT_DATA);
-    }
-  }
 }
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java
@@ -31,7 +31,6 @@
 import net.snowflake.client.core.QueryStatus;
 import net.snowflake.client.core.SFBaseResultSet;
 import net.snowflake.client.core.SFException;
-import net.snowflake.client.core.SfSqlArray;
 import net.snowflake.client.core.arrow.StructObjectWrapper;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
@@ -270,22 +269,21 @@ public Object getObject(int columnIndex) throws SQLException {
     raiseSQLExceptionIfResultSetIsClosed();
     Object object =
         SnowflakeUtil.mapSFExceptionToSQLException(() -> sfBaseResultSet.getObject(columnIndex));
+
     if (object == null) {
       return null;
     }
-    if (object instanceof SfSqlArray) {
-      return ((SfSqlArray) object).getText();
-    }
+
     if (object instanceof StructObjectWrapper) {
       StructObjectWrapper structObjectWrapper = (StructObjectWrapper) object;
       if (resultSetMetaData.isStructuredTypeColumn(columnIndex)
           && structObjectWrapper.getJsonString() != null) {
         return structObjectWrapper.getJsonString();
       }
-      if (structObjectWrapper.getObject() != null) {
-        return structObjectWrapper.getObject();
-      }
+
+      return structObjectWrapper.getObject();
     }
+
     return object;
   }
 

diff --git a/...test/java/net/snowflake/client/jdbc/structuredtypes/ResultSetStructuredTypesLatestIT.java b/...test/java/net/snowflake/client/jdbc/structuredtypes/ResultSetStructuredTypesLatestIT.java
@@ -635,6 +635,19 @@ public void testMapIntegerArray(ResultSetFormatType format) throws SQLException
         format);
   }
 
+  @ParameterizedTest
+  @ArgumentsSource(ResultFormatProvider.class)
+  @DontRunOnGithubActions
+  public void testMapIntegerArrayGetObject(ResultSetFormatType format) throws SQLException {
+    withFirstRow(
+        "SELECT ARRAY_CONSTRUCT(10, 20, 30)::ARRAY(INTEGER)",
+        (resultSet) -> {
+          Object resultArray = resultSet.getObject(1);
+          TestUtil.assertEqualsIgnoringWhitespace("[10,20,30]", (String) resultArray);
+        },
+        format);
+  }
+
   @ParameterizedTest
   @ArgumentsSource(ResultFormatProvider.class)
   @DontRunOnGithubActions