Cap the CloudFetch FetchResults budget at 4x buffer_size_bytes

Without a cap the budget grows with max_download_threads and the largest
link seen, and it ignored a deliberately small buffer_size_bytes. The
default path (about 210 MB, 10 links) is unchanged. Also document that
buffer_size_bytes is no longer the max fetch size for CloudFetch.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>

Author: vuanhphung

CHANGELOG.md

@@ -3,7 +3,7 @@
 # Unreleased
 - 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.
-- Speed up large CloudFetch results on Thrift: each FetchResults call now asks for about one result file per download thread instead of 100 MiB of files, so all download threads stay busy and the result takes fewer round trips.
+- Speed up large CloudFetch results on Thrift: each FetchResults call now asks for about one result file per download thread (at most 4x `buffer_size_bytes`) instead of 100 MiB of files, so all download threads stay busy and the result takes fewer round trips.
 
 # 4.9.0 (2026-11-04)
 - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes.

src/databricks/sql/result_set.py

@@ -28,6 +28,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Upper bound on a Thrift CloudFetch FetchResults request, in multiples of buffer_size_bytes.
+_CLOUDFETCH_MAX_BUFFERS_PER_FETCH = 4
+
 
 class ResultSet(ABC):
     """
@@ -60,7 +63,9 @@ def __init__(
             :param connection: The parent connection that was used to execute this command
             :param backend: The specialised backend client to be invoked in the fetch phase
             :param arraysize: The max number of rows to fetch at a time (PEP-249)
-            :param buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch
+            :param buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch.
+                Thrift CloudFetch fetches may request up to 4x this (see
+                ThriftResultSet._fetch_max_bytes).
             :param command_id: The command ID
             :param status: The command status
             :param has_been_closed_server_side: Whether the command has been closed on the server
@@ -217,7 +222,8 @@ def __init__(
             :param connection: The parent connection
             :param execute_response: Response from the execute command
             :param thrift_client: The ThriftDatabricksClient instance for direct access
-            :param buffer_size_bytes: Buffer size for fetching results
+            :param buffer_size_bytes: Buffer size for fetching results. CloudFetch
+                fetches request between 1x and 4x this (see _fetch_max_bytes).
             :param arraysize: Default number of rows to fetch
             :param use_cloud_fetch: Whether to use cloud fetch for retrieving results
             :param t_row_set: The TRowSet containing result data (if available)
@@ -288,8 +294,12 @@ def _fetch_max_bytes(self) -> int:
         # file data, so buffer_size_bytes alone yields only a few links per round trip.
         # Ask for about one link per download thread: every link in the batch then starts
         # downloading immediately, so none sits waiting while it ages towards expiry.
-        return max(
-            self.buffer_size_bytes, self._max_download_threads * self._max_link_bytes
+        # The cap bounds the batch (and buffered downloads) for a large
+        # max_download_threads, an oversized link, or a deliberately small buffer.
+        wanted = self._max_download_threads * self._max_link_bytes
+        return min(
+            max(self.buffer_size_bytes, wanted),
+            _CLOUDFETCH_MAX_BUFFERS_PER_FETCH * self.buffer_size_bytes,
         )
 
     def _fill_results_buffer(self):

tests/unit/test_fetches.py

@@ -415,7 +415,11 @@ def _cloudfetch_queue(cls, *link_sizes):
         return queue
 
     def _make_result_set(
-        self, initial_link_sizes=None, fetched_queues=(), max_download_threads=10
+        self,
+        initial_link_sizes=None,
+        fetched_queues=(),
+        max_download_threads=10,
+        buffer_size_bytes=BUFFER,
     ):
         """Returns the result set and the kwargs of every fetch_results call."""
         calls = []
@@ -449,7 +453,7 @@ def fetch_results(**kwargs):
                     result_format=ttypes.TSparkRowSetType.URL_BASED_SET,
                 ),
                 thrift_client=backend,
-                buffer_size_bytes=self.BUFFER,
+                buffer_size_bytes=buffer_size_bytes,
                 arraysize=1000,
                 t_row_set=t_row_set,
                 max_download_threads=max_download_threads,
@@ -477,12 +481,34 @@ def test_scales_with_max_download_threads(self):
         rs, calls = self._make_result_set(
             initial_link_sizes=[20 * self.MIB],
             fetched_queues=[self._cloudfetch_queue()],
-            max_download_threads=25,
+            max_download_threads=15,
         )
 
         rs._fill_results_buffer()
 
-        self.assertEqual(calls[0]["max_bytes"], 500 * self.MIB)
+        self.assertEqual(calls[0]["max_bytes"], 300 * self.MIB)
+
+    def test_caps_at_four_buffers(self):
+        rs, calls = self._make_result_set(
+            initial_link_sizes=[20 * self.MIB],
+            fetched_queues=[self._cloudfetch_queue()],
+            max_download_threads=50,
+        )
+
+        rs._fill_results_buffer()
+
+        self.assertEqual(calls[0]["max_bytes"], 4 * self.BUFFER)
+
+    def test_small_buffer_size_caps_batch(self):
+        rs, calls = self._make_result_set(
+            initial_link_sizes=[20 * self.MIB],
+            fetched_queues=[self._cloudfetch_queue()],
+            buffer_size_bytes=10 * self.MIB,
+        )
+
+        rs._fill_results_buffer()
+
+        self.assertEqual(calls[0]["max_bytes"], 40 * self.MIB)
 
     def test_never_requests_less_than_buffer_size(self):
         rs, calls = self._make_result_set(