SNOW-1848734: Fix for structured types as strings performance problems (

#1991) Co-authored-by: Mikołaj Kubik <[email protected]>
snowflakedb · Dec 10, 2024 · e793e61 · e793e61
1 parent 3525adf
commit e793e61
Show file tree

Hide file tree

Showing 10 changed files with 66 additions and 42 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.StructConverter;
 import net.snowflake.client.core.arrow.StructObjectWrapper;
 import net.snowflake.client.core.arrow.VarCharConverter;
 import net.snowflake.client.core.arrow.VectorTypeConverter;
@@ -564,6 +565,16 @@ public Timestamp getTimestamp(int columnIndex, TimeZone tz) throws SFException {
 
   @Override
   public Object getObject(int columnIndex) throws SFException {
+    return getObjectRepresentation(columnIndex, true);
+  }
+
+  @SnowflakeJdbcInternalApi
+  @Override
+  public Object getObjectWithoutString(int columnIndex) throws SFException {
+    return getObjectRepresentation(columnIndex, false);
+  }
+
+  private Object getObjectRepresentation(int columnIndex, boolean withString) throws SFException {
     int type = resultSetMetaData.getColumnType(columnIndex);
     if (type == SnowflakeUtil.EXTRA_TYPES_VECTOR) {
       return getString(columnIndex);
@@ -575,13 +586,28 @@ public Object getObject(int columnIndex) throws SFException {
     converter.setUseSessionTimezone(useSessionTimezone);
     converter.setSessionTimeZone(sessionTimeZone);
     Object obj = converter.toObject(index);
+    if (obj == null) {
+      return null;
+    }
+    String jsonString = withString ? converter.toString(index) : null;
     boolean isStructuredType = resultSetMetaData.isStructuredTypeColumn(columnIndex);
     if (isVarcharConvertedStruct(type, isStructuredType, converter)) {
-      if (obj != null) {
-        return new StructObjectWrapper((String) obj, createJsonSqlInput(columnIndex, obj));
-      }
+      return new StructObjectWrapper(jsonString, createJsonSqlInput(columnIndex, obj));
+    } else if (converter instanceof StructConverter) {
+      return new StructObjectWrapper(
+          jsonString, createArrowSqlInput(columnIndex, (Map<String, Object>) obj));
+    } else {
+      return new StructObjectWrapper(jsonString, obj);
+    }
+  }
+
+  private SQLInput createArrowSqlInput(int columnIndex, Map<String, Object> input)
+      throws SFException {
+    if (input == null) {
+      return null;
     }
-    return obj;
+    return new ArrowSqlInput(
+        input, session, converters, resultSetMetaData.getColumnFields(columnIndex));
   }
 
   private boolean isVarcharConvertedStruct(
@@ -620,11 +646,8 @@ public Array getArray(int columnIndex) throws SFException {
     if (converter instanceof VarCharConverter) {
       return getJsonArray((String) obj, columnIndex);
     } else if (converter instanceof ArrayConverter || converter instanceof VectorTypeConverter) {
-      StructObjectWrapper structObjectWrapper = (StructObjectWrapper) obj;
-      return getArrowArray(
-          structObjectWrapper.getJsonString(),
-          (List<Object>) structObjectWrapper.getObject(),
-          columnIndex);
+      String jsonString = converter.toString(index);
+      return getArrowArray(jsonString, (List<Object>) obj, columnIndex);
     } else {
       throw new SFException(queryId, ErrorCode.INVALID_STRUCT_DATA);
     }

diff --git a/src/main/java/net/snowflake/client/core/SFBaseResultSet.java b/src/main/java/net/snowflake/client/core/SFBaseResultSet.java
@@ -113,6 +113,9 @@ public abstract class SFBaseResultSet {
 
   public abstract Object getObject(int columnIndex) throws SFException;
 
+  @SnowflakeJdbcInternalApi
+  public abstract Object getObjectWithoutString(int columnIndex) throws SFException;
+
   public Array getArray(int columnIndex) throws SFException {
     throw new UnsupportedOperationException();
   }

diff --git a/src/main/java/net/snowflake/client/core/SFJsonResultSet.java b/src/main/java/net/snowflake/client/core/SFJsonResultSet.java
@@ -103,6 +103,11 @@ public Object getObject(int columnIndex) throws SFException {
     }
   }
 
+  @SnowflakeJdbcInternalApi
+  @Override
+  public Object getObjectWithoutString(int columnIndex) throws SFException {
+    return getObject(columnIndex);
+  }
   /**
    * Sometimes large BIGINTS overflow the java Long type. In these cases, return a BigDecimal type
    * instead.

diff --git a/src/main/java/net/snowflake/client/core/arrow/ArrayConverter.java b/src/main/java/net/snowflake/client/core/arrow/ArrayConverter.java
@@ -25,7 +25,7 @@ public ArrayConverter(ListVector valueVector, int vectorIndex, DataConversionCon
 
   @Override
   public Object toObject(int index) throws SFException {
-    return isNull(index) ? null : new StructObjectWrapper(toString(index), vector.getObject(index));
+    return isNull(index) ? null : vector.getObject(index);
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/MapConverter.java b/src/main/java/net/snowflake/client/core/arrow/MapConverter.java
@@ -2,7 +2,6 @@
 
 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;
@@ -33,15 +32,11 @@ public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
     }
-
     List<JsonStringHashMap<String, Object>> entriesList =
         (List<JsonStringHashMap<String, Object>>) vector.getObject(index);
-    Map<String, Object> map =
-        entriesList.stream()
-            .collect(
-                Collectors.toMap(
-                    entry -> entry.get("key").toString(), entry -> entry.get("value")));
-    return new StructObjectWrapper(toString(index), map);
+    return entriesList.stream()
+        .collect(
+            Collectors.toMap(entry -> entry.get("key").toString(), entry -> entry.get("value")));
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/StructConverter.java b/src/main/java/net/snowflake/client/core/arrow/StructConverter.java
@@ -21,9 +21,7 @@ public StructConverter(StructVector vector, int columnIndex, DataConversionConte
 
   @Override
   public Object toObject(int index) throws SFException {
-    return isNull(index)
-        ? null
-        : new StructObjectWrapper(toString(index), structVector.getObject(index));
+    return isNull(index) ? null : structVector.getObject(index);
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/core/arrow/VectorTypeConverter.java b/src/main/java/net/snowflake/client/core/arrow/VectorTypeConverter.java
@@ -27,8 +27,7 @@ public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
     }
-    Object object = vector.getObject(index);
-    return new StructObjectWrapper(object.toString(), object);
+    return vector.getObject(index);
   }
 
   @Override

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBaseResultSet.java
@@ -1398,8 +1398,11 @@ public <T> T getObject(int columnIndex, Class<T> type) throws SQLException {
             SnowflakeUtil.mapSFExceptionToSQLException(
                 () -> {
                   StructObjectWrapper structObjectWrapper =
-                      (StructObjectWrapper) sfBaseResultSet.getObject(columnIndex);
-                  return (SQLInput) createJsonSqlInput(columnIndex, structObjectWrapper);
+                      (StructObjectWrapper) sfBaseResultSet.getObjectWithoutString(columnIndex);
+                  if (structObjectWrapper == null) {
+                    return null;
+                  }
+                  return (SQLInput) structObjectWrapper.getObject();
                 });
         if (sqlInput == null) {
           return null;
@@ -1639,7 +1642,7 @@ public <T> Map<String, T> getMap(int columnIndex, Class<T> type) throws SQLExcep
     StructObjectWrapper structObjectWrapper =
         (StructObjectWrapper)
             SnowflakeUtil.mapSFExceptionToSQLException(
-                () -> sfBaseResultSet.getObject(columnIndex));
+                () -> sfBaseResultSet.getObjectWithoutString(columnIndex));
     if (structObjectWrapper == null) {
       return null;
     }

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java
@@ -28,8 +28,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.TimeZone;
-import net.snowflake.client.core.ArrowSqlInput;
-import net.snowflake.client.core.JsonSqlInput;
 import net.snowflake.client.core.QueryStatus;
 import net.snowflake.client.core.SFBaseResultSet;
 import net.snowflake.client.core.SFException;
@@ -274,18 +272,21 @@ public Object getObject(int columnIndex) throws SQLException {
         SnowflakeUtil.mapSFExceptionToSQLException(() -> sfBaseResultSet.getObject(columnIndex));
     if (object == null) {
       return null;
-    } else if (object instanceof JsonSqlInput) {
-      return ((JsonSqlInput) object).getText();
-    } else if (object instanceof StructObjectWrapper) {
-      return ((StructObjectWrapper) object).getJsonString();
-    } else if (object instanceof SfSqlArray) {
+    }
+    if (object instanceof SfSqlArray) {
       return ((SfSqlArray) object).getText();
-    } else if (object instanceof ArrowSqlInput) {
-      throw new SQLException(
-          "Arrow native struct couldn't be converted to String. To map to SqlData the method getObject(int columnIndex, Class type) should be used");
-    } else {
-      return object;
     }
+    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 object;
   }
 
   public Array getArray(int columnIndex) throws SQLException {

diff --git a/...test/java/net/snowflake/client/jdbc/structuredtypes/ResultSetStructuredTypesLatestIT.java b/...test/java/net/snowflake/client/jdbc/structuredtypes/ResultSetStructuredTypesLatestIT.java
@@ -263,10 +263,7 @@ public void testReturnAsArrayOfNullableFieldsInSqlData(ResultSetFormatType forma
             + "'date', null, 'bd', null, 'bytes', null, 'longValue', null)"
             + "::OBJECT(string VARCHAR, nullableIntValue INTEGER, nullableLongValue INTEGER, date DATE, bd DOUBLE, bytes BINARY, longValue INTEGER)",
         (resultSet) -> {
-          NullableFieldsSqlData result =
-              resultSet
-                  .unwrap(SnowflakeBaseResultSet.class)
-                  .getObject(1, NullableFieldsSqlData.class);
+          NullableFieldsSqlData result = resultSet.getObject(1, NullableFieldsSqlData.class);
           assertNull(result.getString());
           assertNull(result.getNullableIntValue());
           assertNull(result.getNullableLongValue());