perf(thrift): request about one CloudFetch file per download thread

CHANGELOG.md

@@ -3,6 +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 (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

@@ -20,13 +20,17 @@
 from databricks.sql.utils import (
     ColumnTable,
     ColumnQueue,
+    ThriftCloudFetchQueue,
     concat_table_chunks,
 )
 from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse
 from databricks.sql.telemetry.models.event import StatementType
 
 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):
     """
@@ -59,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
@@ -216,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)
@@ -229,6 +236,8 @@ def __init__(
         # Initialize ThriftResultSet-specific attributes
         self._use_cloud_fetch = use_cloud_fetch
         self.has_more_rows = has_more_rows
+        self._max_download_threads = max_download_threads
+        self._max_link_bytes = 0
 
         # Build the results queue if t_row_set is provided
         results_queue = None
@@ -251,6 +260,7 @@ def __init__(
             )
             if t_row_set.resultLinks:
                 self.num_chunks += len(t_row_set.resultLinks)
+                self._note_link_sizes(t_row_set.resultLinks)
 
         # Call parent constructor with common attributes
         super().__init__(
@@ -274,11 +284,32 @@ def __init__(
         if not self.results:
             self._fill_results_buffer()
 
+    def _note_link_sizes(self, links) -> None:
+        for link in links:
+            # The reader doesn't enforce required fields; a missing size only
+            # disables the sizing heuristic.
+            size = link.bytesNum or 0
+            if size > self._max_link_bytes:
+                self._max_link_bytes = size
+
+    def _fetch_max_bytes(self) -> int:
+        # The server packs CloudFetch links into a FetchResults reply up to maxBytes of
+        # 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.
+        # 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):
         results, has_more_rows, result_links_count = self.backend.fetch_results(
             command_id=self.command_id,
             max_rows=self.arraysize,
-            max_bytes=self.buffer_size_bytes,
+            max_bytes=self._fetch_max_bytes(),
             expected_row_start_offset=self._next_row_index,
             lz4_compressed=self.lz4_compressed,
             arrow_schema_bytes=self._arrow_schema_bytes,
@@ -289,6 +320,8 @@ def _fill_results_buffer(self):
         self.results = results
         self.has_more_rows = has_more_rows
         self.num_chunks += result_links_count
+        if isinstance(results, ThriftCloudFetchQueue):
+            self._note_link_sizes(results.result_links)
 
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]

tests/unit/test_fetches.py

@@ -1,6 +1,6 @@
 import unittest
 import pytest
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 try:
     import pyarrow as pa
@@ -9,9 +9,10 @@
 
 import databricks.sql.client as client
 from databricks.sql.backend.types import ExecuteResponse
-from databricks.sql.utils import ArrowQueue
+from databricks.sql.utils import ArrowQueue, ThriftCloudFetchQueue
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.result_set import ThriftResultSet
+from databricks.sql.thrift_api.TCLIService import ttypes
 
 
 class _StubArrowQueue:
@@ -391,5 +392,168 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self):
         self.assertEqual(result.num_rows, 0)
 
 
+class CloudFetchBatchSizeTests(unittest.TestCase):
+    """maxBytes sent by ThriftResultSet._fill_results_buffer."""
+
+    MIB = 2**20
+    BUFFER = 100 * MIB
+
+    @staticmethod
+    def _link(bytes_num):
+        return ttypes.TSparkArrowResultLink(
+            fileLink="https://example.com/file",
+            expiryTime=0,
+            startRowOffset=0,
+            rowCount=1,
+            bytesNum=bytes_num,
+        )
+
+    @classmethod
+    def _cloudfetch_queue(cls, *link_sizes):
+        queue = Mock(spec=ThriftCloudFetchQueue)
+        queue.result_links = [cls._link(b) for b in link_sizes]
+        return queue
+
+    def _make_result_set(
+        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 = []
+        queues = list(fetched_queues)
+
+        def fetch_results(**kwargs):
+            calls.append(kwargs)
+            queue = queues.pop(0)
+            return queue, bool(queues), len(getattr(queue, "result_links", []))
+
+        backend = Mock(spec=ThriftDatabricksClient)
+        backend.fetch_results = fetch_results
+        t_row_set = None
+        if initial_link_sizes is not None:
+            t_row_set = ttypes.TRowSet(
+                startRowOffset=0,
+                rows=[],
+                resultLinks=[self._link(b) for b in initial_link_sizes],
+            )
+        with patch(
+            "databricks.sql.utils.ThriftResultSetQueueFactory.build_queue",
+            return_value=Mock(),
+        ):
+            rs = ThriftResultSet(
+                connection=Mock(),
+                execute_response=ExecuteResponse(
+                    command_id=Mock(),
+                    status=None,
+                    has_been_closed_server_side=False,
+                    description=[],
+                    result_format=ttypes.TSparkRowSetType.URL_BASED_SET,
+                ),
+                thrift_client=backend,
+                buffer_size_bytes=buffer_size_bytes,
+                arraysize=1000,
+                t_row_set=t_row_set,
+                max_download_threads=max_download_threads,
+            )
+        return rs, calls
+
+    def test_requests_one_link_per_download_thread(self):
+        rs, calls = self._make_result_set(
+            initial_link_sizes=[20 * self.MIB, 20 * self.MIB, 5 * self.MIB],
+            fetched_queues=[
+                self._cloudfetch_queue(21 * self.MIB),
+                self._cloudfetch_queue(),
+            ],
+        )
+
+        rs._fill_results_buffer()
+        rs._fill_results_buffer()
+
+        self.assertEqual(
+            [c["max_bytes"] for c in calls], [200 * self.MIB, 210 * self.MIB]
+        )
+        self.assertEqual([c["max_rows"] for c in calls], [1000, 1000])
+
+    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=15,
+        )
+
+        rs._fill_results_buffer()
+
+        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_links_without_size_keep_buffer_size(self):
+        rs, calls = self._make_result_set(
+            initial_link_sizes=[None],
+            fetched_queues=[self._cloudfetch_queue(None), self._cloudfetch_queue()],
+        )
+
+        rs._fill_results_buffer()
+        rs._fill_results_buffer()
+
+        self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, 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(
+            initial_link_sizes=[20 * self.MIB],
+            fetched_queues=[self._cloudfetch_queue()],
+            max_download_threads=2,
+        )
+
+        rs._fill_results_buffer()
+
+        self.assertEqual(calls[0]["max_bytes"], self.BUFFER)
+
+    def test_inline_results_keep_buffer_size(self):
+        # No initial rows: the constructor fetches the first batch itself.
+        rs, calls = self._make_result_set(
+            fetched_queues=[Mock(spec=ArrowQueue), Mock(spec=ArrowQueue)]
+        )
+
+        rs._fill_results_buffer()
+
+        self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, self.BUFFER])
+
+    def test_first_fetch_without_links_uses_buffer_size(self):
+        rs, calls = self._make_result_set(
+            fetched_queues=[
+                self._cloudfetch_queue(20 * self.MIB),
+                self._cloudfetch_queue(),
+            ]
+        )
+
+        rs._fill_results_buffer()
+
+        self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, 200 * self.MIB])
+
+
 if __name__ == "__main__":
     unittest.main()