perf(thrift): prefetch CloudFetch links while earlier files download

CHANGELOG.md

@@ -4,6 +4,7 @@
 - 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.
+- Speed up large CloudFetch results on Thrift further by requesting the next result links while earlier files are still downloading, and stop downloads promptly when a result set is closed early.
 
 # 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/backend/thrift_backend.py

@@ -6,7 +6,7 @@
 import re
 import time
 import threading
-from typing import Dict, List, Optional, Union, Any, TYPE_CHECKING
+from typing import Dict, List, Optional, Tuple, Union, Any, TYPE_CHECKING
 from uuid import UUID
 
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
@@ -160,6 +160,9 @@ def __init__(
         #  (defaults to None)
         # max_download_threads
         #  Number of threads for handling cloud fetch downloads. Defaults to 10
+        # _cloudfetch_prefetch
+        #  Request the next CloudFetch links while earlier files still download.
+        #  False fetches one batch of links at a time. Defaults to True
 
         logger.debug(
             "ThriftBackend.__init__(server_hostname=%s, port=%s, http_path=%s)"
@@ -192,6 +195,10 @@ def __init__(
 
         # Cloud fetch
         self._max_download_threads = kwargs.get("max_download_threads", 10)
+        prefetch = kwargs.get("_cloudfetch_prefetch", True)
+        if isinstance(prefetch, str):  # e.g. from a SQLAlchemy URL query
+            prefetch = prefetch.strip().lower() not in ("false", "0", "no", "off")
+        self._cloudfetch_prefetch = bool(prefetch)
 
         self._ssl_options = ssl_options
         self._auth_provider = auth_provider
@@ -1001,6 +1008,7 @@ def get_execution_result(
             t_row_set=resp.results,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            cloudfetch_prefetch=self._cloudfetch_prefetch,
             has_more_rows=has_more_rows,
         )
 
@@ -1155,6 +1163,7 @@ def execute_command(
                 t_row_set=t_row_set,
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
+                cloudfetch_prefetch=self._cloudfetch_prefetch,
                 has_more_rows=has_more_rows,
             )
 
@@ -1193,6 +1202,7 @@ def get_catalogs(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            cloudfetch_prefetch=self._cloudfetch_prefetch,
             has_more_rows=has_more_rows,
         )
 
@@ -1239,6 +1249,7 @@ def get_schemas(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            cloudfetch_prefetch=self._cloudfetch_prefetch,
             has_more_rows=has_more_rows,
         )
 
@@ -1289,6 +1300,7 @@ def get_tables(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            cloudfetch_prefetch=self._cloudfetch_prefetch,
             has_more_rows=has_more_rows,
         )
 
@@ -1339,6 +1351,7 @@ def get_columns(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            cloudfetch_prefetch=self._cloudfetch_prefetch,
             has_more_rows=has_more_rows,
         )
 
@@ -1374,17 +1387,13 @@ def _handle_execute_response_async(self, resp, cursor):
         cursor.active_command_id = command_id
         self._check_direct_results_for_error(resp.directResults, self._host)
 
-    def fetch_results(
+    def _fetch_next(
         self,
         command_id: CommandId,
         max_rows: int,
         max_bytes: int,
         expected_row_start_offset: int,
-        lz4_compressed: bool,
-        arrow_schema_bytes,
-        description,
-        chunk_id: int,
-        use_cloud_fetch=True,
+        retryable: bool,
     ):
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
@@ -1404,14 +1413,59 @@ def fetch_results(
         )
 
         # Fetch results in Inline mode with FETCH_NEXT orientation are not idempotent and hence not retried
-        resp = self.make_request(self._client.FetchResults, req, use_cloud_fetch)
+        resp = self.make_request(self._client.FetchResults, req, retryable)
         if resp.results.startRowOffset > expected_row_start_offset:
             raise DataError(
                 "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batches must have been skipped".format(
                     expected_row_start_offset, resp.results.startRowOffset
                 ),
                 host_url=self._host,
             )
+        return resp
+
+    def fetch_result_links(
+        self,
+        command_id: CommandId,
+        max_rows: int,
+        max_bytes: int,
+        expected_row_start_offset: int,
+    ) -> Tuple[List[ttypes.TSparkArrowResultLink], bool]:
+        """Fetch the next CloudFetch links of a result: (links, has_more_rows)."""
+        resp = self._fetch_next(
+            command_id, max_rows, max_bytes, expected_row_start_offset, True
+        )
+        results, metadata = resp.results, resp.resultSetMetadata
+        # Rows in any other form would be dropped silently, so fail loudly instead
+        if (
+            results.arrowBatches
+            or results.columns
+            or (
+                metadata is not None
+                and metadata.resultFormat is not None
+                and metadata.resultFormat != ttypes.TSparkRowSetType.URL_BASED_SET
+            )
+        ):
+            raise DataError(
+                "Expected only CloudFetch links when fetching more results of a CloudFetch result",
+                host_url=self._host,
+            )
+        return results.resultLinks or [], resp.hasMoreRows
+
+    def fetch_results(
+        self,
+        command_id: CommandId,
+        max_rows: int,
+        max_bytes: int,
+        expected_row_start_offset: int,
+        lz4_compressed: bool,
+        arrow_schema_bytes,
+        description,
+        chunk_id: int,
+        use_cloud_fetch=True,
+    ):
+        resp = self._fetch_next(
+            command_id, max_rows, max_bytes, expected_row_start_offset, use_cloud_fetch
+        )
 
         queue = ThriftResultSetQueueFactory.build_queue(
             row_set_type=resp.resultSetMetadata.resultFormat,

src/databricks/sql/cloudfetch/download_manager.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import logging
+import threading
 
 from concurrent.futures import ThreadPoolExecutor, Future
-from typing import List, Union, Tuple, Optional, TYPE_CHECKING
+from typing import Callable, List, Union, Tuple, Optional, TYPE_CHECKING
 
 from databricks.sql.cloudfetch.downloader import (
     ResultSetDownloadHandler,
@@ -21,6 +22,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Called with the number of free download slots; returns (new links, has more rows).
+LinkSource = Callable[[int], Tuple[List["TSparkArrowResultLink"], bool]]
+
 
 class ResultFileDownloadManager:
     def __init__(
@@ -56,6 +60,21 @@ def __init__(
         self.session_id_hex = session_id_hex
         self.statement_id = statement_id
         self._http_client = http_client
+        self._link_source: Optional[LinkSource] = None
+        # Set on shutdown so downloads that have not started (or finished) stop early.
+        self._cancelled = threading.Event()
+
+    def set_link_source(self, link_source: LinkSource) -> bool:
+        """Fetch further links on demand.
+
+        Whenever a download slot is free and no links are pending, ``link_source`` is
+        asked for more, so the next links arrive while earlier files still download.
+        Returns False, without attaching it, if the manager is already shut down.
+        """
+        if self._cancelled.is_set():
+            return False
+        self._link_source = link_source
+        return True
 
     def get_next_downloaded_file(
         self, next_row_offset: int
@@ -98,6 +117,28 @@ def get_next_downloaded_file(
         return file
 
     def _schedule_downloads(self):
+        self._submit_pending_links()
+        while (
+            self._link_source is not None
+            and not self._pending_links
+            and len(self._download_tasks) < self._max_download_threads
+        ):
+            scheduled = len(self._download_tasks)
+            links, has_more = self._link_source(self._max_download_threads - scheduled)
+            if self._cancelled.is_set():
+                # Shut down while the request was in flight
+                return
+            if not has_more:
+                self._link_source = None
+            for link in links:
+                self.add_link(link)
+            self._submit_pending_links()
+            if self._download_tasks and len(self._download_tasks) == scheduled:
+                # Nothing new to download (no links, or only empty ones): ask again
+                # on the next call instead of polling the server now
+                break
+
+    def _submit_pending_links(self):
         """
         While download queue has a capacity, peek pending links and submit them to thread pool.
         """
@@ -119,6 +160,7 @@ def _schedule_downloads(self):
                 session_id_hex=self.session_id_hex,
                 statement_id=self.statement_id,
                 http_client=self._http_client,
+                cancel_event=self._cancelled,
             )
             task = self._thread_pool.submit(handler.run)
             self._download_tasks.append(task)
@@ -144,6 +186,8 @@ def add_link(self, link: TSparkArrowResultLink):
 
     def _shutdown_manager(self):
         # Clear download handlers and shutdown the thread pool
+        self._link_source = None
         self._pending_links = []
         self._download_tasks = []
-        self._thread_pool.shutdown(wait=False)
+        self._cancelled.set()
+        self._thread_pool.shutdown(wait=False, cancel_futures=True)

src/databricks/sql/cloudfetch/downloader.py

@@ -19,6 +19,8 @@
     # runtime -- which keeps the (Apache Thrift) ``thrift`` package out of the
     # cloud-fetch code path used by the SEA and kernel backends. See the
     # ``test_lazy_thrift_import`` regression test.
+    import threading
+
     from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 
 logger = logging.getLogger(__name__)
@@ -70,6 +72,7 @@ def __init__(
         session_id_hex: Optional[str],
         statement_id: str,
         http_client,
+        cancel_event: Optional[threading.Event] = None,
     ):
         self.settings = settings
         self.link = link
@@ -78,6 +81,11 @@ def __init__(
         self.chunk_id = chunk_id
         self.session_id_hex = session_id_hex
         self.statement_id = statement_id
+        self._cancel_event = cancel_event
+
+    def _raise_if_cancelled(self):
+        if self._cancel_event is not None and self._cancel_event.is_set():
+            raise Error("CloudFetch download cancelled: the result set was closed")
 
     @log_latency(StatementType.QUERY)
     def run(self) -> DownloadedFile:
@@ -100,6 +108,7 @@ def run(self) -> DownloadedFile:
             self.link, self.settings.link_expiry_buffer_secs
         )
 
+        self._raise_if_cancelled()
         start_time = time.time()
 
         with self._http_client.request_context(
@@ -111,6 +120,8 @@ def run(self) -> DownloadedFile:
             if response.status >= 400:
                 raise Exception(f"HTTP {response.status}: {response.data.decode()}")
             compressed_data = response.data
+        # Skip decompressing a file nobody will read
+        self._raise_if_cancelled()
 
         # Log download metrics
         download_duration = time.time() - start_time

src/databricks/sql/result_set.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import List, Optional, TYPE_CHECKING, Tuple
+from typing import List, Optional, TYPE_CHECKING, Tuple, cast
 
 import logging
 import pandas
@@ -214,6 +214,7 @@ def __init__(
         max_download_threads: int = 10,
         ssl_options=None,
         has_more_rows: bool = True,
+        cloudfetch_prefetch: bool = True,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -230,6 +231,8 @@ def __init__(
             :param max_download_threads: Maximum number of download threads for cloud fetch
             :param ssl_options: SSL options for cloud fetch
             :param has_more_rows: Whether there are more rows to fetch
+            :param cloudfetch_prefetch: Request further CloudFetch links while earlier
+                files still download, instead of one batch at a time
         """
         self.num_chunks = 0
 
@@ -238,6 +241,8 @@ def __init__(
         self.has_more_rows = has_more_rows
         self._max_download_threads = max_download_threads
         self._max_link_bytes = 0
+        self._cloudfetch_prefetch = cloudfetch_prefetch
+        self._links_end_offset = 0
 
         # Build the results queue if t_row_set is provided
         results_queue = None
@@ -283,6 +288,51 @@ def __init__(
         # Initialize results queue if not provided
         if not self.results:
             self._fill_results_buffer()
+        else:
+            self._start_prefetch()
+
+    def _start_prefetch(self) -> None:
+        # Hand the remaining CloudFetch links to the queue: its download manager then
+        # asks for more links whenever a download slot frees up, rather than this result
+        # set fetching the next batch only after the current one is fully consumed.
+        queue = self.results
+        if (
+            not self._cloudfetch_prefetch
+            or not isinstance(queue, ThriftCloudFetchQueue)
+            or not queue.result_links
+            or not self.has_more_rows
+            or self.has_been_closed_server_side
+        ):
+            return
+        self._links_end_offset = max(
+            link.startRowOffset + link.rowCount for link in queue.result_links
+        )
+        if queue.download_manager.set_link_source(self._fetch_more_links):
+            self.has_more_rows = False
+
+    def _fetch_more_links(self, free_slots: int):
+        # One link per free slot plus half a pool queued ahead, so slots refill without a
+        # round trip each time. A queued link waits for at most that many files to be
+        # consumed before it starts downloading.
+        wanted = (free_slots + self._max_download_threads // 2) * self._max_link_bytes
+        max_bytes = min(
+            max(self.buffer_size_bytes, wanted),
+            _CLOUDFETCH_MAX_BUFFERS_PER_FETCH * self.buffer_size_bytes,
+        )
+        backend = cast("ThriftDatabricksClient", self.backend)
+        links, has_more_rows = backend.fetch_result_links(
+            command_id=self.command_id,
+            max_rows=self.arraysize,
+            max_bytes=max_bytes,
+            expected_row_start_offset=self._links_end_offset,
+        )
+        if links:
+            self._note_link_sizes(links)
+            self.num_chunks += len(links)
+            self._links_end_offset = max(
+                link.startRowOffset + link.rowCount for link in links
+            )
+        return links, has_more_rows
 
     def _note_link_sizes(self, links) -> None:
         for link in links:
@@ -322,6 +372,7 @@ def _fill_results_buffer(self):
         self.num_chunks += result_links_count
         if isinstance(results, ThriftCloudFetchQueue):
             self._note_link_sizes(results.result_links)
+        self._start_prefetch()
 
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]