fix: preserve geospatial result metadata and shapes

Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com>

Author: cathleeny

src/databricks/sql/backend/thrift_backend.py

@@ -808,23 +808,31 @@ def _col_to_description(col, field=None, host_url=None):
 
     @staticmethod
     def _hive_schema_to_description(t_table_schema, schema_bytes=None, host_url=None):
+        arrow_fields = []
         field_dict = {}
         if pyarrow and schema_bytes:
             try:
                 arrow_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes))
-                # Build a dictionary mapping column names to fields
-                for field in arrow_schema:
+                arrow_fields = list(arrow_schema)
+                # Retain a name fallback for malformed/partial schemas, but use
+                # result-column position whenever it is available. SQL permits
+                # duplicate aliases, for which a name lookup is ambiguous.
+                for field in arrow_fields:
                     field_dict[field.name] = field
             except Exception as e:
                 logger.debug(f"Could not parse arrow schema: {e}")
 
         return [
             ThriftDatabricksClient._col_to_description(
                 col,
-                field_dict.get(col.columnName) if field_dict else None,
+                (
+                    arrow_fields[index]
+                    if index < len(arrow_fields)
+                    else field_dict.get(col.columnName)
+                ),
                 host_url,
             )
-            for col in t_table_schema.columns
+            for index, col in enumerate(t_table_schema.columns)
         ]
 
     def _results_message_to_execute_response(

tests/unit/test_kernel_result_set.py

@@ -95,7 +95,8 @@ def test_fetchall_arrow_drains_all_batches(int_schema):
     assert rs.has_more_rows is False
 
 
-def test_geospatial_string_and_binary_values_keep_logical_type():
+@pytest.mark.parametrize("disable_pandas", [False, True])
+def test_geospatial_string_and_binary_values_keep_logical_type(disable_pandas):
     wkb = bytes.fromhex("0101000000000000000000F03F0000000000000040")
     geo_metadata = {
         b"databricks.type_name": b"GEOMETRY",
@@ -108,7 +109,8 @@ def test_geospatial_string_and_binary_values_keep_logical_type():
         schema=string_schema,
     )
     string_rows = _make_rs(
-        _FakeKernelHandle(string_schema, [string_batch]), disable_pandas=True
+        _FakeKernelHandle(string_schema, [string_batch]),
+        disable_pandas=disable_pandas,
     ).fetchall()
     assert [row[0] for row in string_rows] == ["SRID=4326;POINT(1 2)", None]
 
@@ -124,10 +126,14 @@ def test_geospatial_string_and_binary_values_keep_logical_type():
         schema=binary_schema,
     )
     binary_rs = _make_rs(
-        _FakeKernelHandle(binary_schema, [binary_batch]), disable_pandas=True
+        _FakeKernelHandle(binary_schema, [binary_batch]),
+        disable_pandas=disable_pandas,
     )
     assert binary_rs.description[0][1] == "geometry"
     binary_rows = binary_rs.fetchall()
+    assert isinstance(binary_rows[0][0], dict)
+    assert isinstance(binary_rows[0][0]["srid"], int)
+    assert isinstance(binary_rows[0][0]["wkb"], bytes)
     assert binary_rows[0][0] == {"srid": 4326, "wkb": wkb}
     assert binary_rows[1][0] is None
 

tests/unit/test_thrift_backend.py

@@ -2627,6 +2627,17 @@ def test_hive_schema_to_description(self):
                 None,  # No arrow schema
                 [("regular_col", "string")],
             ),
+            (
+                [
+                    ("g", ttypes.TTypeId.STRING_TYPE),
+                    ("g", ttypes.TTypeId.STRING_TYPE),
+                ],
+                [
+                    ("g", {}),
+                    ("g", {b"Spark:DataType:SqlName": b"GEOMETRY(0)"}),
+                ],
+                [("g", "string"), ("g", "geometry")],
+            ),
         ]
 
         for columns, arrow_fields, expected_types in test_cases: