From 8ab9fe7f623eebe508a8b2460706c1ee0bd0d41e Mon Sep 17 00:00:00 2001
From: Bing Li <bing.li@snowflake.com>
Date: Wed, 10 Apr 2024 15:50:20 -0700
Subject: [PATCH] structure type

---
 .../snowpark_test/DataTypeSuite.scala         | 62 +++++++++++++++++--
 1 file changed, 57 insertions(+), 5 deletions(-)

diff --git a/src/test/scala/com/snowflake/snowpark_test/DataTypeSuite.scala b/src/test/scala/com/snowflake/snowpark_test/DataTypeSuite.scala
index c82c5881..ec55cc01 100644
--- a/src/test/scala/com/snowflake/snowpark_test/DataTypeSuite.scala
+++ b/src/test/scala/com/snowflake/snowpark_test/DataTypeSuite.scala
@@ -128,22 +128,20 @@ class DataTypeSuite extends SNTestBase {
             StructField("col10", DoubleType),
             StructField("col11", DecimalType(10, 1)))))))
 
-    schema.printTreeString()
-
     assert(
       TestUtils.treeString(schema, 0) ==
         s"""root
            | |--COL1: Binary (nullable = true)
            | |--COL2: Boolean (nullable = true)
-           | |--COL14: Struct (nullable = false)
+           | |--COL14: StructType (nullable = false)
            |   |--COL15: Timestamp (nullable = false)
            | |--COL3: Date (nullable = false)
-           | |--COL4: Struct (nullable = true)
+           | |--COL4: StructType (nullable = true)
            |   |--COL5: Byte (nullable = true)
            |   |--COL6: Short (nullable = true)
            |   |--COL7: Integer (nullable = false)
            |   |--COL8: Long (nullable = true)
-           |   |--COL12: Struct (nullable = false)
+           |   |--COL12: StructType (nullable = false)
            |     |--COL13: String (nullable = true)
            |   |--COL9: Float (nullable = true)
            |   |--COL10: Double (nullable = true)
@@ -379,5 +377,59 @@ class DataTypeSuite extends SNTestBase {
            |       |--C: Long (nullable = true)
            |""".stripMargin)
     // scalastyle:on
+
+    // schema query: not null
+    val query2 =
+    // scalastyle:off
+      """SELECT
+        |  {'a': 1, 'b': 'a'} :: OBJECT(a VARCHAR not null, b NUMBER) as object1,
+        |  {'a': 1, 'b': [1,2,3,4]} :: OBJECT(a VARCHAR, b ARRAY(NUMBER not null) not null) as object2,
+        |  {'a': 1, 'b': [1,2,3,4], 'c': {'1':'a'}} :: OBJECT(a VARCHAR, b ARRAY(NUMBER), c MAP(NUMBER, VARCHAR not null) not null) as object3,
+        |  {'a': {'b': {'c': 1}}} :: OBJECT(a OBJECT(b OBJECT(c NUMBER not null) not null) not null) as object4
+        |""".stripMargin
+    // scalastyle:on
+
+    val df2 = session.sql(query2)
+    assert(
+      TestUtils.treeString(df2.schema, 0) ==
+        // scalastyle:off
+        s"""root
+           | |--OBJECT1: StructType (nullable = true)
+           |   |--A: String (nullable = false)
+           |   |--B: Long (nullable = true)
+           | |--OBJECT2: StructType (nullable = true)
+           |   |--A: String (nullable = true)
+           |   |--B: ArrayType[Long nullable = false] (nullable = false)
+           | |--OBJECT3: StructType (nullable = true)
+           |   |--A: String (nullable = true)
+           |   |--B: ArrayType[Long nullable = true] (nullable = true)
+           |   |--C: MapType[Long, String nullable = false] (nullable = false)
+           | |--OBJECT4: StructType (nullable = true)
+           |   |--A: StructType (nullable = false)
+           |     |--B: StructType (nullable = false)
+           |       |--C: Long (nullable = false)
+           |""".stripMargin)
+    // scalastyle:on
+
+    assert(
+      TestUtils.treeString(df2.select("*").schema, 0) ==
+        // scalastyle:off
+        s"""root
+           | |--OBJECT1: StructType (nullable = true)
+           |   |--A: String (nullable = false)
+           |   |--B: Long (nullable = true)
+           | |--OBJECT2: StructType (nullable = true)
+           |   |--A: String (nullable = true)
+           |   |--B: ArrayType[Long nullable = false] (nullable = false)
+           | |--OBJECT3: StructType (nullable = true)
+           |   |--A: String (nullable = true)
+           |   |--B: ArrayType[Long nullable = true] (nullable = true)
+           |   |--C: MapType[Long, String nullable = false] (nullable = false)
+           | |--OBJECT4: StructType (nullable = true)
+           |   |--A: StructType (nullable = false)
+           |     |--B: StructType (nullable = false)
+           |       |--C: Long (nullable = false)
+           |""".stripMargin)
+    // scalastyle:on
   }
 }