SNOW-1000284: Add support for structured types to ResultMetadataV2 (#1890)

Author: sfc-gh-jrose

DESCRIPTION.md

@@ -20,6 +20,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
     - Add `SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT` flag (usage: `SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT=true`) to make a non-blocking socket.recv call and retry on Error
       - Consider using this if running in a containerized environment and externalbrowser auth frequently hangs while waiting for callback
       - NOTE: this has not been tested extensively, but has been shown to improve the experience when using WSL
+  - Added support for parsing structured type information in schema queries.
   - Bumped platformdirs from >=2.6.0,<4.0.0 to >=2.6.0,<5.0.0
   - Updated diagnostics to use system$allowlist instead of system$whitelist.
   - Update `write_pandas` to skip TABLE IF NOT EXISTS in truncate mode.

src/snowflake/connector/constants.py

@@ -44,38 +44,75 @@ class FieldType(NamedTuple):
 
 
 def vector_pa_type(metadata: ResultMetadataV2) -> DataType:
-    """Generate the Arrow type represented by the given vector column metadata.
-
+    """
+    Generate the Arrow type represented by the given vector column metadata.
     Vectors are represented as Arrow fixed-size lists.
     """
+    assert (
+        metadata.fields is not None and len(metadata.fields) == 1
+    ), "Invalid result metadata for vector type: expected a single field to be defined"
+    assert (
+        metadata.vector_dimension or 0
+    ) > 0, "Invalid result metadata for vector type: expected a positive dimension"
+
+    field_type = FIELD_TYPES[metadata.fields[0].type_code]
+    return pa.list_(field_type.pa_type(metadata.fields[0]), metadata.vector_dimension)
+
+
+def array_pa_type(metadata: ResultMetadataV2) -> DataType:
+    """
+    Generate the Arrow type represented by the given array column metadata.
+    """
+    # If fields is missing then structured types are not enabled.
+    # Fallback to json encoded string
+    if metadata.fields is None:
+        return pa.string()
+
+    assert (
+        len(metadata.fields) == 1
+    ), "Invalid result metadata for array type: expected a single field to be defined"
+
+    field_type = FIELD_TYPES[metadata.fields[0].type_code]
+    return pa.list_(field_type.pa_type(metadata.fields[0]))
+
+
+def map_pa_type(metadata: ResultMetadataV2) -> DataType:
+    """
+    Generate the Arrow type represented by the given map column metadata.
+    """
+    # If fields is missing then structured types are not enabled.
+    # Fallback to json encoded string
+    if metadata.fields is None:
+        return pa.string()
 
+    assert (
+        len(metadata.fields or []) == 2
+    ), "Invalid result metadata for map type: expected a field for key and a field for value"
+    key_type = FIELD_TYPES[metadata.fields[0].type_code]
+    value_type = FIELD_TYPES[metadata.fields[1].type_code]
+    return pa.map_(
+        key_type.pa_type(metadata.fields[0]), value_type.pa_type(metadata.fields[1])
+    )
+
+
+def struct_pa_type(metadata: ResultMetadataV2) -> DataType:
+    """
+    Generate the Arrow type represented by the given struct column metadata.
+    """
+    # If fields is missing then structured types are not enabled.
+    # Fallback to json encoded string
     if metadata.fields is None:
-        raise ValueError(
-            "Invalid result metadata for vector type: expected sub-field metadata"
-        )
-    if len(metadata.fields) != 1:
-        raise ValueError(
-            "Invalid result metadata for vector type: expected a single sub-field metadata"
-        )
-    field_type = FIELD_ID_TO_NAME[metadata.fields[0].type_code]
-
-    if metadata.vector_dimension is None:
-        raise ValueError(
-            "Invalid result metadata for vector type: expected a dimension"
-        )
-    elif metadata.vector_dimension <= 0:
-        raise ValueError(
-            "Invalid result metadata for vector type: expected a positive dimension"
-        )
-
-    if field_type == "FIXED":
-        return pa.list_(pa.int32(), metadata.vector_dimension)
-    elif field_type == "REAL":
-        return pa.list_(pa.float32(), metadata.vector_dimension)
-    else:
-        raise ValueError(
-            f"Invalid result metadata for vector type: invalid element type: {field_type}"
-        )
+        return pa.string()
+
+    assert all(
+        field.name is not None for field in metadata.fields
+    ), "All fields of a stuct type must have a name."
+    return pa.struct(
+        {
+            field.name: FIELD_TYPES[field.type_code].pa_type(field)
+            for field in metadata.fields
+        }
+    )
 
 
 # This type mapping holds column type definitions.
@@ -121,12 +158,8 @@ def vector_pa_type(metadata: ResultMetadataV2) -> DataType:
         dbapi_type=[DBAPI_TYPE_TIMESTAMP],
         pa_type=lambda _: pa.timestamp("ns"),
     ),
-    FieldType(
-        name="OBJECT", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=lambda _: pa.string()
-    ),
-    FieldType(
-        name="ARRAY", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=lambda _: pa.string()
-    ),
+    FieldType(name="OBJECT", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=struct_pa_type),
+    FieldType(name="ARRAY", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=array_pa_type),
     FieldType(
         name="BINARY", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=lambda _: pa.binary()
     ),
@@ -143,6 +176,7 @@ def vector_pa_type(metadata: ResultMetadataV2) -> DataType:
         name="GEOMETRY", dbapi_type=[DBAPI_TYPE_STRING], pa_type=lambda _: pa.string()
     ),
     FieldType(name="VECTOR", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=vector_pa_type),
+    FieldType(name="MAP", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=map_pa_type),
 )
 
 FIELD_NAME_TO_ID: DefaultDict[Any, int] = defaultdict(int)

src/snowflake/connector/cursor.py

@@ -201,28 +201,35 @@ def from_column(cls, col: dict[str, Any]) -> ResultMetadataV2:
         This differs from ResultMetadata in that it has newly-added fields which cannot be added to
         ResultMetadata since it is a named tuple.
         """
-        type_code = FIELD_NAME_TO_ID[
+        col_type = (
             col["extTypeName"].upper()
             if col.get("extTypeName")
             else col["type"].upper()
-        ]
+        )
 
-        fields = None
-        if type_code == FIELD_NAME_TO_ID["VECTOR"] and col.get("fields") is not None:
-            fields = [
-                ResultMetadataV2.from_column({"name": None, **f}) for f in col["fields"]
-            ]
+        fields = col.get("fields")
+        processed_fields: Optional[List[ResultMetadataV2]] = None
+        if fields is not None:
+            if col_type in {"VECTOR", "ARRAY", "OBJECT", "MAP"}:
+                processed_fields = [
+                    ResultMetadataV2.from_column({"name": None, **f})
+                    for f in col["fields"]
+                ]
+            else:
+                raise ValueError(
+                    f"Field parsing is not supported for columns of type {col_type}."
+                )
 
         return cls(
             col["name"],
-            type_code,
+            FIELD_NAME_TO_ID[col_type],
             col["nullable"],
             None,
             col["length"],
             col["precision"],
             col["scale"],
             col.get("vectorDimension"),
-            fields,
+            processed_fields,
         )
 
     def _to_result_metadata_v1(self):

test/integ/pandas/test_arrow_pandas.py

@@ -1202,6 +1202,93 @@ def test_batch_to_pandas_arrow(conn_cnx, result_format):
                 assert arrow_table.to_pydict()["FOO"] == list(range(rowcount))
 
 
+@pytest.mark.internal
+@pytest.mark.parametrize("enable_structured_types", [True, False])
+def test_to_arrow_datatypes(enable_structured_types, conn_cnx):
+    with conn_cnx() as cnx:
+        with cnx.cursor() as cur:
+            cur.execute(SQL_ENABLE_ARROW)
+
+            expected_types = (
+                pyarrow.int64(),
+                pyarrow.float64(),
+                pyarrow.string(),
+                pyarrow.date64(),
+                pyarrow.timestamp("ns"),
+                pyarrow.string(),
+                pyarrow.timestamp("ns"),
+                pyarrow.timestamp("ns"),
+                pyarrow.timestamp("ns"),
+                pyarrow.binary(),
+                pyarrow.time64("ns"),
+                pyarrow.bool_(),
+                pyarrow.string(),
+                pyarrow.string(),
+                pyarrow.list_(pyarrow.float64(), 5),
+            )
+
+            if enable_structured_types:
+                structured_params = {
+                    "ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE",
+                    "IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE",
+                    "FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT",
+                }
+                try:
+                    for param in structured_params:
+                        cur.execute(f"alter session set {param}=true")
+                    expected_types += (
+                        pyarrow.map_(pyarrow.string(), pyarrow.int64()),
+                        pyarrow.struct(
+                            {"city": pyarrow.string(), "population": pyarrow.float64()}
+                        ),
+                        pyarrow.list_(pyarrow.float64()),
+                    )
+                finally:
+                    cur.execute(f"alter session unset {param}")
+            else:
+                expected_types += (
+                    pyarrow.string(),
+                    pyarrow.string(),
+                    pyarrow.string(),
+                )
+
+            # Ensure an empty batch to use default typing
+            # Otherwise arrow will resize types to save space
+            cur.execute(
+                """
+            select
+            1 :: INTEGER as FIXED_type,
+            2.0 :: FLOAT as REAL_type,
+            'test' :: TEXT as TEXT_type,
+            '2024-02-28' :: DATE as DATE_type,
+            '2020-03-12 01:02:03.123456789' :: TIMESTAMP as TIMESTAMP_type,
+            '{"foo": "bar"}' :: VARIANT as VARIANT_type,
+            '2020-03-12 01:02:03.123456789' :: TIMESTAMP_LTZ as TIMESTAMP_LTZ_type,
+            '2020-03-12 01:02:03.123456789' :: TIMESTAMP_TZ as TIMESTAMP_TZ_type,
+            '2020-03-12 01:02:03.123456789' :: TIMESTAMP_NTZ as TIMESTAMP_NTZ_type,
+            '0xAAAA' :: BINARY as BINARY_type,
+            '01:02:03.123456789' :: TIME as TIME_type,
+            true :: BOOLEAN as BOOLEAN_type,
+            TO_GEOGRAPHY('LINESTRING(13.4814 52.5015, -121.8212 36.8252)') as GEOGRAPHY_type,
+            TO_GEOMETRY('LINESTRING(13.4814 52.5015, -121.8212 36.8252)') as GEOMETRY_type,
+            [1,2,3,4,5] :: vector(float, 5) as VECTOR_type,
+            object_construct('k1', 1, 'k2', 2, 'k3', 3, 'k4', 4, 'k5', 5) :: map(varchar, int) as MAP_type,
+            object_construct('city', 'san jose', 'population', 0.05) :: object(city varchar, population float) as OBJECT_type,
+            [1.0, 3.1, 4.5] :: array(float) as ARRAY_type
+            WHERE 1=0
+            """
+            )
+
+            batches = cur.get_result_batches()
+            assert len(batches) == 1
+            batch = batches[0]
+            arrow_table = batch.to_arrow()
+            for actual, expected in zip(arrow_table.schema, expected_types):
+                assert (
+                    actual.type == expected
+                ), f"Expected {actual.name} :: {actual.type} column to be of type {expected}"
+
+
 def test_simple_arrow_fetch(conn_cnx):
     rowcount = 250_000
     with conn_cnx() as cnx: