SNOW-3237416: Support structured type schema string parsing (#4155)

Author: sfc-gh-wshangguan

src/snowflake/snowpark/_internal/type_utils.py

@@ -1363,6 +1363,193 @@ def is_likely_struct(s: str) -> bool:
     return top_level_space_found
 
 
+# ---------------------------------------------------------------------------
+# Structured-type INFER_SCHEMA parser
+#
+# These helpers parse Snowflake type strings as returned by the server's
+# INFER_SCHEMA function, including the structured forms
+# ``OBJECT(field type, ...)``, ``MAP(key, value)``, ``ARRAY(element)`` with
+# arbitrary nesting and ``NOT NULL`` annotations.
+#
+# Used by ``DataFrameReader._infer_schema_for_file_format`` when the session
+# flag ``Session._use_structured_type_infer_schema`` is enabled.
+# ---------------------------------------------------------------------------
+
+_STRUCTURED_TYPE_KEYWORDS = frozenset({"OBJECT", "MAP", "ARRAY"})
+
+_SF_EXTRA_TYPE_MAPPINGS = {
+    "text": StringType,
+    "real": DoubleType,
+    "fixed": LongType,
+}
+
+
+def _lookup_simple_type(name: str, original: str) -> DataType:
+    """Look up a simple Snowflake type name in the standard and extra mappings.
+
+    ``name`` is the candidate keyword to look up (lower-cased and
+    whitespace-stripped before lookup). ``original`` is the caller's full
+    type string, used verbatim in the error message so that parenthesized
+    inputs like ``"FOO(1)"`` show up correctly when ``"FOO"`` isn't found.
+
+    Raises ``ValueError`` if the name isn't recognized.
+    """
+    normalized = name.replace(" ", "").lower()
+    if normalized in DATA_TYPE_STRING_OBJECT_MAPPINGS:
+        return DATA_TYPE_STRING_OBJECT_MAPPINGS[normalized]()
+    if normalized in _SF_EXTRA_TYPE_MAPPINGS:
+        return _SF_EXTRA_TYPE_MAPPINGS[normalized]()
+    raise ValueError(f"'{original}' is not a supported type")
+
+
+def _extract_paren_content(type_str: str) -> Optional[Tuple[str, str]]:
+    """Extract the base keyword and content inside matching parentheses.
+
+    Returns (base, inner_content) if matching parens are found.
+    Returns None if ``type_str`` contains no ``(`` at all (normal for simple
+    types like ``VARCHAR`` or ``BOOLEAN``).
+
+    Raises ``ValueError`` if ``type_str`` contains a ``(`` that is never
+    closed. Reaching this branch implies a malformed type string from the
+    backend (``INFER_SCHEMA``), so we fail loudly rather than silently
+    degrade to ``VariantType``.
+
+    E.g. "OBJECT(city VARCHAR, zip NUMBER(38,0))" -> ("OBJECT", "city VARCHAR, zip NUMBER(38,0)")
+    """
+    paren_idx = type_str.find("(")
+    if paren_idx == -1:
+        return None
+    base = type_str[:paren_idx].strip()
+    depth = 0
+    for i in range(paren_idx, len(type_str)):
+        if type_str[i] == "(":
+            depth += 1
+        elif type_str[i] == ")":
+            depth -= 1
+            if depth == 0:
+                return base, type_str[paren_idx + 1 : i]
+    raise ValueError(f"Unbalanced parentheses in type string: '{type_str}'")
+
+
+def _sf_type_to_type_object(type_str: str) -> DataType:
+    """Parse a Snowflake SQL type string directly into a Snowpark DataType.
+
+    Handles both simple types and structured types returned by INFER_SCHEMA:
+    - Simple: VARCHAR, NUMBER(38,0), BOOLEAN, TIMESTAMP_NTZ, etc.
+    - ARRAY(element_type [NOT NULL])
+    - MAP(key_type, value_type [NOT NULL])
+    - OBJECT(field1 type1, field2 type2 NOT NULL, ...)
+    - Nested combinations of the above
+
+    NOT NULL annotations are respected:
+    - On ARRAY elements: sets ArrayType.contains_null = False
+    - On MAP values: sets MapType.value_contains_null = False
+    - On OBJECT fields: sets StructField.nullable = False
+    """
+    type_str = type_str.strip()
+    if not type_str:
+        raise ValueError("Empty type string")
+
+    # Strip a trailing top-level NOT NULL if present and discard the bool:
+    # top-level column nullability is carried by INFER_SCHEMA row metadata
+    # (handled in _infer_schema_for_file_format), not by the type string.
+    # Nested NOT NULL (inside ARRAY/MAP/OBJECT) is already consumed by those
+    # branches below before they recurse into this function, so the bool is
+    # redundant here.
+    type_str, _ = extract_nullable_keyword(type_str)
+
+    result = _extract_paren_content(type_str)
+    if result is None:
+        return _lookup_simple_type(type_str, type_str)
+
+    base, inner = result
+    base_upper = base.upper()
+
+    if base_upper == "ARRAY":
+        element_str, element_nullable = extract_nullable_keyword(inner)
+        element_type = _sf_type_to_type_object(element_str)
+        return ArrayType(element_type, structured=True, contains_null=element_nullable)
+
+    if base_upper == "MAP":
+        parts = split_top_level_comma_fields(inner)
+        if len(parts) != 2:
+            raise ValueError(f"Invalid MAP type definition: '{type_str}'")
+        key_type = _sf_type_to_type_object(parts[0])
+        value_str, value_nullable = extract_nullable_keyword(parts[1])
+        value_type = _sf_type_to_type_object(value_str)
+        return MapType(
+            key_type, value_type, structured=True, value_contains_null=value_nullable
+        )
+
+    if base_upper == "OBJECT":
+        # OBJECT() with no inner content is valid per the Snowflake grammar:
+        # "a structured OBJECT that contains no keys."
+        if not inner.strip():
+            return StructType([], structured=True)
+        fields = split_top_level_comma_fields(inner)
+        struct_fields = []
+        for field_def in fields:
+            field_def = field_def.strip()
+            if not field_def:
+                # A trailing comma or empty fragment is not valid Snowflake
+                # SQL grammar for OBJECT types; raise so backend bugs or
+                # malformed input surface loudly.
+                raise ValueError(f"Empty field in OBJECT type: '{type_str}'")
+            parts = field_def.split(None, 1)
+            if len(parts) != 2:
+                raise ValueError(f"Cannot parse OBJECT field definition: '{field_def}'")
+            field_name = parts[0]
+            type_part, nullable = extract_nullable_keyword(parts[1])
+            field_type = _sf_type_to_type_object(type_part)
+            struct_fields.append(StructField(field_name, field_type, nullable=nullable))
+        return StructType(struct_fields, structured=True)
+
+    precision_scale = get_number_precision_scale(type_str)
+    if precision_scale:
+        return DecimalType(*precision_scale)
+    length = get_string_length(type_str)
+    if length:
+        return StringType(length)
+
+    return _lookup_simple_type(base_upper, type_str)
+
+
+def _parse_structured_type_str(type_str: str, max_string_size: int) -> DataType:
+    """Parse a Snowflake type string from INFER_SCHEMA into a Snowpark DataType.
+
+    For structured types (OBJECT, MAP, ARRAY), uses the recursive parser.
+    For simple types, delegates to convert_sf_to_sp_type for precision/scale.
+    """
+    type_str = type_str.strip()
+    if not type_str:
+        return VariantType()
+
+    result = _extract_paren_content(type_str)
+    base_upper = result[0].upper() if result else type_str.upper()
+
+    if base_upper in _STRUCTURED_TYPE_KEYWORDS:
+        if result is not None:
+            return _sf_type_to_type_object(type_str)
+        # Bare structured keyword (e.g. "OBJECT", "MAP", "ARRAY") without
+        # inner type details — older backends may return these.  Return
+        # VariantType so column names are preserved and callers (e.g. SAS)
+        # can apply their own structured-type discovery.
+        return VariantType()
+
+    if result is None:
+        return convert_sf_to_sp_type(base_upper, 0, 0, 0, max_string_size)
+
+    inner = result[1]
+    parts = inner.split(",")
+    try:
+        precision = int(parts[0].strip())
+        scale = int(parts[1].strip()) if len(parts) > 1 else 0
+    except (ValueError, IndexError):
+        precision = 0
+        scale = 0
+    return convert_sf_to_sp_type(base_upper, precision, scale, 0, max_string_size)
+
+
 def type_string_to_type_object(type_str: str) -> DataType:
     type_str = type_str.strip()
     if not type_str:

src/snowflake/snowpark/dataframe_reader.py

@@ -54,6 +54,7 @@
     convert_sp_to_sf_type,
     most_permissive_type,
     type_string_to_type_object,
+    _parse_structured_type_str,
 )
 from snowflake.snowpark._internal.udf_utils import get_types_from_type_hints
 from snowflake.snowpark._internal.xml_reader import DEFAULT_CHUNK_SIZE
