object

snowflakedb · Apr 8, 2024 · 7e5d47e · 7e5d47e
1 parent 7cf91c5
commit 7e5d47e
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 1 deletion.
diff --git a/src/main/scala/com/snowflake/snowpark/internal/ServerConnection.scala b/src/main/scala/com/snowflake/snowpark/internal/ServerConnection.scala
@@ -124,7 +124,20 @@ private[snowpark] object ServerConnection {
               signed = true,
               field(1).getFields.asScala.toList))
         } else {
-          null // object
+          // 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

diff --git a/src/main/scala/com/snowflake/snowpark/types/package.scala b/src/main/scala/com/snowflake/snowpark/types/package.scala
@@ -68,6 +68,7 @@ package object types {
       case VariantType => "VARIANT"
       case GeographyType => "GEOGRAPHY"
       case GeometryType => "GEOMETRY"
+      case StructType(_) => "OBJECT"
       case _ =>
         throw new UnsupportedOperationException(s"Unsupported data type: ${dataType.typeName}")
     }

diff --git a/src/test/scala/com/snowflake/snowpark_test/DataTypeSuite.scala b/src/test/scala/com/snowflake/snowpark_test/DataTypeSuite.scala
@@ -236,4 +236,37 @@ class DataTypeSuite extends SNTestBase {
            | |--MAP0: MapType[String, String] (nullable = true)
            |""".stripMargin)
   }
+
+  test("ObjectType v2") {
+    val query =
+      // scalastyle:off
+      """SELECT
+        |  {'a': 1, 'b': 'a'} :: OBJECT(a VARCHAR, b NUMBER) as object1,
+        |  {'a': 1, 'b': [1,2,3,4]} :: OBJECT(a VARCHAR, b ARRAY(NUMBER)) as object2,
+        |  {'a': 1, 'b': [1,2,3,4], 'c': {'1':'a'}} :: OBJECT(a VARCHAR, b ARRAY(NUMBER), c MAP(NUMBER, VARCHAR)) as object3,
+        |  {'a': {'b': {'c': 1}}} :: OBJECT(a OBJECT(b OBJECT(c NUMBER))) as object4
+        |""".stripMargin
+    // scalastyle:on
+    val df = session.sql(query)
+    assert(
+      TestUtils.treeString(df.schema, 0) ==
+        // scalastyle:off
+        s"""root
+           | |--OBJECT1: StructType[StructField(A, String, Nullable = true), StructField(B, Long, Nullable = true)] (nullable = true)
+           |   |--A: String (nullable = true)
+           |   |--B: Long (nullable = true)
+           | |--OBJECT2: StructType[StructField(A, String, Nullable = true), StructField(B, ArrayType[Long], Nullable = true)] (nullable = true)
+           |   |--A: String (nullable = true)
+           |   |--B: ArrayType[Long] (nullable = true)
+           | |--OBJECT3: StructType[StructField(A, String, Nullable = true), StructField(B, ArrayType[Long], Nullable = true), StructField(C, MapType[Long, String], Nullable = true)] (nullable = true)
+           |   |--A: String (nullable = true)
+           |   |--B: ArrayType[Long] (nullable = true)
+           |   |--C: MapType[Long, String] (nullable = true)
+           | |--OBJECT4: StructType[StructField(A, StructType[StructField(B, StructType[StructField(C, Long, Nullable = true)], Nullable = true)], Nullable = true)] (nullable = true)
+           |   |--A: StructType[StructField(B, StructType[StructField(C, Long, Nullable = true)], Nullable = true)] (nullable = true)
+           |     |--B: StructType[StructField(C, Long, Nullable = true)] (nullable = true)
+           |       |--C: Long (nullable = true)
+           |""".stripMargin)
+    // scalastyle:on
+  }
 }