From 2b566034e9a2ab9297a3ca6ff7f2a1c688f44e2f Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Wed, 30 Sep 2026 06:18:27 +0000 Subject: [PATCH] feat: support geospatial results on the kernel path Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- CHANGELOG.md | 3 + CONNECTION_PARAMETERS.md | 1 + KERNEL_REV | 2 +- src/databricks/sql/backend/kernel/client.py | 31 +++++++ src/databricks/sql/client.py | 8 ++ src/databricks/sql/session.py | 1 + tests/unit/test_kernel_client.py | 92 ++++++++++++++++++++- tests/unit/test_kernel_result_set.py | 40 ++++++++- tests/unit/test_session.py | 2 + 9 files changed, 177 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bce9e16e4..bf9065694 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index eaf69c57b..a146fbbc3 100644 --- a/CONNECTION_PARAMETERS.md +++ b/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. | diff --git a/KERNEL_REV b/KERNEL_REV index f22f21685..f6a0e2687 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -80f2aee7d884994d7b0af9a9ea6078872859a9cd +ad3bc6993bca95b810839feade77ccd0ab98ece5 diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index ab1ef0dd7..771fb72ee 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/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, ) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 3e06aa39b..0cb914b2f 100755 --- a/src/databricks/sql/client.py +++ b/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. diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index db512dd0b..b95768005 100644 --- a/src/databricks/sql/session.py +++ b/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, ) diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 3eecb0145..d9b04fab6 100644 --- a/tests/unit/test_kernel_client.py +++ b/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(): diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index 9ec69380a..81ef6a91b 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/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) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 69522bc60..ff090ad5f 100644 --- a/tests/unit/test_session.py +++ b/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()