From ca94150ed165f78c751f4e34d7777bb20300001c Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Sat, 14 Nov 2026 03:22:26 +0000 Subject: [PATCH 1/3] Honor standard CloudFetch download settings with the kernel Signed-off-by: Aakash Saravanan --- CHANGELOG.md | 1 + CONNECTION_PARAMETERS.md | 4 +- src/databricks/sql/backend/kernel/client.py | 24 +++++--- src/databricks/sql/client.py | 9 +++ src/databricks/sql/session.py | 1 + tests/unit/test_kernel_client.py | 62 +++++++++++++++++++-- tests/unit/test_session.py | 2 + 7 files changed, 86 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..102777efc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History # Unreleased +- Honor `max_download_threads` on kernel connections, with a default download window of 10 and explicit chunk-setting overrides. - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index c15b4d9d4..e90da3b35 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -150,7 +150,7 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | ------------------------------------- | ------ | :----: | :----: | ------------- | ------------------------------------------------------------------------------------------------------------- | | `use_cloud_fetch` | `bool` | ✅ | ✅ | `True` | Set `False` at connection creation to fetch complete results without cloud-storage downloads. On the kernel path, changing this attribute after connecting does not change result delivery. | -| `max_download_threads` | `int` | ✅ | ❌ | `10` | Worker threads for cloud-fetch downloads. Not forwarded to the kernel. | +| `max_download_threads` | `int` | ✅ | ✅ | `10` | CloudFetch download window; on the kernel, counts downloading and queued chunks ahead of consumption, excluding consumer buffers. Not a byte limit. | | `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. | | `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. | @@ -204,7 +204,7 @@ are forwarded when supported by the installed kernel binding. 2. `experimental_oauth_persistence` (custom OAuth token store). 3. TLS-client-cert *authentication* (`_use_cert_as_auth`) — note the TLS *transport* options (`_tls_*`) themselves **are** honored on both backends. -4. Result-transport tuning: `max_download_threads`, +4. Result-transport tuning: `enable_query_result_lz4_compression`. 5. Arrow-native rendering for `_use_arrow_native_decimals` / `_use_arrow_native_timestamps` (complex types **are** forwarded). diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index df52f34a5..3c8b49699 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -347,6 +347,13 @@ def __init__( self._request_timeout_secs = kwargs.get("request_timeout_secs") self._max_connections = kwargs.get("max_connections") self._use_cloud_fetch = kwargs.get("use_cloud_fetch", True) + self._max_download_threads = kwargs.get("max_download_threads", 10) + if ( + isinstance(self._max_download_threads, bool) + or not isinstance(self._max_download_threads, int) + or self._max_download_threads <= 0 + ): + raise ValueError("max_download_threads must be a positive integer") # Opt-in native GEOMETRY / GEOGRAPHY support. True requests # the canonical Arrow struct and surfaces as ``{"srid": int, "wkb": # bytes}`` through pyarrow; False requests WKT / EWKT strings. This is @@ -448,11 +455,14 @@ def open_session( if self._kernel_session is not None: raise InterfaceError("KernelDatabricksClient already has an open session.") # The kernel consumes client settings locally before forwarding server parameters. - session_conf: Optional[Dict[str, str]] = None - if session_configuration: - session_conf = {k: str(v) for k, v in session_configuration.items()} + session_conf = {k: str(v) for k, v in (session_configuration or {}).items()} + if not any( + k.lower() == "cloudfetch_max_chunks_in_memory" for k in session_conf + ): + session_conf["cloudfetch_max_chunks_in_memory"] = str( + self._max_download_threads + ) if not self._use_cloud_fetch: - session_conf = session_conf or {} session_conf["cloudfetch_enabled"] = "false" # The kwarg builds run INSIDE the try so the ``finally`` scrub # below always fires — including when ``kernel_auth_kwargs`` @@ -1167,11 +1177,7 @@ def get_columns( @property def max_download_threads(self) -> int: - # CloudFetch parallelism lives kernel-side. This property is - # consulted by Thrift code paths that don't run for - # use_kernel=True; return a non-zero default so anything that - # peeks at it does not divide by zero. - return 10 + return self._max_download_threads _STATE_TO_COMMAND_STATE: Dict[str, CommandState] = { diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 613b62386..bf68f914d 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -247,6 +247,15 @@ def __init__( user_agent_entry: `str`, optional A custom tag to append to the User-Agent header. This is typically used by partners to identify their applications.. If not specified, it will use the default user agent PyDatabricksSqlConnector + max_download_threads: `int`, optional (default 10) + CloudFetch download window. With use_kernel=True, bounds downloading + and queued chunks ahead of consumption, excluding consumer buffers; + it is not a byte limit. Must be a positive integer on the kernel + backend and is fixed when connecting. Values above 256 are clamped. + An explicit cloudfetch_max_chunks_in_memory in session_configuration + takes precedence, case-insensitively. Invalid raw overrides warn and + retain the kernel default of 16. + experimental_oauth_persistence: configures preferred storage for persisting oauth tokens. This has to be a class implementing `OAuthPersistence`. When `auth_type` is set to `databricks-oauth` or `azure-oauth` without persisting the oauth token in a diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index d20fe5df0..a15c77546 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -338,6 +338,7 @@ def _create_backend( request_timeout_secs=kwargs.get("_socket_timeout"), max_connections=kwargs.get("_pool_maxsize") or None, use_cloud_fetch=kwargs.get("use_cloud_fetch", True), + max_download_threads=kwargs.get("max_download_threads", 10), enable_geospatial_support=kwargs.get( "enable_geospatial_support", False ), diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index f94aef0c5..686ef14ec 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -400,7 +400,10 @@ def fake_session(**options): ) original = dict(configuration) if configuration is not None else None client.open_session(configuration, catalog=None, schema=None) - assert captured["session_conf"] == expected + assert captured["session_conf"] == { + "cloudfetch_max_chunks_in_memory": "10", + **expected, + } assert configuration == original @@ -1613,11 +1616,58 @@ def test_get_query_state_returns_state_when_no_failure(): # --------------------------------------------------------------------------- -def test_max_download_threads_is_nonzero(): - """Property is consulted by Thrift code paths that don't run for - ``use_kernel=True``; a non-zero default avoids divide-by-zero.""" - c = _make_client() - assert c.max_download_threads > 0 +@pytest.mark.parametrize( + "threads, override", + [ + (1, None), + (4, None), + (1000, None), + (4, "cloudfetch_max_chunks_in_memory"), + (4, "CLOUDFETCH_MAX_CHUNKS_IN_MEMORY"), + ], +) +def test_max_download_threads_configures_kernel(monkeypatch, threads, override): + captured = {} + + def fake_session(**options): + captured.update(options) + return MagicMock(session_id="sess-id") + + monkeypatch.setattr(kernel_client._kernel, "Session", fake_session) + client = kernel_client.KernelDatabricksClient( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + auth_provider=AccessTokenAuthProvider("dapi-test"), + ssl_options=None, + max_download_threads=threads, + ) + configuration = {"TIMEZONE": "UTC"} + if override: + configuration[override] = "2" + original = dict(configuration) + client.open_session(configuration, catalog=None, schema=None) + expected = ( + original + if override + else {**original, "cloudfetch_max_chunks_in_memory": str(threads)} + ) + assert captured["session_conf"] == expected + assert configuration == original + assert client.max_download_threads == threads + + +@pytest.mark.parametrize("threads", [0, -1, None, "4", 1.5, True]) +def test_max_download_threads_rejects_invalid_values(threads): + with pytest.raises( + ValueError, match="max_download_threads must be a positive integer" + ): + kernel_client.KernelDatabricksClient( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + auth_provider=AccessTokenAuthProvider("dapi-test"), + ssl_options=None, + max_download_threads=threads, + ) def test_synthetic_command_id_is_uuid_shaped(): diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 2a555a968..b47d17ddf 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -539,6 +539,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self): _socket_timeout=12.5, _pool_maxsize=41, use_cloud_fetch=False, + max_download_threads=4, enable_geospatial_support=False, ) try: @@ -551,6 +552,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self): assert kwargs["request_timeout_secs"] == 12.5 assert kwargs["max_connections"] == 41 assert kwargs["use_cloud_fetch"] is False + assert kwargs["max_download_threads"] == 4 assert kwargs["enable_geospatial_support"] is False finally: conn.close() From 52290d50cfd6b13210fee357025c69e59a2275c0 Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Sat, 14 Nov 2026 03:32:33 +0000 Subject: [PATCH 2/3] Clamp unbounded Python download settings before native parsing Signed-off-by: Aakash Saravanan --- CHANGELOG.md | 2 +- src/databricks/sql/backend/kernel/client.py | 2 +- src/databricks/sql/client.py | 3 ++- tests/unit/test_kernel_client.py | 3 ++- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 102777efc..5d1f42a8e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ # Release History # Unreleased -- Honor `max_download_threads` on kernel connections, with a default download window of 10 and explicit chunk-setting overrides. +- Honor `max_download_threads` on kernel connections, with a default download window of 10 (previously the kernel default of 16) and explicit chunk-setting overrides. - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 3c8b49699..7c60e7326 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -460,7 +460,7 @@ def open_session( k.lower() == "cloudfetch_max_chunks_in_memory" for k in session_conf ): session_conf["cloudfetch_max_chunks_in_memory"] = str( - self._max_download_threads + min(self._max_download_threads, 256) ) if not self._use_cloud_fetch: session_conf["cloudfetch_enabled"] = "false" diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index bf68f914d..a5d0ac6bd 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -251,7 +251,8 @@ def __init__( CloudFetch download window. With use_kernel=True, bounds downloading and queued chunks ahead of consumption, excluding consumer buffers; it is not a byte limit. Must be a positive integer on the kernel - backend and is fixed when connecting. Values above 256 are clamped. + backend and is fixed when connecting. On the kernel backend, values above + 256 are clamped. An explicit cloudfetch_max_chunks_in_memory in session_configuration takes precedence, case-insensitively. Invalid raw overrides warn and retain the kernel default of 16. diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 686ef14ec..af7667d0b 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1622,6 +1622,7 @@ def test_get_query_state_returns_state_when_no_failure(): (1, None), (4, None), (1000, None), + (10**100, None), (4, "cloudfetch_max_chunks_in_memory"), (4, "CLOUDFETCH_MAX_CHUNKS_IN_MEMORY"), ], @@ -1649,7 +1650,7 @@ def fake_session(**options): expected = ( original if override - else {**original, "cloudfetch_max_chunks_in_memory": str(threads)} + else {**original, "cloudfetch_max_chunks_in_memory": str(min(threads, 256))} ) assert captured["session_conf"] == expected assert configuration == original From a8fb2980b6a5ab0bb1d52c5e06a81b918f0971be Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Sat, 14 Nov 2026 04:09:46 +0000 Subject: [PATCH 3/3] Document kernel download setting validation timing Signed-off-by: Aakash Saravanan --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d1f42a8e..bf7230fcb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ # Release History # Unreleased -- Honor `max_download_threads` on kernel connections, with a default download window of 10 (previously the kernel default of 16) and explicit chunk-setting overrides. +- Honor `max_download_threads` on kernel connections, with a default download window of 10 (previously the kernel default of 16) and explicit chunk-setting overrides. Invalid values raise `ValueError` when connecting with `use_kernel=True`. - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most.