SNOW-1847626: Add support for contains_null to ArrayType

CHANGELOG.md

@@ -22,6 +22,7 @@
 - Added support for `DataFrame.map`.
 - Added support for `DataFrame.from_dict` and `DataFrame.from_records`.
 - Added support for mixed case field names in struct type columns.
+- Added support for `contains_null` parameter to ArrayType.
 
 #### Improvements
 - Improve performance of `DataFrame.map`, `Series.apply` and `Series.map` methods by mapping numpy functions to snowpark functions if possible.

src/snowflake/snowpark/_internal/type_utils.py

@@ -141,6 +141,7 @@ def convert_metadata_to_sp_type(
             return ArrayType(
                 convert_metadata_to_sp_type(metadata.fields[0], max_string_size),
                 structured=True,
+                contains_null=metadata.fields[0]._is_nullable,
             )
         elif column_type_name == "MAP":
             assert (
@@ -285,7 +286,8 @@ def convert_sp_to_sf_type(datatype: DataType) -> str:
         return "BINARY"
     if isinstance(datatype, ArrayType):
         if datatype.structured:
-            return f"ARRAY({convert_sp_to_sf_type(datatype.element_type)})"
+            nullable = "" if datatype.contains_null else " NOT NULL"
+            return f"ARRAY({convert_sp_to_sf_type(datatype.element_type)}{nullable})"
         else:
             return "ARRAY"
     if isinstance(datatype, MapType):

src/snowflake/snowpark/types.py

@@ -335,9 +335,11 @@ def __init__(
         self,
         element_type: Optional[DataType] = None,
         structured: bool = False,
+        contains_null: bool = True,
     ) -> None:
         self.structured = structured
         self.element_type = element_type if element_type else StringType()
+        self.contains_null = contains_null
 
     def __repr__(self) -> str:
         return f"ArrayType({repr(self.element_type) if self.element_type else ''})"

tests/integ/scala/test_datatype_suite.py

@@ -991,6 +991,33 @@ def test_structured_type_print_schema(
     )
 
 
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="local testing does not fully support structured types yet.",
+)
+def test_structured_array_contains_null(
+    structured_type_session, structured_type_support
+):
+    if not structured_type_support:
+        pytest.skip("Test requires structured type support.")
+
+    # SNOW-1862947 create DDL test once save as table supported
+    array_df = structured_type_session.sql(
+        "select [1, 2, 3] :: ARRAY(INT NOT NULL) as A, [1, 2, 3] :: ARRAY(INT) as A_N"
+    )
+    expected_schema = StructType(
+        [
+            StructField(
+                "A", ArrayType(LongType(), structured=True, contains_null=False)
+            ),
+            StructField(
+                "A_N", ArrayType(LongType(), structured=True, contains_null=True)
+            ),
+        ]
+    )
+    assert array_df.schema == expected_schema
+
+
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="local testing does not fully support structured types yet.",

tests/unit/test_types.py

@@ -963,6 +963,15 @@ def test_convert_sp_to_sf_type():
     )
     assert convert_sp_to_sf_type(BinaryType()) == "BINARY"
     assert convert_sp_to_sf_type(ArrayType()) == "ARRAY"
+    assert (
+        convert_sp_to_sf_type(ArrayType(IntegerType(), structured=True)) == "ARRAY(INT)"
+    )
+    assert (
+        convert_sp_to_sf_type(
+            ArrayType(IntegerType(), structured=True, contains_null=False)
+        )
+        == "ARRAY(INT NOT NULL)"
+    )
     assert convert_sp_to_sf_type(MapType()) == "OBJECT"
     assert convert_sp_to_sf_type(StructType()) == "OBJECT"
     assert convert_sp_to_sf_type(VariantType()) == "VARIANT"