feat: support geospatial results on the kernel path

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

Author: cathleeny

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Release History
 
+# Unreleased
+- Add the kernel-only `geospatial_as_string` connection option. GEOMETRY / GEOGRAPHY results are exposed as EWKT strings when true or `{"srid": int, "wkb": bytes}` values when false.
+
 # 4.6.0 (2026-09-24)
 - Upgrade Databricks SQL Kernel to 1.1.0; the kernel dependency is now stable and no longer experimental.
 - Transparently auto-recover Thrift connections to Reyden / Real-Time warehouses: when a warehouse rejects the default Thrift protocol (SQLSTATE `KP001`), the session is re-opened on the kernel backend and the warehouse is remembered so later connections skip Thrift. Applies only when no backend was chosen explicitly.

CONNECTION_PARAMETERS.md

@@ -153,6 +153,7 @@ to change without notice.
 | `max_download_threads`                | `int`  |   ✅   |   ❌   | `10`          | Worker threads for cloud-fetch downloads. Not forwarded to the kernel.                                        |
 | `enable_query_result_lz4_compression` | `bool` |   ✅   |   ❌   | `True`        | LZ4-compress result payloads. Not forwarded; the kernel handles compression internally.                       |
 | `_disable_pandas`                     | `bool` |   ✅   |   ✅   | `False`       | Skip the pandas-based Arrow→row deserialization and materialize rows directly with PyArrow. This is a **Python-side** result-conversion toggle, not a wire option: the kernel returns results as Arrow (`RecordBatch`es) and the connector runs the *same* `_convert_arrow_table` for both backends, so the flag is honored on the kernel path too. Affects only row fetches (`fetchone`/`fetchmany`/`fetchall`); the `fetch*_arrow` methods return the Arrow table unchanged regardless of this flag. |
+| `geospatial_as_string`                | `bool \| None` |   ❌   |   ✅   | `None` (kernel default: `True`) | Return GEOMETRY / GEOGRAPHY as EWKT strings when `True`, or as `{"srid": int, "wkb": bytes}` values when `False`. `None` leaves the kernel default in force. This is a local result conversion and is never forwarded to SEA. |
 | `_use_arrow_native_complex_types`     | `bool` |   ✅   |   ✅   | `True`        | Return `ARRAY`/`MAP`/`STRUCT` as native Arrow types instead of JSON strings. Forwarded to the kernel.         |
 | `_use_arrow_native_decimals`          | `bool` |   ✅   |   ❌   | `True`        | Thrift wire encoding for `DECIMAL`: `True` → native Arrow `decimal128`, `False` → Arrow string. **No value-level effect**, though: the connector unconditionally re-casts the column back to `decimal128` (`convert_decimals_in_arrow_table`, `thrift_backend.py`), so both `fetchall()` and `fetchall_arrow()` yield `Decimal` / `decimal128(p,s)` either way (verified live). Not forwarded to the kernel, which always returns native Arrow decimals. |
 | `_use_arrow_native_timestamps`        | `bool` |   ✅   |   ❌   | `True`        | Thrift wire encoding for `TIMESTAMP`: `True` → native Arrow timestamp (→ Python `datetime`), `False` → Arrow string (→ Python **`str`**). **Unlike decimals there is no re-cast**, so `False` genuinely surfaces strings — and `cursor.description` still reports the type code as `'timestamp'`, a mismatch to watch for (verified live). Note the connector always also sends the `spark.thriftserver.arrowBasedRowSet.timestampAsString=false` conf, but the `timestampAsArrow=False` flag wins. Not forwarded to the kernel, which always returns native Arrow timestamps. |

KERNEL_REV

@@ -1 +1 @@
-80f2aee7d884994d7b0af9a9ea6078872859a9cd
+ad3bc6993bca95b810839feade77ccd0ab98ece5

src/databricks/sql/backend/kernel/client.py

@@ -169,6 +169,29 @@ def _kernel_session_accepts_kwarg(name: str) -> bool:
     return name in params
 
 
