SNOW-1829870: Allow structured types to be enabled by default

src/snowflake/snowpark/_internal/analyzer/datatype_mapper.py

@@ -156,10 +156,16 @@ def to_sql(value: Any, datatype: DataType, from_values_statement: bool = False)
         return f"'{binascii.hexlify(bytes(value)).decode()}' :: BINARY"
 
     if isinstance(value, (list, tuple, array)) and isinstance(datatype, ArrayType):
-        return f"PARSE_JSON({str_to_sql(json.dumps(value, cls=PythonObjJSONEncoder))}) :: ARRAY"
+        type_str = "ARRAY"
+        if datatype.structured:
+            type_str = convert_sp_to_sf_type(datatype)
+        return f"PARSE_JSON({str_to_sql(json.dumps(value, cls=PythonObjJSONEncoder))}) :: {type_str}"
 
     if isinstance(value, dict) and isinstance(datatype, MapType):
-        return f"PARSE_JSON({str_to_sql(json.dumps(value, cls=PythonObjJSONEncoder))}) :: OBJECT"
+        type_str = "OBJECT"
+        if datatype.structured:
+            type_str = convert_sp_to_sf_type(datatype)
+        return f"PARSE_JSON({str_to_sql(json.dumps(value, cls=PythonObjJSONEncoder))}) :: {type_str}"
 
     if isinstance(datatype, VariantType):
         # PARSE_JSON returns VARIANT, so no need to append :: VARIANT here explicitly.
@@ -214,11 +220,14 @@ def schema_expression(data_type: DataType, is_nullable: bool) -> str:
             return "to_timestamp('2020-09-16 06:30:00')"
     if isinstance(data_type, ArrayType):
         if data_type.structured:
+            assert isinstance(data_type.element_type, DataType)
             element = schema_expression(data_type.element_type, is_nullable)
             return f"to_array({element}) :: {convert_sp_to_sf_type(data_type)}"
         return "to_array(0)"
     if isinstance(data_type, MapType):
         if data_type.structured:
+            assert isinstance(data_type.key_type, DataType)
+            assert isinstance(data_type.value_type, DataType)
             key = schema_expression(data_type.key_type, is_nullable)
             value = schema_expression(data_type.value_type, is_nullable)
             return f"object_construct_keep_null({key}, {value}) :: {convert_sp_to_sf_type(data_type)}"

src/snowflake/snowpark/_internal/type_utils.py

@@ -30,6 +30,7 @@
     get_origin,
 )
 
+import snowflake.snowpark.context as context
 import snowflake.snowpark.types  # type: ignore
 from snowflake.connector.constants import FIELD_ID_TO_NAME
 from snowflake.connector.cursor import ResultMetadata
@@ -183,12 +184,15 @@ def convert_sf_to_sp_type(
     max_string_size: int,
 ) -> DataType:
     """Convert the Snowflake logical type to the Snowpark type."""
+    semi_structured_fill = (
+        None if context._should_use_structured_type_semantics else StringType()
+    )
     if column_type_name == "ARRAY":
-        return ArrayType(StringType())
+        return ArrayType(semi_structured_fill)
     if column_type_name == "VARIANT":
         return VariantType()
     if column_type_name in {"OBJECT", "MAP"}:
-        return MapType(StringType(), StringType())
+        return MapType(semi_structured_fill, semi_structured_fill)
     if column_type_name == "GEOGRAPHY":
         return GeographyType()
     if column_type_name == "GEOMETRY":
@@ -530,7 +534,10 @@ def merge_type(a: DataType, b: DataType, name: Optional[str] = None) -> DataType
         return a
 
 
-def python_value_str_to_object(value, tp: DataType) -> Any:
+def python_value_str_to_object(value, tp: Optional[DataType]) -> Any:
+    if tp is None:
+        return None
+
     if isinstance(tp, StringType):
         return value
 
