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 in flight when a result set is closed early or its statement is cancelled.
 
 # 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,24 @@ def __init__(
 
         # Cloud fetch
         self._max_download_threads = kwargs.get("max_download_threads", 10)
+        prefetch = kwargs.get("_cloudfetch_prefetch")
+        if prefetch is None:  # unset, or forwarded as None by a wrapper
+            prefetch = True
+        elif isinstance(prefetch, str):  # e.g. from a SQLAlchemy URL query
+            value = prefetch.strip().lower()
+            if value in ("false", "0", "no", "off"):
+                prefetch = False
+            elif value in ("true", "1", "yes", "on"):
+                prefetch = True
+            else:
+                # Someone setting the kill switch likely wants it off: fail safe
+                logger.warning(
+                    "Unrecognized _cloudfetch_prefetch=%r; turning CloudFetch "
+                    "prefetch off (use 'true' or 'false')",
+                    prefetch,
+                )
+                prefetch = False
+        self._cloudfetch_prefetch = bool(prefetch)
 
         self._ssl_options = ssl_options
         self._auth_provider = auth_provider
@@ -1001,6 +1022,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 +1177,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 +1216,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 +1263,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 +1314,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 +1365,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 +1401,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 +1427,60 @@ 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 results.rows
+            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/client.py

@@ -1958,7 +1958,17 @@ def cancel(self) -> None:
         This method can be called from another thread.
         """
         if self.active_command_id is not None:
-            self.backend.cancel_command(self.active_command_id)
+            try:
+                self.backend.cancel_command(self.active_command_id)
+            finally:
+                result_set = self.active_result_set
+                if result_set is not None:
+                    # Stop downloads already under way; reads then fail like the server
+                    # would. Done even if the cancel RPC fails, without hiding its error.
+                    try:
+                        result_set._cancel_downloads()
+                    except Exception as e:
+                        logger.debug("Failed to cancel CloudFetch downloads: %s", e)
             return
         # No command id yet. A backend whose synchronous ``execute()``
         # blocks without first publishing a command id (the kernel

src/databricks/sql/cloudfetch/download_manager.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 import logging
+import threading
 
-from concurrent.futures import ThreadPoolExecutor, Future
-from typing import List, Union, Tuple, Optional, TYPE_CHECKING
+from concurrent.futures import CancelledError, ThreadPoolExecutor, Future
+from typing import Callable, List, NoReturn, Union, Tuple, Optional, TYPE_CHECKING
 
+from databricks.sql.exc import OperationalError
 from databricks.sql.cloudfetch.downloader import (
     ResultSetDownloadHandler,
     DownloadableResultSettings,
@@ -21,6 +23,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 +61,37 @@ 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()
+        # Set by cancel(): reads fail once they need data the cancel stopped.
+        self._cancel_requested = False
+        # A failed link fetch, raised once the downloaded files run out
+        self._link_error: Optional[BaseException] = None
+        self._link_error_raised = False
+
+    def cancel(self) -> None:
+        """Stop downloading because the statement was cancelled.
+
+        May be called from another thread while the consumer is reading. Files that
+        finished downloading are still returned; a read that needs anything else
+        raises instead of reporting the end of the result.
+        """
+        self._cancel_requested = True
+        self._cancelled.set()
+        self._thread_pool.shutdown(wait=False, cancel_futures=True)
+
+    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
@@ -73,19 +109,30 @@ def get_next_downloaded_file(
             next_row_offset (int): The offset of the starting row of the next file we want data from.
         """
 
-        # Make sure the download queue is always full
-        self._schedule_downloads()
+        if not self._cancel_requested and self._link_error is None:
+            # Make sure the download queue is always full
+            self._schedule_downloads()
 
         # No more files to download from this batch of links
         if len(self._download_tasks) == 0:
+            if self._link_error is not None:
+                self._raise_link_error(self._link_error)
+            if self._cancel_requested and (
+                self._pending_links or self._link_source is not None
+            ):
+                # The cancel stopped part of the result: fail rather than end early
+                raise OperationalError("CloudFetch download cancelled")
             self._shutdown_manager()
             return None
 
         task = self._download_tasks.pop(0)
         # Future's `result()` method will wait for the call to complete, and return
         # the value returned by the call. If the call throws an exception - `result()`
         # will throw the same exception
-        file = task.result()
+        try:
+            file = task.result()
+        except CancelledError:
+            raise OperationalError("CloudFetch download cancelled") from None
         if (next_row_offset < file.start_row_offset) or (
             next_row_offset > file.start_row_offset + file.row_count
         ):
@@ -97,7 +144,49 @@ def get_next_downloaded_file(
 
         return file
 
+    def _raise_link_error(self, error: BaseException) -> NoReturn:
+        # That read loses the rows it had collected, so every later read fails too
+        # and a retry can't continue past them.
+        self._shutdown_manager()
+        if not self._link_error_raised:
+            self._link_error_raised = True
+            raise error
+        raise OperationalError(
+            "Fetching more CloudFetch results failed earlier; the result can't "
+            "be read further, execute the statement again"
+        ) from error
+
     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)
+            try:
+                links, has_more = self._link_source(
+                    self._max_download_threads - scheduled
+                )
+            except Exception as e:
+                # Stop scheduling; files already downloading are still returned and
+                # the read that needs rows past them raises (_raise_link_error)
+                self._link_error = e
+                return
+            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,8 +208,17 @@ 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)
+            try:
+                task = self._thread_pool.submit(handler.run)
+            except RuntimeError:
+                # cancel() shut the pool down from another thread; keep the link so
+                # the result still counts as incomplete
+                if self._cancel_requested:
+                    self._pending_links.insert(0, (chunk_id, link))
+                    return
+                raise
             self._download_tasks.append(task)
 
     def add_link(self, link: TSparkArrowResultLink):
@@ -144,6 +242,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)