+def _kernel_geospatial_kwargs(value: Optional[bool]) -> Dict[str, bool]:
+    """Build the optional geospatial result-representation kwarg.
+
+    ``None`` deliberately omits the option so the installed kernel owns its
+    default. An explicit value must never be silently ignored: older kernel
+    wheels do not declare ``geospatial_as_string`` and would otherwise return a
+    different public value shape than the caller requested.
+    """
+    if value is None:
+        return {}
+    if not isinstance(value, bool):
+        raise ValueError(
+            "geospatial_as_string must be a bool or None; "
+            f"got {type(value).__name__}"
+        )
+    if not _kernel_session_accepts_kwarg("geospatial_as_string"):
+        raise NotSupportedError(
+            "geospatial_as_string requires a newer databricks-sql-kernel "
+            "wheel that exposes geospatial result representation support."
+        )
+    return {"geospatial_as_string": value}
+
+
 def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]:
     """Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``.
 
@@ -262,6 +285,12 @@ def __init__(
         # The kernel binding owns type and range validation.
         self._request_timeout_secs = kwargs.get("request_timeout_secs")
         self._max_connections = kwargs.get("max_connections")
+        # Client-side result representation for GEOMETRY / GEOGRAPHY. None
+        # leaves the kernel default in force (EWKT strings); False requests the
+        # canonical Arrow struct and surfaces as ``{"srid": int, "wkb":
+        # bytes}`` through pyarrow. This is intentionally separate from
+        # ``session_configuration``: it is never forwarded to SEA.
+        self._geospatial_as_string = kwargs.get("geospatial_as_string")
         # Kernel telemetry phase 7 adds binding/runtime identity and
         # telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
         self._telemetry_options = kwargs.get("telemetry_options") or {}
@@ -379,6 +408,7 @@ def open_session(
             # kernel's ``retry_*`` kwargs. Empty when at defaults.
             retry_kwargs = _kernel_retry_kwargs(self._retry_options)
             telemetry_kwargs = _kernel_telemetry_kwargs(self._telemetry_options)
+            geospatial_kwargs = _kernel_geospatial_kwargs(self._geospatial_as_string)
             max_connections_kwargs: Dict[str, Any] = {}
             if _kernel_session_accepts_kwarg("max_connections"):
                 max_connections_kwargs["max_connections"] = self._max_connections
@@ -426,6 +456,7 @@ def open_session(
                 **tls_kwargs,
                 **retry_kwargs,
                 **telemetry_kwargs,
+                **geospatial_kwargs,
                 **max_connections_kwargs,
                 **http_headers_kwargs,
             )

src/databricks/sql/client.py

@@ -194,6 +194,14 @@ def __init__(
                 decision. This is an intentional divergence from the
                 Thrift/SEA paths, where an explicit ``True`` can still
                 be suppressed by the feature flag.
+            :param geospatial_as_string: `bool | None`, optional (default is None)
+                Kernel backend only. Controls the public representation of
+                ``GEOMETRY`` and ``GEOGRAPHY`` result values. ``True`` returns
+                EWKT strings (for example ``"SRID=4326;POINT(1 2)"``);
+                ``False`` returns ``{"srid": int, "wkb": bytes}``; and
+                ``None`` uses the kernel default (currently EWKT strings).
+                The conversion is local to the kernel/driver and this option is
+                never sent to the SQL Execution API.
             :param use_hybrid_disposition: `bool`, optional (default is False)
                 Use the hybrid disposition instead of the inline disposition.
             :param server_hostname: Databricks instance host name.

src/databricks/sql/session.py

@@ -314,6 +314,7 @@ def _create_backend(
                 retry_options=kernel_retry_options,
                 request_timeout_secs=kwargs.get("_socket_timeout"),
                 max_connections=kwargs.get("_pool_maxsize") or None,
+                geospatial_as_string=kwargs.get("geospatial_as_string"),
                 telemetry_options=kernel_telemetry_options,
             )
 

tests/unit/test_kernel_client.py

@@ -392,6 +392,94 @@ def fake_session(**kw):
     assert captured["max_connections"] == max_connections
 
 
+@pytest.mark.parametrize("as_string", [True, False])
+def test_open_session_passes_geospatial_representation_to_kernel(
+    monkeypatch, as_string
+):
+    captured = {}
+
+    def fake_session(*, geospatial_as_string=None, **kw):
+        captured["geospatial_as_string"] = geospatial_as_string
+        sess = MagicMock()
+        sess.session_id = "sess-id"
+        return sess
+
+    monkeypatch.setattr(kernel_client._kernel, "Session", fake_session)
+    c = kernel_client.KernelDatabricksClient(
+        server_hostname="example.cloud.databricks.com",
+        http_path="/sql/1.0/warehouses/abc",
+        auth_provider=AccessTokenAuthProvider("dapi-test"),
+        ssl_options=None,
+        geospatial_as_string=as_string,
+    )
+
+    c.open_session(session_configuration=None, catalog=None, schema=None)
+
+    assert captured["geospatial_as_string"] is as_string
+
+
+def test_open_session_omits_unset_geospatial_representation(monkeypatch):
+    captured = {}
+
+    def fake_session(**kw):
+        captured.update(kw)
+        sess = MagicMock()
+        sess.session_id = "sess-id"
+        return sess
+
+    monkeypatch.setattr(kernel_client._kernel, "Session", fake_session)
+    c = kernel_client.KernelDatabricksClient(
+        server_hostname="example.cloud.databricks.com",
+        http_path="/sql/1.0/warehouses/abc",
+        auth_provider=AccessTokenAuthProvider("dapi-test"),
+        ssl_options=None,
+    )
+
+    c.open_session(session_configuration=None, catalog=None, schema=None)
+
+    assert "geospatial_as_string" not in captured
+
+
+def test_open_session_rejects_explicit_geospatial_option_with_old_kernel(
+    monkeypatch,
+):
+    def fake_session_without_geospatial(
+        host,
+        http_path,
+        *,
+        catalog=None,
+        schema=None,
+        session_conf=None,
+        complex_types_as_json=False,
+        intervals_as_string=False,
+        request_timeout_secs=None,
+        auth_type=None,
+        access_token=None,
+    ):
+        sess = MagicMock()
+        sess.session_id = "sess-id"
+        return sess
+
+    monkeypatch.setattr(
+        kernel_client._kernel, "Session", fake_session_without_geospatial
+    )
+    c = kernel_client.KernelDatabricksClient(
+        server_hostname="example.cloud.databricks.com",
+        http_path="/sql/1.0/warehouses/abc",
+        auth_provider=AccessTokenAuthProvider("dapi-test"),
+        ssl_options=None,
+        geospatial_as_string=False,
+    )
+
+    with pytest.raises(NotSupportedError, match="newer databricks-sql-kernel"):
+        c.open_session(session_configuration=None, catalog=None, schema=None)
+
+
+def test_geospatial_option_rejects_non_bool():
+    with pytest.raises(ValueError, match="must be a bool or None"):
+        kernel_client._kernel_geospatial_kwargs("false")
+
+
 def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch):
     """Kernel telemetry phase 7 added binding/runtime identity and
     telemetry config kwargs to ``databricks_sql_kernel.Session``."""
@@ -589,7 +677,9 @@ def raise_value_error(_obj):
     kwargs = kernel_client._kernel_telemetry_kwargs(
         {"enable_telemetry": True, "telemetry_batch_size": 17}
     )
-    assert kwargs == {}, f"expected no phase-7 kwargs when signature unreadable, got {kwargs}"
+    assert (
+        kwargs == {}
+    ), f"expected no phase-7 kwargs when signature unreadable, got {kwargs}"
 
 
 def test_execute_command_forwards_parameters_to_bind_param():

tests/unit/test_kernel_result_set.py

@@ -43,11 +43,12 @@ def close(self):
         self.closed = True
 
 
-def _make_rs(handle) -> KernelResultSet:
+def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet:
     # The base ResultSet __init__ takes a `connection` ref it never
     # actually dereferences during these buffer tests, so a Mock is
     # fine.
     connection = MagicMock()
+    connection.disable_pandas = disable_pandas
     backend = MagicMock()
     return KernelResultSet(
         connection=connection,
@@ -94,6 +95,43 @@ 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():
+    wkb = bytes.fromhex("0101000000000000000000F03F0000000000000040")
+    geo_metadata = {
+        b"databricks.type_name": b"GEOMETRY",
+        b"databricks.type_text": b"GEOMETRY(ANY)",
+    }
+
+    string_schema = pa.schema([pa.field("g", pa.string(), metadata=geo_metadata)])
+    string_batch = pa.RecordBatch.from_arrays(
+        [pa.array(["SRID=4326;POINT(1 2)", None], type=pa.string())],
+        schema=string_schema,
+    )
+    string_rows = _make_rs(
+        _FakeKernelHandle(string_schema, [string_batch]), disable_pandas=True
+    ).fetchall()
+    assert [row[0] for row in string_rows] == ["SRID=4326;POINT(1 2)", None]
+
+    binary_type = pa.struct(
+        [
+            pa.field("srid", pa.int32(), nullable=False),
+            pa.field("wkb", pa.binary(), nullable=False),
+        ]
+    )
+    binary_schema = pa.schema([pa.field("g", binary_type, metadata=geo_metadata)])
+    binary_batch = pa.RecordBatch.from_arrays(
+        [pa.array([{"srid": 4326, "wkb": wkb}, None], type=binary_type)],
+        schema=binary_schema,
+    )
+    binary_rs = _make_rs(
+        _FakeKernelHandle(binary_schema, [binary_batch]), disable_pandas=True
+    )
+    assert binary_rs.description[0][1] == "geometry"
+    binary_rows = binary_rs.fetchall()
+    assert binary_rows[0][0] == {"srid": 4326, "wkb": wkb}
+    assert binary_rows[1][0] is None
+
+
 def test_fetchmany_arrow_slices_within_batch(int_schema):
     handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [10, 20, 30, 40])])
     rs = _make_rs(handle)

tests/unit/test_session.py

@@ -476,6 +476,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self):
                 _retry_stop_after_attempts_duration=600.0,
                 _socket_timeout=12.5,
                 _pool_maxsize=41,
+                geospatial_as_string=False,
             )
             try:
                 _, kwargs = mock_kernel_client.call_args
@@ -486,6 +487,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self):
                 assert opts["retry_stop_after_attempts_duration"] == 600.0
                 assert kwargs["request_timeout_secs"] == 12.5
                 assert kwargs["max_connections"] == 41
+                assert kwargs["geospatial_as_string"] is False
             finally:
                 conn.close()