Skip to content

Commit afc6cb8

Browse files
authored
feat: support geospatial results on the kernel path (#990)
* feat: support geospatial results on the kernel path Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * feat: align geospatial support configuration Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * test: handle pending kernel geospatial release Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * test: import kernel compatibility error Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * fix(thrift): report geospatial cursor metadata Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * fix: preserve geospatial result metadata and shapes Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * chore: update geospatial kernel pin Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * fix: default Python geospatial results to strings Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * docs: correct geospatial support default Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * test: match kernel table-type filters by exact case Align Kernel E2E with case-sensitive filtering from databricks-sql-kernel#393. Check uppercase matches and lowercase/mixed-case empty results, and update stale filter comments. Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> * ci: refresh merge-queue check records Refresh status contexts after GitHub reported the Kernel E2E requirement as expected despite matching successful checks. The source tree is unchanged from the validated fix in 3c800a5. Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --------- Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com>
1 parent b1978d4 commit afc6cb8

12 files changed

Lines changed: 336 additions & 27 deletions

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Release History
22

33
# Unreleased
4+
- Add the kernel-only `enable_geospatial_support` connection option (default `False`). GEOMETRY / GEOGRAPHY results are exposed as `{"srid": int, "wkb": bytes}` values when true or WKT / EWKT strings when false.
45
- Add the `escape_catalog_name` connection parameter (default `False`) to opt in to escaping unescaped `_` in metadata catalog names on Thrift and kernel connections. Existing escapes and `%` wildcards are preserved.
56

67
# 4.6.0 (2026-09-24)

‎CONNECTION_PARAMETERS.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@ to change without notice.
153153
| `max_download_threads` | `int` | ✅ | ❌ | `10` | Worker threads for cloud-fetch downloads. Not forwarded to the kernel. |
154154
| `enable_query_result_lz4_compression` | `bool` | ✅ | ❌ | `True` | LZ4-compress result payloads. Not forwarded; the kernel handles compression internally. |
155155
| `_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. |
156+
| `enable_geospatial_support` | `bool` | ❌ | ✅ | `False` | Return GEOMETRY / GEOGRAPHY as `{"srid": int, "wkb": bytes}` values when `True`, or as WKT / EWKT strings when `False`. This is a local result conversion and is never forwarded to SEA. |
156157
| `_use_arrow_native_complex_types` | `bool` | ✅ | ✅ | `True` | Return `ARRAY`/`MAP`/`STRUCT` as native Arrow types instead of JSON strings. Forwarded to the kernel. |
157158
| `_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. |
158159
| `_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‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
ab1e6a707e4a15a3a7a8c78d8fdd12c752c51102
1+
9c2e2e6182fc7285bcfff39b2177941bc804897c

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

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,26 @@ def _kernel_session_accepts_kwarg(name: str) -> bool:
170170
return name in params
171171

172172

173+
def _kernel_geospatial_kwargs(value: bool) -> Dict[str, bool]:
174+
"""Build the explicit native geospatial-support kwarg.
175+
176+
The value is always passed explicitly so the driver and kernel contracts
177+
cannot drift. Older kernel wheels do not declare
178+
``enable_geospatial_support`` and must fail clearly rather than silently
179+
return a different public value shape.
180+
"""
181+
if not isinstance(value, bool):
182+
raise ValueError(
183+
"enable_geospatial_support must be a bool; " f"got {type(value).__name__}"
184+
)
185+
if not _kernel_session_accepts_kwarg("enable_geospatial_support"):
186+
raise NotSupportedError(
187+
"enable_geospatial_support requires a newer databricks-sql-kernel "
188+
"wheel that exposes geospatial result representation support."
189+
)
190+
return {"enable_geospatial_support": value}
191+
192+
173193
def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]:
174194
"""Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``.
175195
@@ -263,6 +283,12 @@ def __init__(
263283
# The kernel binding owns type and range validation.
264284
self._request_timeout_secs = kwargs.get("request_timeout_secs")
265285
self._max_connections = kwargs.get("max_connections")
286+
# Opt-in native GEOMETRY / GEOGRAPHY support. True requests
287+
# the canonical Arrow struct and surfaces as ``{"srid": int, "wkb":
288+
# bytes}`` through pyarrow; False requests WKT / EWKT strings. This is
289+
# intentionally separate from ``session_configuration``: it is never
290+
# forwarded to SEA.
291+
self._enable_geospatial_support = kwargs.get("enable_geospatial_support", False)
266292
# Kernel telemetry phase 7 adds binding/runtime identity and
267293
# telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
268294
self._telemetry_options = kwargs.get("telemetry_options") or {}
@@ -381,6 +407,9 @@ def open_session(
381407
# kernel's ``retry_*`` kwargs. Empty when at defaults.
382408
retry_kwargs = _kernel_retry_kwargs(self._retry_options)
383409
telemetry_kwargs = _kernel_telemetry_kwargs(self._telemetry_options)
410+
geospatial_kwargs = _kernel_geospatial_kwargs(
411+
self._enable_geospatial_support
412+
)
384413
max_connections_kwargs: Dict[str, Any] = {}
385414
if _kernel_session_accepts_kwarg("max_connections"):
386415
max_connections_kwargs["max_connections"] = self._max_connections
@@ -428,6 +457,7 @@ def open_session(
428457
**tls_kwargs,
429458
**retry_kwargs,
430459
**telemetry_kwargs,
460+
**geospatial_kwargs,
431461
**max_connections_kwargs,
432462
**http_headers_kwargs,
433463
)
@@ -1015,11 +1045,9 @@ def get_tables(
10151045
if self._kernel_session is None:
10161046
raise InterfaceError("get_tables requires an open session.")
10171047
try:
1018-
# ``table_types`` is filtered kernel-side (the kernel applies
1019-
# it to the reshaped result, case-insensitively as of the
1020-
# batch-3 kernel change), so we forward it and let the kernel
1021-
# do the work — no connector-side drain + refilter. Passing it
1022-
# through preserves streaming for large schemas.
1048+
# The kernel filters ``table_types`` with case-sensitive
1049+
# matching for Thrift parity. Forward the values unchanged
1050+
# to preserve those semantics and streaming for large schemas.
10231051
stream = self._kernel_session.metadata().list_tables(
10241052
catalog=escape_metadata_catalog_name(
10251053
_catalog_or_none(catalog_name), self._escape_catalog_name

‎src/databricks/sql/backend/thrift_backend.py‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import errno
44
import logging
55
import math
6+
import re
67
import time
78
import threading
89
from typing import Dict, List, Optional, Union, Any, TYPE_CHECKING
@@ -807,38 +808,50 @@ def _col_to_description(col, field=None, host_url=None):
807808
else:
808809
precision, scale = None, None
809810

810-
# Extract variant type from field if available
811+
# Extract logical types that Thrift surfaces as strings from Arrow metadata.
811812
if field is not None:
812813
try:
813-
# Check for variant type in metadata
814814
if field.metadata and b"Spark:DataType:SqlName" in field.metadata:
815815
sql_type = field.metadata.get(b"Spark:DataType:SqlName")
816-
if sql_type == b"VARIANT":
817-
cleaned_type = "variant"
816+
logical_type = re.fullmatch(
817+
rb"(variant|geometry|geography)(?:\([^()]*\))?",
818+
sql_type.strip(),
819+
re.IGNORECASE,
820+
)
821+
if logical_type:
822+
cleaned_type = logical_type.group(1).decode("ascii").lower()
818823
except Exception as e:
819-
logger.debug(f"Could not extract variant type from field: {e}")
824+
logger.debug(f"Could not extract logical type from field: {e}")
820825

821826
return col.columnName, cleaned_type, None, None, precision, scale, None
822827

823828
@staticmethod
824829
def _hive_schema_to_description(t_table_schema, schema_bytes=None, host_url=None):
830+
arrow_fields = []
825831
field_dict = {}
826832
if pyarrow and schema_bytes:
827833
try:
828834
arrow_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes))
829-
# Build a dictionary mapping column names to fields
830-
for field in arrow_schema:
835+
arrow_fields = list(arrow_schema)
836+
# Retain a name fallback for malformed/partial schemas, but use
837+
# result-column position whenever it is available. SQL permits
838+
# duplicate aliases, for which a name lookup is ambiguous.
839+
for field in arrow_fields:
831840
field_dict[field.name] = field
832841
except Exception as e:
833842
logger.debug(f"Could not parse arrow schema: {e}")
834843

835844
return [
836845
ThriftDatabricksClient._col_to_description(
837846
col,
838-
field_dict.get(col.columnName) if field_dict else None,
847+
(
848+
arrow_fields[index]
849+
if index < len(arrow_fields)
850+
else field_dict.get(col.columnName)
851+
),
839852
host_url,
840853
)
841-
for col in t_table_schema.columns
854+
for index, col in enumerate(t_table_schema.columns)
842855
]
843856

844857
def _results_message_to_execute_response(

‎src/databricks/sql/client.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,13 @@ def __init__(
194194
decision. This is an intentional divergence from the
195195
Thrift/SEA paths, where an explicit ``True`` can still
196196
be suppressed by the feature flag.
197+
:param enable_geospatial_support: `bool`, optional (default is False)
198+
Kernel backend only. Controls the public representation of
199+
``GEOMETRY`` and ``GEOGRAPHY`` result values. ``True`` returns
200+
``{"srid": int, "wkb": bytes}``; ``False`` returns WKT / EWKT
201+
strings (for example ``"SRID=4326;POINT(1 2)"``). The
202+
conversion is local to the kernel/driver and this option is
203+
never sent to the SQL Execution API.
197204
:param use_hybrid_disposition: `bool`, optional (default is False)
198205
Use the hybrid disposition instead of the inline disposition.
199206
:param server_hostname: Databricks instance host name.

‎src/databricks/sql/session.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,9 @@ def _create_backend(
332332
retry_options=kernel_retry_options,
333333
request_timeout_secs=kwargs.get("_socket_timeout"),
334334
max_connections=kwargs.get("_pool_maxsize") or None,
335+
enable_geospatial_support=kwargs.get(
336+
"enable_geospatial_support", False
337+
),
335338
telemetry_options=kernel_telemetry_options,
336339
escape_catalog_name=kwargs.get("escape_catalog_name", False),
337340
)

‎tests/e2e/test_kernel_backend.py‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -383,12 +383,12 @@ def test_columns_with_empty_string_filter_matches_nothing(conn, empty_filter):
383383
assert cur.fetchall() == []
384384

385385

386-
def test_tables_table_types_filter_is_case_insensitive(conn):
387-
"""Lowercase ``table_types=['view']`` / uppercase ``['TABLE']``
388-
each match the right object regardless of case — locks the
389-
kernel-side case-insensitive ``table_types`` match (batch-3 kernel
390-
B2) end-to-end plus the connector drain removal (the filter now
391-
runs kernel-side, not client-side).
386+
def test_tables_table_types_filter_is_case_sensitive(conn):
387+
"""Table types match exact case, consistent with Thrift.
388+
389+
Uppercase ``['VIEW']`` / ``['TABLE']`` match the corresponding
390+
object; lowercase and mixed-case values match no objects. The
391+
kernel applies the filter to values forwarded by the connector.
392392
393393
Self-contained: creates a table + a view over it in the session's
394394
default (writable) schema, scopes the lookup to a unique name
@@ -411,12 +411,12 @@ def _names_and_types():
411411
ni, ti = cols.index("TABLE_NAME"), cols.index("TABLE_TYPE")
412412
return {(r[ni], r[ti]) for r in rows}
413413

414-
# Lowercase 'view' must match the VIEW (and only it).
414+
# Uppercase 'VIEW' must match the VIEW (and only it).
415415
cur.tables(
416416
catalog_name=cat,
417417
schema_name=sch,
418418
table_name=name_pat,
419-
table_types=["view"],
419+
table_types=["VIEW"],
420420
)
421421
assert _names_and_types() == {(vw, "VIEW")}
422422

@@ -428,6 +428,16 @@ def _names_and_types():
428428
table_types=["TABLE"],
429429
)
430430
assert _names_and_types() == {(tbl, "TABLE")}
431+
432+
# Lowercase and mixed-case requests must not match either type.
433+
for table_type in ("view", "View", "table", "Table"):
434+
cur.tables(
435+
catalog_name=cat,
436+
schema_name=sch,
437+
table_name=name_pat,
438+
table_types=[table_type],
439+
)
440+
assert _names_and_types() == set(), table_type
431441
finally:
432442
cur.execute(f"DROP VIEW IF EXISTS {vw}")
433443
cur.execute(f"DROP TABLE IF EXISTS {tbl}")

‎tests/unit/test_kernel_client.py‎

Lines changed: 93 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -392,6 +392,94 @@ def fake_session(**kw):
392392
assert captured["max_connections"] == max_connections
393393

394394

395+
@pytest.mark.parametrize("enabled", [True, False])
396+
def test_open_session_passes_geospatial_representation_to_kernel(
397+
monkeypatch, enabled
398+
):
399+
captured = {}
400+
401+
def fake_session(*, enable_geospatial_support=True, **kw):
402+
captured["enable_geospatial_support"] = enable_geospatial_support
403+
sess = MagicMock()
404+
sess.session_id = "sess-id"
405+
return sess
406+
407+
monkeypatch.setattr(kernel_client._kernel, "Session", fake_session)
408+
c = kernel_client.KernelDatabricksClient(
409+
server_hostname="example.cloud.databricks.com",
410+
http_path="/sql/1.0/warehouses/abc",
411+
auth_provider=AccessTokenAuthProvider("dapi-test"),
412+
ssl_options=None,
413+
enable_geospatial_support=enabled,
414+
)
415+
416+
c.open_session(session_configuration=None, catalog=None, schema=None)
417+
418+
assert captured["enable_geospatial_support"] is enabled
419+
420+
421+
def test_open_session_disables_native_geospatial_support_by_default(monkeypatch):
422+
captured = {}
423+
424+
def fake_session(**kw):
425+
captured.update(kw)
426+
sess = MagicMock()
427+
sess.session_id = "sess-id"
428+
return sess
429+
430+
monkeypatch.setattr(kernel_client._kernel, "Session", fake_session)
431+
c = kernel_client.KernelDatabricksClient(
432+
server_hostname="example.cloud.databricks.com",
433+
http_path="/sql/1.0/warehouses/abc",
434+
auth_provider=AccessTokenAuthProvider("dapi-test"),
435+
ssl_options=None,
436+
)
437+
438+
c.open_session(session_configuration=None, catalog=None, schema=None)
439+
440+
assert captured["enable_geospatial_support"] is False
441+
442+
443+
def test_open_session_rejects_explicit_geospatial_option_with_old_kernel(
444+
monkeypatch,
445+
):
446+
def fake_session_without_geospatial(
447+
host,
448+
http_path,
449+
*,
450+
catalog=None,
451+
schema=None,
452+
session_conf=None,
453+
complex_types_as_json=False,
454+
intervals_as_string=False,
455+
request_timeout_secs=None,
456+
auth_type=None,
457+
access_token=None,
458+
):
459+
sess = MagicMock()
460+
sess.session_id = "sess-id"
461+
return sess
462+
463+
monkeypatch.setattr(
464+
kernel_client._kernel, "Session", fake_session_without_geospatial
465+
)
466+
c = kernel_client.KernelDatabricksClient(
467+
server_hostname="example.cloud.databricks.com",
468+
http_path="/sql/1.0/warehouses/abc",
469+
auth_provider=AccessTokenAuthProvider("dapi-test"),
470+
ssl_options=None,
471+
enable_geospatial_support=False,
472+
)
473+
474+
with pytest.raises(NotSupportedError, match="newer databricks-sql-kernel"):
475+
c.open_session(session_configuration=None, catalog=None, schema=None)
476+
477+
478+
def test_geospatial_option_rejects_non_bool():
479+
with pytest.raises(ValueError, match="must be a bool"):
480+
kernel_client._kernel_geospatial_kwargs("false")
481+
482+
395483
def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch):
396484
"""Kernel telemetry phase 7 added binding/runtime identity and
397485
telemetry config kwargs to ``databricks_sql_kernel.Session``."""
@@ -490,6 +578,7 @@ def fake_session_without_optional_kwargs(
490578
catalog=None,
491579
schema=None,
492580
session_conf=None,
581+
enable_geospatial_support=True,
493582
complex_types_as_json=False,
494583
intervals_as_string=False,
495584
request_timeout_secs=None,
@@ -589,7 +678,9 @@ def raise_value_error(_obj):
589678
kwargs = kernel_client._kernel_telemetry_kwargs(
590679
{"enable_telemetry": True, "telemetry_batch_size": 17}
591680
)
592-
assert kwargs == {}, f"expected no phase-7 kwargs when signature unreadable, got {kwargs}"
681+
assert (
682+
kwargs == {}
683+
), f"expected no phase-7 kwargs when signature unreadable, got {kwargs}"
593684

594685

595686
def test_execute_command_forwards_parameters_to_bind_param():
@@ -1535,7 +1626,7 @@ def _make_tables_stream() -> MagicMock:
15351626

15361627
def test_get_tables_forwards_table_types_to_kernel():
15371628
"""``table_types`` is forwarded verbatim to the kernel's
1538-
``list_tables`` (which filters case-insensitively) — the
1629+
``list_tables`` (which filters case-sensitively) — the
15391630
connector no longer drains + refilters client-side. The stream
15401631
flows back through the normal ``KernelResultSet`` path unchanged."""
15411632
c = _make_client()

0 commit comments

Comments
 (0)