SNOW-1314319 Read Structured Types From Query Results (#96)

* fix type name * array v2 * map type * object * update JDBC * support structed array * map type * support map type * struct type * structure type * update jdbc * fix test * fix test * fix bc * temporary disable some sproc tests
snowflakedb · May 1, 2024 · fcc29b2 · fcc29b2
1 parent dfc1c78
commit fcc29b2
Show file tree

Hide file tree

Showing 13 changed files with 390 additions and 15 deletions.
diff --git a/fips-pom.xml b/fips-pom.xml
@@ -37,7 +37,7 @@
         <scala.version>2.12.18</scala.version>
         <scala.compat.version>2.12</scala.compat.version>
         <spec2.version>4.2.0</spec2.version>
-        <snowflake.jdbc.version>3.14.4</snowflake.jdbc.version>
+        <snowflake.jdbc.version>3.16.0</snowflake.jdbc.version>
         <version.scala.binary>${scala.compat.version}</version.scala.binary>
         <doctitle>Snowpark ${project.version}</doctitle>
         <bouncycastle.version>1.64</bouncycastle.version>

diff --git a/pom.xml b/pom.xml
@@ -37,7 +37,7 @@
         <scala.version>2.12.18</scala.version>
         <scala.compat.version>2.12</scala.compat.version>
         <spec2.version>4.2.0</spec2.version>
-        <snowflake.jdbc.version>3.14.4</snowflake.jdbc.version>
+        <snowflake.jdbc.version>3.16.0</snowflake.jdbc.version>
         <version.scala.binary>${scala.compat.version}</version.scala.binary>
         <doctitle>Snowpark ${project.version}</doctitle>
         <scoverage.plugin.version>1.4.11</scoverage.plugin.version>

diff --git a/src/main/scala/com/snowflake/snowpark/internal/ServerConnection.scala b/src/main/scala/com/snowflake/snowpark/internal/ServerConnection.scala
@@ -26,18 +26,20 @@ import com.snowflake.snowpark.internal.ParameterUtils.{
 import com.snowflake.snowpark.internal.Utils.PackageNameDelimiter
 import com.snowflake.snowpark.internal.analyzer.{Attribute, Query, SnowflakePlan}
 import net.snowflake.client.jdbc.{
+  FieldMetadata,
   SnowflakeConnectString,
   SnowflakeConnectionV1,
   SnowflakeReauthenticationRequest,
   SnowflakeResultSet,
-  SnowflakeSQLException,
+  SnowflakeResultSetMetaData,
   SnowflakeStatement
 }
 import com.snowflake.snowpark.types._
 import net.snowflake.client.core.QueryStatus
 
 import scala.collection.mutable
 import scala.reflect.runtime.universe.TypeTag
+import scala.collection.JavaConverters._
 
 private[snowpark] case class QueryResult(
     rows: Option[Array[Row]],
@@ -55,6 +57,11 @@ private[snowpark] object ServerConnection {
 
   def convertResultMetaToAttribute(meta: ResultSetMetaData): Seq[Attribute] =
     (1 to meta.getColumnCount).map(index => {
+      val fieldMetadata = meta
+        .asInstanceOf[SnowflakeResultSetMetaData]
+        .getColumnFields(index)
+        .asScala
+        .toList
       val columnName = analyzer.quoteNameWithoutUpperCasing(meta.getColumnLabel(index))
       val dataType = meta.getColumnType(index)
       val fieldSize = meta.getPrecision(index)
@@ -64,7 +71,8 @@ private[snowpark] object ServerConnection {
       // This field is useful for snowflake types that are not JDBC types like
       // variant, object and array
       val columnTypeName = meta.getColumnTypeName(index)
-      val columnType = getDataType(dataType, columnTypeName, fieldSize, fieldScale, isSigned)
+      val columnType =
+        getDataType(dataType, columnTypeName, fieldSize, fieldScale, isSigned, fieldMetadata)
 
       Attribute(columnName, columnType, nullable)
     })
@@ -74,11 +82,61 @@ private[snowpark] object ServerConnection {
       columnTypeName: String,
       precision: Int,
       scale: Int,
-      signed: Boolean): DataType = {
+      signed: Boolean,
+      field: List[FieldMetadata] = List.empty): DataType = {
     columnTypeName match {
-      case "ARRAY" => ArrayType(StringType)
+      case "ARRAY" =>
+        if (field.isEmpty) {
+          ArrayType(StringType)
+        } else {
+          StructuredArrayType(
+            getDataType(
+              field.head.getType,
+              field.head.getTypeName,
+              field.head.getPrecision,
+              field.head.getScale,
+              signed = true, // no sign info in the fields
+              field.head.getFields.asScala.toList),
+            field.head.isNullable)
+        }
       case "VARIANT" => VariantType
-      case "OBJECT" => MapType(StringType, StringType)
+      case "OBJECT" =>
+        if (field.isEmpty) {
+          MapType(StringType, StringType)
+        } else if (field.size == 2 && field.head.getName.isEmpty) {
+          // Map
+          StructuredMapType(
+            getDataType(
+              field.head.getType,
+              field.head.getTypeName,
+              field.head.getPrecision,
+              field.head.getScale,
+              signed = true,
+              field.head.getFields.asScala.toList),
+            getDataType(
+              field(1).getType,
+              field(1).getTypeName,
+              field(1).getPrecision,
+              field(1).getScale,
+              signed = true,
+              field(1).getFields.asScala.toList),
+            field(1).isNullable)
+        } else {
+          // object
+          StructType(
+            field.map(
+              f =>
+                StructField(
+                  f.getName,
+                  getDataType(
+                    f.getType,
+                    f.getTypeName,
+                    f.getPrecision,
+                    f.getScale,
+                    signed = true,
+                    f.getFields.asScala.toList),
+                  f.isNullable)))
+        }
       case "GEOGRAPHY" => GeographyType
       case "GEOMETRY" => GeometryType
       case _ => getTypeFromJDBCType(sqlType, precision, scale, signed)

diff --git a/src/main/scala/com/snowflake/snowpark/internal/analyzer/DataTypeMapper.scala b/src/main/scala/com/snowflake/snowpark/internal/analyzer/DataTypeMapper.scala
@@ -90,7 +90,6 @@ object DataTypeMapper {
       dataType match {
         case GeographyType => "TRY_TO_GEOGRAPHY(NULL)"
         case GeometryType => "TRY_TO_GEOMETRY(NULL)"
-        case ArrayType(_) => "PARSE_JSON('NULL')::ARRAY"
         case _ => "NULL :: " + convertToSFType(dataType)
       }
     } else {
@@ -102,8 +101,8 @@ object DataTypeMapper {
         case DateType => "date('2020-9-16')"
         case TimeType => "to_time('04:15:29.999')"
         case TimestampType => "to_timestamp_ntz('2020-09-16 06:30:00')"
-        case _: ArrayType => "to_array(0)"
-        case _: MapType => "to_object(parse_json('0'))"
+        case _: ArrayType => "[]::" + convertToSFType(dataType)
+        case _: MapType => "{}::" + convertToSFType(dataType)
         case VariantType => "to_variant(0)"
         case GeographyType => "to_geography('POINT(-122.35 37.55)')"
         case GeometryType => "to_geometry('POINT(-122.35 37.55)')"

diff --git a/src/main/scala/com/snowflake/snowpark/types/ArrayType.scala b/src/main/scala/com/snowflake/snowpark/types/ArrayType.scala
@@ -9,4 +9,27 @@ case class ArrayType(elementType: DataType) extends DataType {
   override def toString: String = {
     s"ArrayType[${elementType.toString}]"
   }
+
+  override def schemaString: String =
+    s"Array"
+}
+
+/* temporary solution for Structured and Semi Structured data types.
+Two types will be merged in the future BCR. */
+private[snowpark] class StructuredArrayType(
+    override val elementType: DataType,
+    val nullable: Boolean)
+    extends ArrayType(elementType) {
+  override def toString: String = {
+    s"ArrayType[${elementType.toString} nullable = $nullable]"
+  }
+
+  override def schemaString: String =
+    s"Array[${elementType.schemaString} nullable = $nullable]"
+}
+
+private[snowpark] object StructuredArrayType {
+
+  def apply(elementType: DataType, nullable: Boolean): StructuredArrayType =
+    new StructuredArrayType(elementType, nullable)
 }
diff --git a/src/main/scala/com/snowflake/snowpark/types/DataType.scala b/src/main/scala/com/snowflake/snowpark/types/DataType.scala
@@ -18,6 +18,8 @@ abstract class DataType {
    * @since 0.1.0
    */
   override def toString: String = typeName
+
+  private[snowpark] def schemaString: String = toString
 }
 
 private[snowpark] abstract class AtomicType extends DataType
diff --git a/src/main/scala/com/snowflake/snowpark/types/GeographyType.scala b/src/main/scala/com/snowflake/snowpark/types/GeographyType.scala
@@ -8,4 +8,6 @@ object GeographyType extends DataType {
   override def toString: String = {
     s"GeographyType"
   }
+
+  override def schemaString: String = s"Geography"
 }
diff --git a/src/main/scala/com/snowflake/snowpark/types/GeometryType.scala b/src/main/scala/com/snowflake/snowpark/types/GeometryType.scala
@@ -8,4 +8,6 @@ object GeometryType extends DataType {
   override def toString: String = {
     s"GeometryType"
   }
+
+  override def schemaString: String = s"Geometry"
 }
diff --git a/src/main/scala/com/snowflake/snowpark/types/MapType.scala b/src/main/scala/com/snowflake/snowpark/types/MapType.scala
@@ -9,4 +9,25 @@ case class MapType(keyType: DataType, valueType: DataType) extends DataType {
   override def toString: String = {
     s"MapType[${keyType.toString}, ${valueType.toString}]"
   }
+
+  override private[snowpark] def schemaString =
+    s"Map"
+}
+
+private[snowpark] class StructuredMapType(
+    override val keyType: DataType,
+    override val valueType: DataType,
+    val isValueNullable: Boolean)
+    extends MapType(keyType, valueType) {
+  override def toString: String = {
+    s"MapType[${keyType.toString}, ${valueType.toString} nullable = $isValueNullable]"
+  }
+
+  override private[snowpark] def schemaString =
+    s"Map[${keyType.schemaString}, ${valueType.schemaString} nullable = $isValueNullable]"
+}
+
+private[snowpark] object StructuredMapType {
+  def apply(keyType: DataType, valueType: DataType, isValueType: Boolean): StructuredMapType =
+    new StructuredMapType(keyType, valueType, isValueType)
 }
diff --git a/src/main/scala/com/snowflake/snowpark/types/StructType.scala b/src/main/scala/com/snowflake/snowpark/types/StructType.scala
@@ -65,6 +65,8 @@ case class StructType(fields: Array[StructField] = Array())
   override def toString: String =
     s"StructType[${fields.map(_.toString).mkString(", ")}]"
 
+  override private[snowpark] def schemaString: String = "Struct"
+
   /**
    * Appends a new [[StructField]] to the end of this object.
    * @since 0.1.0
@@ -168,7 +170,7 @@ case class StructField(
 
   private[types] def treeString(layer: Int): String = {
     val prepended: String = (1 to (1 + 2 * layer)).map(x => " ").mkString + "|--"
-    val body: String = s"$name: ${dataType.typeName} (nullable = $nullable)\n" +
+    val body: String = s"$name: ${dataType.schemaString} (nullable = $nullable)\n" +
       (dataType match {
         case st: StructType => st.treeString(layer + 1)
         case _ => ""

diff --git a/src/main/scala/com/snowflake/snowpark/types/package.scala b/src/main/scala/com/snowflake/snowpark/types/package.scala
@@ -63,11 +63,26 @@ package object types {
       case TimeType => "TIME"
       case TimestampType => "TIMESTAMP"
       case BinaryType => "BINARY"
+      case sa: StructuredArrayType =>
+        val nullable = if (sa.nullable) "" else " not null"
+        s"ARRAY(${convertToSFType(sa.elementType)}$nullable)"
+      case sm: StructuredMapType =>
+        val isValueNullable = if (sm.isValueNullable) "" else " not null"
+        s"MAP(${convertToSFType(sm.keyType)}, ${convertToSFType(sm.valueType)}$isValueNullable)"
+      case StructType(fields) =>
+        val fieldStr = fields
+          .map(
+            field =>
+              s"${field.name} ${convertToSFType(field.dataType)} " +
+                (if (field.nullable) "" else "not null"))
+          .mkString(",")
+        s"OBJECT($fieldStr)"
       case ArrayType(_) => "ARRAY"
       case MapType(_, _) => "OBJECT"
       case VariantType => "VARIANT"
       case GeographyType => "GEOGRAPHY"
       case GeometryType => "GEOMETRY"
+      case StructType(_) => "OBJECT"
       case _ =>
         throw new UnsupportedOperationException(s"Unsupported data type: ${dataType.typeName}")
     }