diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..4038e882d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ # Unreleased - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. +- Speed up large CloudFetch results on Thrift: each FetchResults call now asks for about one result file per download thread (at most 4x `buffer_size_bytes`) instead of 100 MiB of files, so all download threads stay busy and the result takes fewer round trips. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index f24a07505..2a2599259 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -20,6 +20,7 @@ from databricks.sql.utils import ( ColumnTable, ColumnQueue, + ThriftCloudFetchQueue, concat_table_chunks, ) from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse @@ -27,6 +28,9 @@ logger = logging.getLogger(__name__) +# Upper bound on a Thrift CloudFetch FetchResults request, in multiples of buffer_size_bytes. +_CLOUDFETCH_MAX_BUFFERS_PER_FETCH = 4 + class ResultSet(ABC): """ @@ -59,7 +63,9 @@ def __init__( :param connection: The parent connection that was used to execute this command :param backend: The specialised backend client to be invoked in the fetch phase :param arraysize: The max number of rows to fetch at a time (PEP-249) - :param buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch + :param buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch. + Thrift CloudFetch fetches may request up to 4x this (see + ThriftResultSet._fetch_max_bytes). :param command_id: The command ID :param status: The command status :param has_been_closed_server_side: Whether the command has been closed on the server @@ -216,7 +222,8 @@ def __init__( :param connection: The parent connection :param execute_response: Response from the execute command :param thrift_client: The ThriftDatabricksClient instance for direct access - :param buffer_size_bytes: Buffer size for fetching results + :param buffer_size_bytes: Buffer size for fetching results. CloudFetch + fetches request between 1x and 4x this (see _fetch_max_bytes). :param arraysize: Default number of rows to fetch :param use_cloud_fetch: Whether to use cloud fetch for retrieving results :param t_row_set: The TRowSet containing result data (if available) @@ -229,6 +236,8 @@ def __init__( # Initialize ThriftResultSet-specific attributes self._use_cloud_fetch = use_cloud_fetch self.has_more_rows = has_more_rows + self._max_download_threads = max_download_threads + self._max_link_bytes = 0 # Build the results queue if t_row_set is provided results_queue = None @@ -251,6 +260,7 @@ def __init__( ) if t_row_set.resultLinks: self.num_chunks += len(t_row_set.resultLinks) + self._note_link_sizes(t_row_set.resultLinks) # Call parent constructor with common attributes super().__init__( @@ -274,11 +284,32 @@ def __init__( if not self.results: self._fill_results_buffer() + def _note_link_sizes(self, links) -> None: + for link in links: + # The reader doesn't enforce required fields; a missing size only + # disables the sizing heuristic. + size = link.bytesNum or 0 + if size > self._max_link_bytes: + self._max_link_bytes = size + + def _fetch_max_bytes(self) -> int: + # The server packs CloudFetch links into a FetchResults reply up to maxBytes of + # file data, so buffer_size_bytes alone yields only a few links per round trip. + # Ask for about one link per download thread: every link in the batch then starts + # downloading immediately, so none sits waiting while it ages towards expiry. + # The cap bounds the batch (and buffered downloads) for a large + # max_download_threads, an oversized link, or a deliberately small buffer. + wanted = self._max_download_threads * self._max_link_bytes + return min( + max(self.buffer_size_bytes, wanted), + _CLOUDFETCH_MAX_BUFFERS_PER_FETCH * self.buffer_size_bytes, + ) + def _fill_results_buffer(self): results, has_more_rows, result_links_count = self.backend.fetch_results( command_id=self.command_id, max_rows=self.arraysize, - max_bytes=self.buffer_size_bytes, + max_bytes=self._fetch_max_bytes(), expected_row_start_offset=self._next_row_index, lz4_compressed=self.lz4_compressed, arrow_schema_bytes=self._arrow_schema_bytes, @@ -289,6 +320,8 @@ def _fill_results_buffer(self): self.results = results self.has_more_rows = has_more_rows self.num_chunks += result_links_count + if isinstance(results, ThriftCloudFetchQueue): + self._note_link_sizes(results.result_links) def _convert_columnar_table(self, table): column_names = [c[0] for c in self.description] diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py index e3a963355..d0d74df1d 100644 --- a/tests/unit/test_fetches.py +++ b/tests/unit/test_fetches.py @@ -1,6 +1,6 @@ import unittest import pytest -from unittest.mock import Mock +from unittest.mock import Mock, patch try: import pyarrow as pa @@ -9,9 +9,10 @@ import databricks.sql.client as client from databricks.sql.backend.types import ExecuteResponse -from databricks.sql.utils import ArrowQueue +from databricks.sql.utils import ArrowQueue, ThriftCloudFetchQueue from databricks.sql.backend.thrift_backend import ThriftDatabricksClient from databricks.sql.result_set import ThriftResultSet +from databricks.sql.thrift_api.TCLIService import ttypes class _StubArrowQueue: @@ -391,5 +392,168 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self): self.assertEqual(result.num_rows, 0) +class CloudFetchBatchSizeTests(unittest.TestCase): + """maxBytes sent by ThriftResultSet._fill_results_buffer.""" + + MIB = 2**20 + BUFFER = 100 * MIB + + @staticmethod + def _link(bytes_num): + return ttypes.TSparkArrowResultLink( + fileLink="https://example.com/file", + expiryTime=0, + startRowOffset=0, + rowCount=1, + bytesNum=bytes_num, + ) + + @classmethod + def _cloudfetch_queue(cls, *link_sizes): + queue = Mock(spec=ThriftCloudFetchQueue) + queue.result_links = [cls._link(b) for b in link_sizes] + return queue + + def _make_result_set( + self, + initial_link_sizes=None, + fetched_queues=(), + max_download_threads=10, + buffer_size_bytes=BUFFER, + ): + """Returns the result set and the kwargs of every fetch_results call.""" + calls = [] + queues = list(fetched_queues) + + def fetch_results(**kwargs): + calls.append(kwargs) + queue = queues.pop(0) + return queue, bool(queues), len(getattr(queue, "result_links", [])) + + backend = Mock(spec=ThriftDatabricksClient) + backend.fetch_results = fetch_results + t_row_set = None + if initial_link_sizes is not None: + t_row_set = ttypes.TRowSet( + startRowOffset=0, + rows=[], + resultLinks=[self._link(b) for b in initial_link_sizes], + ) + with patch( + "databricks.sql.utils.ThriftResultSetQueueFactory.build_queue", + return_value=Mock(), + ): + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=Mock(), + status=None, + has_been_closed_server_side=False, + description=[], + result_format=ttypes.TSparkRowSetType.URL_BASED_SET, + ), + thrift_client=backend, + buffer_size_bytes=buffer_size_bytes, + arraysize=1000, + t_row_set=t_row_set, + max_download_threads=max_download_threads, + ) + return rs, calls + + def test_requests_one_link_per_download_thread(self): + rs, calls = self._make_result_set( + initial_link_sizes=[20 * self.MIB, 20 * self.MIB, 5 * self.MIB], + fetched_queues=[ + self._cloudfetch_queue(21 * self.MIB), + self._cloudfetch_queue(), + ], + ) + + rs._fill_results_buffer() + rs._fill_results_buffer() + + self.assertEqual( + [c["max_bytes"] for c in calls], [200 * self.MIB, 210 * self.MIB] + ) + self.assertEqual([c["max_rows"] for c in calls], [1000, 1000]) + + def test_scales_with_max_download_threads(self): + rs, calls = self._make_result_set( + initial_link_sizes=[20 * self.MIB], + fetched_queues=[self._cloudfetch_queue()], + max_download_threads=15, + ) + + rs._fill_results_buffer() + + self.assertEqual(calls[0]["max_bytes"], 300 * self.MIB) + + def test_caps_at_four_buffers(self): + rs, calls = self._make_result_set( + initial_link_sizes=[20 * self.MIB], + fetched_queues=[self._cloudfetch_queue()], + max_download_threads=50, + ) + + rs._fill_results_buffer() + + self.assertEqual(calls[0]["max_bytes"], 4 * self.BUFFER) + + def test_links_without_size_keep_buffer_size(self): + rs, calls = self._make_result_set( + initial_link_sizes=[None], + fetched_queues=[self._cloudfetch_queue(None), self._cloudfetch_queue()], + ) + + rs._fill_results_buffer() + rs._fill_results_buffer() + + self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, self.BUFFER]) + + def test_small_buffer_size_caps_batch(self): + rs, calls = self._make_result_set( + initial_link_sizes=[20 * self.MIB], + fetched_queues=[self._cloudfetch_queue()], + buffer_size_bytes=10 * self.MIB, + ) + + rs._fill_results_buffer() + + self.assertEqual(calls[0]["max_bytes"], 40 * self.MIB) + + def test_never_requests_less_than_buffer_size(self): + rs, calls = self._make_result_set( + initial_link_sizes=[20 * self.MIB], + fetched_queues=[self._cloudfetch_queue()], + max_download_threads=2, + ) + + rs._fill_results_buffer() + + self.assertEqual(calls[0]["max_bytes"], self.BUFFER) + + def test_inline_results_keep_buffer_size(self): + # No initial rows: the constructor fetches the first batch itself. + rs, calls = self._make_result_set( + fetched_queues=[Mock(spec=ArrowQueue), Mock(spec=ArrowQueue)] + ) + + rs._fill_results_buffer() + + self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, self.BUFFER]) + + def test_first_fetch_without_links_uses_buffer_size(self): + rs, calls = self._make_result_set( + fetched_queues=[ + self._cloudfetch_queue(20 * self.MIB), + self._cloudfetch_queue(), + ] + ) + + rs._fill_results_buffer() + + self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, 200 * self.MIB]) + + if __name__ == "__main__": unittest.main()