@@ -112,6 +113,9 @@
 
 logger = getLogger(__name__)
 
+_NOT_NULL_RE = re.compile(r"\s+NOT\s+NULL", re.IGNORECASE)
+
+
 LOCAL_TESTING_SUPPORTED_FILE_FORMAT = ("JSON",)
 
 _TIME_TRAVEL_OPTIONS_PARAMS_MAP = {
@@ -1341,25 +1345,39 @@ def _infer_schema_for_file_format(
             schema_to_cast = []
             transformations: List["snowflake.snowpark.column.Column"] = []
             read_file_transformations = None
+            use_structured_type_infer_schema = (
+                self._session._use_structured_type_infer_schema
+            )
             for r in results:
                 # Columns for r [column_name, type, nullable, expression, filenames, order_id]
                 column_name, type, nullable, expression = r[0], r[1], r[2], r[3]
                 name = quote_name_without_upper_casing(column_name)
-                # Parse the type returned by infer_schema command to
-                # pass to determine datatype for schema
-                data_type_parts = type.split("(")
-                parts_length = len(data_type_parts)
-                if parts_length == 1:
-                    data_type = type
-                    precision = 0
-                    scale = 0
+                # Parse the type returned by infer_schema command.
+                if use_structured_type_infer_schema:
+                    # handles both simple types and structured
+                    # types (OBJECT, MAP, ARRAY with inner details).
+                    datatype = _parse_structured_type_str(
+                        type, self._session._conn.max_string_size
+                    )
                 else:
-                    data_type = data_type_parts[0]
-                    precision = int(data_type_parts[1].split(",")[0])
-                    scale = int(data_type_parts[1].split(",")[1][:-1])
-                datatype = convert_sf_to_sp_type(
-                    data_type, precision, scale, 0, self._session._conn.max_string_size
-                )
+                    data_type_parts = type.split("(")
+                    parts_length = len(data_type_parts)
+                    if parts_length == 1:
+                        data_type = type
+                        precision = 0
+                        scale = 0
+                    else:
+                        data_type = data_type_parts[0]
+                        precision = int(data_type_parts[1].split(",")[0])
+                        scale = int(data_type_parts[1].split(",")[1][:-1])
+                    datatype = convert_sf_to_sp_type(
+                        data_type,
+                        precision,
+                        scale,
+                        0,
+                        self._session._conn.max_string_size,
+                    )
+
                 if use_relaxed_types:
                     datatype = most_permissive_type(datatype)
                 new_schema.append(
@@ -1379,6 +1397,26 @@ def _infer_schema_for_file_format(
                     "GET_IGNORE_CASE"
                 ):
                     identifier = expression
+                elif (
+                    format.lower() in ("parquet", "json")
+                    and use_structured_type_infer_schema
+                ):
+                    if isinstance(datatype, VariantType):
+                        # Bare structured keyword was returned by the backend
+                        # (no inner details).  Skip the cast — $1:{name}
+                        # extracts as variant and lets callers handle
+                        # structured-type discovery.
+                        identifier = f"$1:{name}"
+                    elif use_relaxed_types:
+                        identifier = f"$1:{name}::{convert_sp_to_sf_type(datatype)}"
+                    else:
+                        # INFER_SCHEMA may return NOT NULL annotations in
+                        # structured type strings (e.g.,
+                        # "ARRAY(NUMBER(10,0) NOT NULL)"). Strip them before
+                        # embedding in SQL — nullable info is already captured
+                        # in the parsed DataType objects.
+                        cast_type = _NOT_NULL_RE.sub("", type)
+                        identifier = f"$1:{name}::{cast_type}"
                 else:
                     identifier = f"$1:{name}::{convert_sp_to_sf_type(datatype) if use_relaxed_types else type}"
 

src/snowflake/snowpark/session.py

@@ -710,6 +710,15 @@ def __init__(
             )
         )
 
+        # When True, INFER_SCHEMA results for parquet/json files with
+        # structured types (OBJECT/MAP/ARRAY with inner type details) are
+        # parsed into full Snowpark DataType objects by DataFrameReader and
+        # TRY_CAST is used in the SELECT expression. Defaults to False;
+        # opt in by direct assignment, e.g.
+        #     session._use_structured_type_infer_schema = True
+        # (SAS / snowpark-connect enables this during session configuration.)
+        self._use_structured_type_infer_schema: bool = False
+
         self._large_query_breakdown_enabled: bool = self.is_feature_enabled_for_version(
             _PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION_VERSION
         )