@@ -639,7 +646,7 @@ def python_type_to_snow_type(
         element_type = (
             python_type_to_snow_type(tp_args[0], is_return_type_of_sproc)[0]
             if tp_args
-            else StringType()
+            else None
         )
         return ArrayType(element_type), False
 
@@ -649,12 +656,12 @@ def python_type_to_snow_type(
         key_type = (
             python_type_to_snow_type(tp_args[0], is_return_type_of_sproc)[0]
             if tp_args
-            else StringType()
+            else None
         )
         value_type = (
             python_type_to_snow_type(tp_args[1], is_return_type_of_sproc)[0]
             if tp_args
-            else StringType()
+            else None
         )
         return MapType(key_type, value_type), False
 

src/snowflake/snowpark/context.py

@@ -21,6 +21,10 @@
 _should_continue_registration: Optional[Callable[..., bool]] = None
 
 
+# Global flag that determines if structured type semantics should be used
+_should_use_structured_type_semantics = False
+
+
 def get_active_session() -> "snowflake.snowpark.Session":
     """Returns the current active Snowpark session.
 

src/snowflake/snowpark/types.py

@@ -11,6 +11,7 @@
 from enum import Enum
 from typing import Generic, List, Optional, Type, TypeVar, Union, Dict, Any
 
+import snowflake.snowpark.context as context
 import snowflake.snowpark._internal.analyzer.expression as expression
 import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto
 
@@ -333,10 +334,16 @@ class ArrayType(DataType):
     def __init__(
         self,
         element_type: Optional[DataType] = None,
-        structured: bool = False,
+        structured: Optional[bool] = None,
     ) -> None:
-        self.structured = structured
-        self.element_type = element_type if element_type else StringType()
+        if context._should_use_structured_type_semantics:
+            self.structured = (
+                structured if structured is not None else element_type is not None
+            )
+            self.element_type = element_type
+        else:
+            self.structured = structured or False
+            self.element_type = element_type if element_type else StringType()
 
     def __repr__(self) -> str:
         return f"ArrayType({repr(self.element_type) if self.element_type else ''})"
@@ -379,14 +386,30 @@ def __init__(
         self,
         key_type: Optional[DataType] = None,
         value_type: Optional[DataType] = None,
-        structured: bool = False,
+        structured: Optional[bool] = None,
     ) -> None:
-        self.structured = structured
-        self.key_type = key_type if key_type else StringType()
-        self.value_type = value_type if value_type else StringType()
+        if context._should_use_structured_type_semantics:
+            if (key_type is None and value_type is not None) or (
+                key_type is not None and value_type is None
+            ):
+                raise ValueError(
+                    "Must either set both key_type and value_type or leave both unset."
+                )
+            self.structured = (
+                structured if structured is not None else key_type is not None
+            )
+            self.key_type = key_type
+            self.value_type = value_type
+        else:
+            self.structured = structured or False
+            self.key_type = key_type if key_type else StringType()
+            self.value_type = value_type if value_type else StringType()
 
     def __repr__(self) -> str:
-        return f"MapType({repr(self.key_type) if self.key_type else ''}, {repr(self.value_type) if self.value_type else ''})"
+        type_str = ""
+        if self.key_type and self.value_type:
+            type_str = f"{repr(self.key_type)}, {repr(self.value_type)}"
+        return f"MapType({type_str})"
 
     def is_primitive(self):
         return False
@@ -617,12 +640,20 @@ class StructType(DataType):
     """Represents a table schema or structured column. Contains :class:`StructField` for each field."""
 
     def __init__(
-        self, fields: Optional[List["StructField"]] = None, structured=False
+        self,
+        fields: Optional[List["StructField"]] = None,
+        structured: Optional[bool] = False,
     ) -> None:
-        self.structured = structured
-        if fields is None:
-            fields = []
-        self.fields = fields
+        if context._should_use_structured_type_semantics:
+            self.structured = (
+                structured if structured is not None else fields is not None
+            )
+            self.fields = fields or []
+        else:
+            self.structured = structured or False
+            if fields is None:
+                fields = []
+            self.fields = fields
 
     def add(
         self,

src/snowflake/snowpark/udaf.py

@@ -37,9 +37,10 @@
     TempObjectType,
     parse_positional_args_to_list,
     publicapi,
+    warning,
 )
 from snowflake.snowpark.column import Column
-from snowflake.snowpark.types import DataType
+from snowflake.snowpark.types import DataType, MapType
 
 # Python 3.8 needs to use typing.Iterable because collections.abc.Iterable is not subscriptable
 # Python 3.9 can use both
@@ -710,6 +711,14 @@ def _do_register_udaf(
             name,
         )
 
+        if isinstance(return_type, MapType):
+            if return_type.structured:
+                warning(
+                    "_do_register_udaf",
+                    "Snowflake does not support structured maps as return type for UDAFs. Downcasting to semi-structured object.",
+                )
+                return_type = MapType()
+
         # Capture original parameters.
         if _emit_ast:
             stmt = self._session._ast_batch.assign()

src/snowflake/snowpark/udtf.py

@@ -969,6 +969,10 @@ def _do_register_udtf(
             output_schema=output_schema,
         )
 
+        # Structured Struct is interpreted as Object by function registration
+        # Force unstructured to ensure Table return type.
+        output_schema.structured = False
+
         # Capture original parameters.
         if _emit_ast:
             stmt = self._session._ast_batch.assign()

tests/integ/test_column_names.py

@@ -361,7 +361,7 @@ def test_literal(session, local_testing_mode):
         BooleanType(),
         # snowflake doesn't enforce the inner type of ArrayType, so it is expected that
         # it returns StringType() as inner type.
-        ArrayType(LongType()) if local_testing_mode else ArrayType(StringType()),
+        ArrayType(LongType()) if local_testing_mode else ArrayType(),
     ]
     verify_column_result(
         session,

tests/integ/test_dataframe.py

@@ -2797,7 +2797,7 @@ def test_save_as_table_with_table_sproc_output(session, save_mode, table_type):
             lambda session_: session_.sql("SELECT 1 as A"),
             packages=["snowflake-snowpark-python"],
             name=temp_sp_name,
-            return_type=StructType([StructField("A", IntegerType())]),
+            return_type=StructType([StructField("A", IntegerType())], structured=False),
             input_types=[],
             replace=True,
         )

tests/integ/test_df_to_pandas.py

@@ -359,7 +359,7 @@ def test_df_to_pandas_df(session):
                 StructField("n", TimestampType()),
                 StructField("o", TimeType()),
                 StructField("p", VariantType()),
-                StructField("q", MapType(StringType(), StringType())),
+                StructField("q", MapType()),
             ]
         ),
     )

tests/integ/test_function.py

@@ -1291,9 +1291,9 @@ def test_to_date_to_array_to_variant_to_object(session, local_testing_mode):
     Utils.assert_rows(res1, expected)
 
     assert df1.schema.fields[0].datatype == DateType()
-    assert df1.schema.fields[1].datatype == ArrayType(StringType())
+    assert df1.schema.fields[1].datatype == ArrayType()
     assert df1.schema.fields[2].datatype == VariantType()
-    assert df1.schema.fields[3].datatype == MapType(StringType(), StringType())
+    assert df1.schema.fields[3].datatype == MapType()
 
 
 def test_to_binary(session):