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..e03bce69a 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,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,7 +1413,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 +1421,51 @@ 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 + ) + 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, diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index ece7406b7..0ca491a64 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,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) 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..73c3c0811 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 + ) + 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] diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py new file mode 100644 index 000000000..260f4d1a5 --- /dev/null +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -0,0 +1,537 @@ +"""CloudFetch link prefetching across batches and cancel-on-close.""" + +import threading +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 +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 _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( + 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) + + @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_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)] + ) + 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_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() + 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 + 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: + 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_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)) + + 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) + + +@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() + + 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): + 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) + 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) + 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, + ) + + @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() 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