Honor CloudFetch download settings on kernel connections

CHANGELOG.md

@@ -1,6 +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. 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.

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).

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(
+                min(self._max_download_threads, 256)
+            )
         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] = {

src/databricks/sql/client.py

@@ -247,6 +247,16 @@ 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. 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.
+
             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

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
                 ),

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,59 @@ 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),
+        (10**100, 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(min(threads, 256))}
+    )
+    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():

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()