From 9c6cb9eab11aae17d6f34d85b6e98390324a104e Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 10:39:12 +0000 Subject: [PATCH 1/5] perf(thrift): prefetch CloudFetch links while earlier files download Each FetchResults batch waited for its slowest file before the next batch was requested, so on average only about half of the download threads were busy. The download manager now asks the result set for more links whenever a slot frees up and no links are pending (one per free slot plus half a pool), using one thread pool for the whole result. _cloudfetch_prefetch=False restores per-batch fetching. Closing a result set now sets a cancel event that downloads check before the GET and before decompression, and cancels queued downloads. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/backend/thrift_backend.py | 52 ++- .../sql/cloudfetch/download_manager.py | 41 ++- src/databricks/sql/cloudfetch/downloader.py | 11 + src/databricks/sql/result_set.py | 53 ++- tests/unit/test_cloudfetch_prefetch.py | 334 ++++++++++++++++++ tests/unit/test_fetches.py | 3 +- 7 files changed, 483 insertions(+), 12 deletions(-) create mode 100644 tests/unit/test_cloudfetch_prefetch.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4038e882d..d6240c029 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index ee904c368..05f929f77 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/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,7 @@ def __init__( # Cloud fetch self._max_download_threads = kwargs.get("max_download_threads", 10) + self._cloudfetch_prefetch = bool(kwargs.get("_cloudfetch_prefetch", True)) self._ssl_options = ssl_options self._auth_provider = auth_provider @@ -1001,6 +1005,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 +1160,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 +1199,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 +1246,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 +1297,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 +1348,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 +1384,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,7 +1410,7 @@ 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( @@ -1412,6 +1418,36 @@ def fetch_results( ), 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 + ) + return resp.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, diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index ece7406b7..8b00731e1 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/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,17 @@ 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) -> None: + """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. + """ + self._link_source = link_source def get_next_downloaded_file( self, next_row_offset: int @@ -98,6 +113,25 @@ 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 + ): + links, has_more = self._link_source( + self._max_download_threads - len(self._download_tasks) + ) + if not has_more: + self._link_source = None + for link in links: + self.add_link(link) + self._submit_pending_links() + if not links and self._download_tasks: + # 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 +153,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 +179,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) diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index 295f147cd..9a140ea63 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/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 diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index 2a2599259..2212aa6b2 100644 --- a/src/databricks/sql/result_set.py +++ b/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 + ) + 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] diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py new file mode 100644 index 000000000..ff113ecbd --- /dev/null +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -0,0 +1,334 @@ +"""CloudFetch link prefetching across batches and cancel-on-close.""" + +import threading +import unittest +from unittest.mock import MagicMock, Mock, patch + +import databricks.sql.cloudfetch.download_manager as download_manager +import databricks.sql.cloudfetch.downloader as downloader +from databricks.sql.auth.authenticators import AuthProvider +from databricks.sql.backend.thrift_backend import ThriftDatabricksClient +from databricks.sql.backend.types import ExecuteResponse +from databricks.sql.cloudfetch.downloader import DownloadedFile +from databricks.sql.exc import DataError, Error +from databricks.sql.result_set import ThriftResultSet +from databricks.sql.thrift_api.TCLIService import ttypes +from databricks.sql.types import SSLOptions +from databricks.sql.utils import ArrowQueue, ThriftCloudFetchQueue + +MIB = 2**20 + + +def _link(start, rows=10, size=20 * MIB): + return ttypes.TSparkArrowResultLink( + fileLink=f"https://example.com/{start}", + expiryTime=2**62, + startRowOffset=start, + rowCount=rows, + bytesNum=size, + ) + + +class _InstantHandler: + """Stands in for ResultSetDownloadHandler: "downloads" a link immediately.""" + + def __init__(self, link, **_): + self.link = link + + def run(self): + return DownloadedFile(b"", self.link.startRowOffset, self.link.rowCount) + + +class DownloadManagerPrefetchTests(unittest.TestCase): + def _manager(self, links, max_download_threads): + return download_manager.ResultFileDownloadManager( + links, + max_download_threads, + lz4_compressed=False, + ssl_options=SSLOptions(), + session_id_hex=None, + statement_id="s", + chunk_id=0, + http_client=MagicMock(), + ) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_refills_from_link_source_whenever_a_slot_frees(self): + manager = self._manager([_link(0), _link(10)], max_download_threads=4) + replies = [ + ([_link(20), _link(30), _link(40)], True), + ([_link(50)], True), + ([], False), + ] + free_slots = [] + + def source(free): + free_slots.append(free) + self.assertLessEqual(len(manager._download_tasks), 4) + return replies.pop(0) + + manager.set_link_source(source) + offsets = [] + while True: + file = manager.get_next_downloaded_file(len(offsets) * 10) + if file is None: + break + self.assertLessEqual(len(manager._download_tasks), 4) + offsets.append(file.start_row_offset) + + self.assertEqual(offsets, [0, 10, 20, 30, 40, 50]) + # Asked only when a slot was free and no link was pending + self.assertEqual(free_slots, [2, 1, 1]) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_empty_reply_is_retried_on_the_next_call_not_in_a_loop(self): + manager = self._manager([_link(0)], max_download_threads=2) + source = Mock(side_effect=[([], True), ([_link(10)], False)]) + manager.set_link_source(source) + + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + self.assertEqual(source.call_count, 1) + self.assertEqual(manager.get_next_downloaded_file(10).start_row_offset, 10) + self.assertIsNone(manager.get_next_downloaded_file(20)) + self.assertEqual(source.call_count, 2) + + def test_shutdown_cancels_downloads_and_stops_fetching_links(self): + manager = self._manager([_link(0), _link(10)], max_download_threads=1) + manager.set_link_source(Mock()) + with patch.object( + manager._thread_pool, "shutdown", wraps=manager._thread_pool.shutdown + ) as shutdown: + manager._shutdown_manager() + + self.assertTrue(manager._cancelled.is_set()) + self.assertIsNone(manager._link_source) + self.assertEqual(manager._pending_links, []) + shutdown.assert_called_once_with(wait=False, cancel_futures=True) + + def test_handlers_share_the_managers_cancel_event(self): + manager = self._manager([_link(0)], max_download_threads=1) + with patch.object(download_manager, "ResultSetDownloadHandler") as handler: + manager._schedule_downloads() + self.assertIs(handler.call_args.kwargs["cancel_event"], manager._cancelled) + manager._shutdown_manager() + + +class DownloaderCancelTests(unittest.TestCase): + def _handler(self, http_client, event): + return downloader.ResultSetDownloadHandler( + downloader.DownloadableResultSettings(is_lz4_compressed=True), + _link(0), + ssl_options=SSLOptions(), + chunk_id=0, + session_id_hex=None, + statement_id="s", + http_client=http_client, + cancel_event=event, + ) + + def test_cancelled_before_the_request_skips_the_get(self): + http_client = MagicMock() + event = threading.Event() + event.set() + + with self.assertRaises(Error): + self._handler(http_client, event).run() + http_client.request_context.assert_not_called() + + def test_cancelled_during_the_download_skips_decompression(self): + event = threading.Event() + + class Response: + status = 200 + + @property + def data(self): + event.set() + return b"compressed" + + http_client = MagicMock() + http_client.request_context.return_value.__enter__.return_value = Response() + with patch.object( + downloader.ResultSetDownloadHandler, "_decompress_data" + ) as decompress: + with self.assertRaises(Error): + self._handler(http_client, event).run() + decompress.assert_not_called() + + +class ResultSetPrefetchTests(unittest.TestCase): + BUFFER = 100 * MIB + + def _result_set(self, initial_links, prefetch=True, has_more_rows=True, queue=None): + backend = Mock(spec=ThriftDatabricksClient) + if queue is None: + queue = Mock(spec=ThriftCloudFetchQueue) + queue.result_links = initial_links + queue.download_manager = Mock() + t_row_set = ttypes.TRowSet(startRowOffset=0, rows=[], resultLinks=initial_links) + with patch( + "databricks.sql.utils.ThriftResultSetQueueFactory.build_queue", + return_value=queue, + ): + 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=self.BUFFER, + arraysize=1000, + t_row_set=t_row_set, + max_download_threads=10, + has_more_rows=has_more_rows, + cloudfetch_prefetch=prefetch, + ) + return rs, backend, queue + + def test_hands_link_fetching_to_the_queue(self): + rs, backend, queue = self._result_set([_link(0, 100), _link(100, 100)]) + + queue.download_manager.set_link_source.assert_called_once_with( + rs._fetch_more_links + ) + self.assertFalse(rs.has_more_rows) + backend.fetch_results.assert_not_called() + + def test_link_source_asks_for_free_slots_plus_half_a_pool(self): + rs, backend, queue = self._result_set([_link(0, 100), _link(100, 100)]) + source = queue.download_manager.set_link_source.call_args.args[0] + backend.fetch_result_links.side_effect = [ + ([_link(200, 50)], True), + ([_link(250, 50)], False), + ] + + self.assertEqual(source(3), ([_link(200, 50)], True)) + self.assertEqual(source(1), ([_link(250, 50)], False)) + + first, second = backend.fetch_result_links.call_args_list + self.assertEqual(first.kwargs["max_bytes"], 8 * 20 * MIB) + self.assertEqual(first.kwargs["expected_row_start_offset"], 200) + self.assertEqual(first.kwargs["max_rows"], 1000) + self.assertEqual(second.kwargs["max_bytes"], 6 * 20 * MIB) + self.assertEqual(second.kwargs["expected_row_start_offset"], 250) + + def test_link_source_budget_stays_between_one_and_four_buffers(self): + rs, backend, queue = self._result_set([_link(0, size=60 * MIB)]) + source = queue.download_manager.set_link_source.call_args.args[0] + backend.fetch_result_links.return_value = ([], True) + + source(10) + rs._max_link_bytes = MIB + source(1) + + budgets = [ + c.kwargs["max_bytes"] for c in backend.fetch_result_links.call_args_list + ] + self.assertEqual(budgets, [4 * self.BUFFER, self.BUFFER]) + + def test_kill_switch_keeps_per_batch_fetching(self): + rs, backend, queue = self._result_set([_link(0)], prefetch=False) + + queue.download_manager.set_link_source.assert_not_called() + self.assertTrue(rs.has_more_rows) + + def test_inline_results_are_not_prefetched(self): + rs, backend, _ = self._result_set([], queue=Mock(spec=ArrowQueue)) + + self.assertTrue(rs.has_more_rows) + backend.fetch_result_links.assert_not_called() + + def test_complete_results_are_not_prefetched(self): + rs, _, queue = self._result_set([_link(0)], has_more_rows=False) + + queue.download_manager.set_link_source.assert_not_called() + + def test_prefetch_starts_after_the_first_fetch_without_direct_results(self): + backend = Mock(spec=ThriftDatabricksClient) + queue = Mock(spec=ThriftCloudFetchQueue) + queue.result_links = [_link(0)] + queue.download_manager = Mock() + backend.fetch_results.return_value = (queue, True, 1) + + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=Mock(), + status=None, + has_been_closed_server_side=False, + description=[], + ), + thrift_client=backend, + ) + + queue.download_manager.set_link_source.assert_called_once() + self.assertFalse(rs.has_more_rows) + + +class ThriftClientPrefetchTests(unittest.TestCase): + def _client(self, **kwargs): + return ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + **kwargs, + ) + + def test_prefetch_is_on_by_default_and_can_be_disabled(self): + self.assertTrue(self._client()._cloudfetch_prefetch) + self.assertFalse(self._client(_cloudfetch_prefetch=False)._cloudfetch_prefetch) + + @patch("databricks.sql.backend.thrift_backend.ThriftResultSet") + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) + def test_execute_passes_the_setting_to_the_result_set(self, _, result_set): + client = self._client(_cloudfetch_prefetch=False) + client._handle_execute_response = Mock(return_value=(Mock(), Mock())) + + client.execute_command("foo", Mock(), 100, 200, Mock(), Mock(), Mock()) + + self.assertIs(result_set.call_args.kwargs["cloudfetch_prefetch"], False) + + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) + def test_fetch_result_links(self, tcli_service_class): + links = [_link(200, 50)] + tcli_service_class.return_value.FetchResults.return_value = ( + ttypes.TFetchResultsResp( + status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), + hasMoreRows=True, + results=ttypes.TRowSet(startRowOffset=200, rows=[], resultLinks=links), + ) + ) + client = self._client() + + got = client.fetch_result_links( + command_id=Mock(), + max_rows=1000, + max_bytes=123, + expected_row_start_offset=200, + ) + + self.assertEqual(got, (links, True)) + req = tcli_service_class.return_value.FetchResults.call_args.args[0] + self.assertEqual( + (req.maxRows, req.maxBytes, req.orientation), + (1000, 123, ttypes.TFetchOrientation.FETCH_NEXT), + ) + with self.assertRaises(DataError): + client.fetch_result_links( + command_id=Mock(), + max_rows=1, + max_bytes=1, + expected_row_start_offset=100, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py index d0d74df1d..8bffc8183 100644 --- a/tests/unit/test_fetches.py +++ b/tests/unit/test_fetches.py @@ -393,7 +393,7 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self): class CloudFetchBatchSizeTests(unittest.TestCase): - """maxBytes sent by ThriftResultSet._fill_results_buffer.""" + """maxBytes sent by ThriftResultSet._fill_results_buffer (prefetch disabled).""" MIB = 2**20 BUFFER = 100 * MIB @@ -457,6 +457,7 @@ def fetch_results(**kwargs): arraysize=1000, t_row_set=t_row_set, max_download_threads=max_download_threads, + cloudfetch_prefetch=False, ) return rs, calls From 117a095c63c568caf0b24899bd51bc510fa845f4 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:18:00 +0000 Subject: [PATCH 2/5] Guard prefetched link fetches and a shut-down queue fetch_result_links raises DataError if a reply carries rows that are not CloudFetch links, instead of dropping them. set_link_source refuses a manager that is already shut down, so such a result keeps per-batch fetching. Pin that a failed link fetch reaches the caller and that fetching can resume. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 17 ++++- .../sql/cloudfetch/download_manager.py | 6 +- src/databricks/sql/result_set.py | 4 +- tests/unit/test_cloudfetch_prefetch.py | 69 +++++++++++++++++++ 4 files changed, 92 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 05f929f77..5a42058cd 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -1431,7 +1431,22 @@ def fetch_result_links( resp = self._fetch_next( command_id, max_rows, max_bytes, expected_row_start_offset, True ) - return resp.results.resultLinks or [], resp.hasMoreRows + 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, diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 8b00731e1..f40dee9a1 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -64,13 +64,17 @@ def __init__( # 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) -> None: + 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 diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index 2212aa6b2..73c3c0811 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -307,8 +307,8 @@ def _start_prefetch(self) -> None: self._links_end_offset = max( link.startRowOffset + link.rowCount for link in queue.result_links ) - queue.download_manager.set_link_source(self._fetch_more_links) - self.has_more_rows = False + 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 diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index ff113ecbd..c176946d2 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -105,6 +105,32 @@ def test_shutdown_cancels_downloads_and_stops_fetching_links(self): self.assertEqual(manager._pending_links, []) shutdown.assert_called_once_with(wait=False, cancel_futures=True) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): + manager = self._manager([_link(0)], max_download_threads=2) + source = Mock( + side_effect=[RuntimeError("FetchResults failed"), ([_link(10)], False)] + ) + manager.set_link_source(source) + + with self.assertRaises(RuntimeError): + manager.get_next_downloaded_file(0) + # The scheduled download is kept and the next call asks for links again + offsets = [manager.get_next_downloaded_file(0).start_row_offset] + offsets.append(manager.get_next_downloaded_file(10).start_row_offset) + self.assertIsNone(manager.get_next_downloaded_file(20)) + self.assertEqual(offsets, [0, 10]) + self.assertEqual(source.call_count, 2) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_link_source_is_refused_after_shutdown(self): + # No schedulable links: the first request already shuts the manager down + manager = self._manager([_link(0, rows=0)], max_download_threads=2) + self.assertIsNone(manager.get_next_downloaded_file(0)) + + self.assertFalse(manager.set_link_source(Mock())) + self.assertIsNone(manager._link_source) + def test_handlers_share_the_managers_cancel_event(self): manager = self._manager([_link(0)], max_download_threads=1) with patch.object(download_manager, "ResultSetDownloadHandler") as handler: @@ -236,6 +262,16 @@ def test_kill_switch_keeps_per_batch_fetching(self): queue.download_manager.set_link_source.assert_not_called() self.assertTrue(rs.has_more_rows) + def test_shut_down_queue_keeps_per_batch_fetching(self): + queue = Mock(spec=ThriftCloudFetchQueue) + queue.result_links = [_link(0)] + queue.download_manager = Mock() + queue.download_manager.set_link_source.return_value = False + + rs, _, _ = self._result_set([_link(0)], queue=queue) + + self.assertTrue(rs.has_more_rows) + def test_inline_results_are_not_prefetched(self): rs, backend, _ = self._result_set([], queue=Mock(spec=ArrowQueue)) @@ -329,6 +365,39 @@ def test_fetch_result_links(self, tcli_service_class): expected_row_start_offset=100, ) + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) + def test_fetch_result_links_rejects_rows_that_are_not_links( + self, tcli_service_class + ): + ok = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) + inline_rows = ttypes.TRowSet( + startRowOffset=0, + rows=[], + arrowBatches=[ttypes.TSparkArrowBatch(batch=b"x", rowCount=1)], + ) + arrow_format = ttypes.TGetResultSetMetadataResp( + status=ok, resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET + ) + client = self._client() + + for resp in ( + ttypes.TFetchResultsResp(status=ok, hasMoreRows=False, results=inline_rows), + ttypes.TFetchResultsResp( + status=ok, + hasMoreRows=False, + results=ttypes.TRowSet(startRowOffset=0, rows=[]), + resultSetMetadata=arrow_format, + ), + ): + tcli_service_class.return_value.FetchResults.return_value = resp + with self.assertRaises(DataError): + client.fetch_result_links( + command_id=Mock(), + max_rows=1, + max_bytes=1, + expected_row_start_offset=0, + ) + if __name__ == "__main__": unittest.main() From 0dede5d210960646b2e3e253b032a9c1292a8580 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:24:43 +0000 Subject: [PATCH 3/5] Don't re-poll for links when a reply only had empty links Base the early exit on whether a download was scheduled, so a reply whose links all have zero rows waits for the next call like an empty reply. Add end-to-end tests that run ThriftResultSet over a real CloudFetch queue and download manager across several link fetches. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/cloudfetch/download_manager.py | 10 +- tests/unit/test_cloudfetch_prefetch.py | 109 ++++++++++++++++++ 2 files changed, 114 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index f40dee9a1..1696ef302 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -123,16 +123,16 @@ def _schedule_downloads(self): and not self._pending_links and len(self._download_tasks) < self._max_download_threads ): - links, has_more = self._link_source( - self._max_download_threads - len(self._download_tasks) - ) + scheduled = len(self._download_tasks) + links, has_more = self._link_source(self._max_download_threads - scheduled) if not has_more: self._link_source = None for link in links: self.add_link(link) self._submit_pending_links() - if not links and self._download_tasks: - # Ask again on the next call instead of polling the server now + 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): diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index c176946d2..48e481db2 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -4,6 +4,13 @@ import unittest from unittest.mock import MagicMock, Mock, patch +import pytest + +try: + import pyarrow as pa +except ImportError: + pa = None + import databricks.sql.cloudfetch.download_manager as download_manager import databricks.sql.cloudfetch.downloader as downloader from databricks.sql.auth.authenticators import AuthProvider @@ -39,6 +46,18 @@ def run(self): return DownloadedFile(b"", self.link.startRowOffset, self.link.rowCount) +class _ArrowFileHandler(_InstantHandler): + """Returns an Arrow IPC stream with ids startRowOffset..startRowOffset+rowCount-1.""" + + def run(self): + start, rows = self.link.startRowOffset, self.link.rowCount + table = pa.table({"id": pa.array(range(start, start + rows), pa.int64())}) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + return DownloadedFile(sink.getvalue().to_pybytes(), start, rows) + + class DownloadManagerPrefetchTests(unittest.TestCase): def _manager(self, links, max_download_threads): return download_manager.ResultFileDownloadManager( @@ -105,6 +124,17 @@ def test_shutdown_cancels_downloads_and_stops_fetching_links(self): self.assertEqual(manager._pending_links, []) shutdown.assert_called_once_with(wait=False, cancel_futures=True) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_reply_with_only_empty_links_is_retried_on_the_next_call(self): + manager = self._manager([_link(0)], max_download_threads=2) + source = Mock(side_effect=[([_link(10, rows=0)], True), ([_link(10)], False)]) + manager.set_link_source(source) + + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + self.assertEqual(source.call_count, 1) + self.assertEqual(manager.get_next_downloaded_file(10).start_row_offset, 10) + self.assertIsNone(manager.get_next_downloaded_file(20)) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): manager = self._manager([_link(0)], max_download_threads=2) @@ -305,6 +335,85 @@ def test_prefetch_starts_after_the_first_fetch_without_direct_results(self): self.assertFalse(rs.has_more_rows) +@pytest.mark.skipif(pa is None, reason="PyArrow is not installed") +@patch.object(download_manager, "ResultSetDownloadHandler", _ArrowFileHandler) +class ResultSetPrefetchEndToEndTests(unittest.TestCase): + """Real ThriftCloudFetchQueue and download manager; only the download is stubbed.""" + + ROWS_PER_LINK = 10 + + def _result_set(self, initial_links=2, total_links=11, links_per_reply=3): + links = [ + _link(i * self.ROWS_PER_LINK, self.ROWS_PER_LINK) + for i in range(total_links) + ] + rest = links[initial_links:] + expected_offsets = [] + + def fetch_result_links(**kwargs): + expected_offsets.append(kwargs["expected_row_start_offset"]) + reply = rest[:links_per_reply] + del rest[:links_per_reply] + return reply, bool(rest) + + backend = Mock(spec=ThriftDatabricksClient) + backend.fetch_result_links.side_effect = fetch_result_links + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=Mock(), + status=None, + has_been_closed_server_side=False, + description=[("id", "bigint", None, None, None, None, None)], + lz4_compressed=False, + arrow_schema_bytes=pa.schema([("id", pa.int64())]) + .serialize() + .to_pybytes(), + result_format=ttypes.TSparkRowSetType.URL_BASED_SET, + ), + thrift_client=backend, + arraysize=1000, + t_row_set=ttypes.TRowSet( + startRowOffset=0, rows=[], resultLinks=links[:initial_links] + ), + max_download_threads=4, + ) + return rs, backend, expected_offsets + + def test_fetchall_arrow_reads_every_link_in_order(self): + rs, backend, expected_offsets = self._result_set() + + self.assertFalse(rs.has_more_rows) + ids = rs.fetchall_arrow().column(0).to_pylist() + + self.assertEqual(ids, list(range(110))) + self.assertEqual(expected_offsets, [20, 50, 80]) + backend.fetch_results.assert_not_called() + + def test_fetchmany_and_fetchone_cross_link_batches(self): + rs, _, _ = self._result_set() + + ids = [rs.fetchone()[0] for _ in range(15)] + while True: + table = rs.fetchmany_arrow(25) + if table.num_rows == 0: + break + ids += table.column(0).to_pylist() + + self.assertEqual(ids, list(range(110))) + + def test_close_mid_stream_cancels_the_download_manager(self): + rs, backend, _ = self._result_set() + manager = rs.results.download_manager + rs.fetchmany_arrow(35) + + rs.close() + + self.assertTrue(manager._cancelled.is_set()) + self.assertIsNone(manager._link_source) + backend.close_command.assert_called_once() + + class ThriftClientPrefetchTests(unittest.TestCase): def _client(self, **kwargs): return ThriftDatabricksClient( From e114c4830322a6ba518ea50d03c6010cdc4d2c84 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:29:21 +0000 Subject: [PATCH 4/5] Stop scheduling if the manager is shut down during a link fetch A close() on another thread while FetchResults is in flight now ends the stream instead of submitting to a shut-down pool. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/cloudfetch/download_manager.py | 3 +++ tests/unit/test_cloudfetch_prefetch.py | 13 +++++++++++++ 2 files changed, 16 insertions(+) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 1696ef302..0ca491a64 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -125,6 +125,9 @@ def _schedule_downloads(self): ): 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: diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 48e481db2..7ace54cf0 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -152,6 +152,19 @@ def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): self.assertEqual(offsets, [0, 10]) self.assertEqual(source.call_count, 2) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_shutdown_during_a_link_fetch_ends_the_stream_cleanly(self): + manager = self._manager([_link(0)], max_download_threads=2) + + def source(free): + manager._shutdown_manager() # e.g. close() on another thread + return [_link(10)], True + + manager.set_link_source(source) + + self.assertIsNone(manager.get_next_downloaded_file(0)) + self.assertEqual(manager._pending_links, []) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) def test_link_source_is_refused_after_shutdown(self): # No schedulable links: the first request already shuts the manager down From 5ef9b34cf889f92954b997f2056176f23c5fb7ec Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:35:13 +0000 Subject: [PATCH 5/5] Accept string values for _cloudfetch_prefetch "false", "0", "no" and "off" now turn prefetch off, so the kill switch also works when set from a URL query. Add a queue-level test that a failed link fetch propagates out of the fetch call. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 5 ++++- tests/unit/test_cloudfetch_prefetch.py | 18 +++++++++++++++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 5a42058cd..e03bce69a 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -195,7 +195,10 @@ def __init__( # Cloud fetch self._max_download_threads = kwargs.get("max_download_threads", 10) - self._cloudfetch_prefetch = bool(kwargs.get("_cloudfetch_prefetch", True)) + 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 diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 7ace54cf0..260f4d1a5 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -136,7 +136,7 @@ def test_reply_with_only_empty_links_is_retried_on_the_next_call(self): self.assertIsNone(manager.get_next_downloaded_file(20)) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) - def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): + def test_link_source_errors_reach_the_caller_and_the_next_call_asks_again(self): manager = self._manager([_link(0)], max_download_threads=2) source = Mock( side_effect=[RuntimeError("FetchResults failed"), ([_link(10)], False)] @@ -157,7 +157,7 @@ def test_shutdown_during_a_link_fetch_ends_the_stream_cleanly(self): manager = self._manager([_link(0)], max_download_threads=2) def source(free): - manager._shutdown_manager() # e.g. close() on another thread + manager._shutdown_manager() return [_link(10)], True manager.set_link_source(source) @@ -426,6 +426,13 @@ def test_close_mid_stream_cancels_the_download_manager(self): self.assertIsNone(manager._link_source) backend.close_command.assert_called_once() + def test_link_fetch_errors_propagate_out_of_the_fetch_call(self): + rs, backend, _ = self._result_set() + backend.fetch_result_links.side_effect = DataError("FetchResults failed") + + with self.assertRaises(DataError): + rs.fetchall_arrow() + class ThriftClientPrefetchTests(unittest.TestCase): def _client(self, **kwargs): @@ -442,7 +449,12 @@ def _client(self, **kwargs): def test_prefetch_is_on_by_default_and_can_be_disabled(self): self.assertTrue(self._client()._cloudfetch_prefetch) - self.assertFalse(self._client(_cloudfetch_prefetch=False)._cloudfetch_prefetch) + for off in (False, "false", "False", "0", "no", "off"): + self.assertFalse( + self._client(_cloudfetch_prefetch=off)._cloudfetch_prefetch + ) + for on in (True, "true", "1"): + self.assertTrue(self._client(_cloudfetch_prefetch=on)._cloudfetch_prefetch) @patch("databricks.sql.backend.thrift_backend.ThriftResultSet") @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)