From 650fd5dc38202c74bcd4b36ff5d417ce7d47f3d0 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 07:59:01 +0000 Subject: [PATCH 1/3] perf(thrift): request about one CloudFetch file per download thread The server packs CloudFetch links into a FetchResults reply up to maxBytes of file data and ignores maxRows for links. With the 100 MiB default that is 4-5 links of ~21 MB, so a 4.9 GB result took 60 round trips and kept only 4-5 of the 10 download threads busy. Request max(buffer_size_bytes, max_download_threads * largest link seen) instead. Every link in a batch then starts downloading immediately, so no link waits towards its expiry. Inline results are unchanged. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/result_set.py | 22 +++++- tests/unit/test_fetches.py | 131 ++++++++++++++++++++++++++++++- 3 files changed, 151 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..879a03bf0 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 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..1006f7038 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 @@ -229,6 +230,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 +254,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 +278,25 @@ def __init__( if not self.results: self._fill_results_buffer() + def _note_link_sizes(self, links) -> None: + for link in links: + if link.bytesNum > self._max_link_bytes: + self._max_link_bytes = link.bytesNum + + 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. + return max( + self.buffer_size_bytes, self._max_download_threads * self._max_link_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 +307,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..f3a7487f3 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,131 @@ 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 + ): + """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=self.BUFFER, + 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=25, + ) + + rs._fill_results_buffer() + + self.assertEqual(calls[0]["max_bytes"], 500 * 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() From 9ae354838f4b49937e15ffe79be28bef938f3c99 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 08:13:36 +0000 Subject: [PATCH 2/3] Cap the CloudFetch FetchResults budget at 4x buffer_size_bytes Without a cap the budget grows with max_download_threads and the largest link seen, and it ignored a deliberately small buffer_size_bytes. The default path (about 210 MB, 10 links) is unchanged. Also document that buffer_size_bytes is no longer the max fetch size for CloudFetch. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 2 +- src/databricks/sql/result_set.py | 18 +++++++++++++---- tests/unit/test_fetches.py | 34 ++++++++++++++++++++++++++++---- 3 files changed, 45 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a03bf0..4038e882d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +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 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: 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 1006f7038..05976e236 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -28,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): """ @@ -60,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 @@ -217,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) @@ -288,8 +294,12 @@ def _fetch_max_bytes(self) -> int: # 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. - return max( - self.buffer_size_bytes, self._max_download_threads * self._max_link_bytes + # 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): diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py index f3a7487f3..98dfd9e81 100644 --- a/tests/unit/test_fetches.py +++ b/tests/unit/test_fetches.py @@ -415,7 +415,11 @@ def _cloudfetch_queue(cls, *link_sizes): return queue def _make_result_set( - self, initial_link_sizes=None, fetched_queues=(), max_download_threads=10 + 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 = [] @@ -449,7 +453,7 @@ def fetch_results(**kwargs): result_format=ttypes.TSparkRowSetType.URL_BASED_SET, ), thrift_client=backend, - buffer_size_bytes=self.BUFFER, + buffer_size_bytes=buffer_size_bytes, arraysize=1000, t_row_set=t_row_set, max_download_threads=max_download_threads, @@ -477,12 +481,34 @@ 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=25, + max_download_threads=15, ) rs._fill_results_buffer() - self.assertEqual(calls[0]["max_bytes"], 500 * self.MIB) + 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_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( From 0a75039a17cbf657d9f8aabff7e4c9a6e1cc5294 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 08:18:29 +0000 Subject: [PATCH 3/3] Treat a CloudFetch link without bytesNum as size 0 The Thrift reader doesn't enforce required fields, so a missing size must not fail the fetch; it only disables the sizing heuristic. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/result_set.py | 7 +++++-- tests/unit/test_fetches.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index 05976e236..2a2599259 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -286,8 +286,11 @@ def __init__( def _note_link_sizes(self, links) -> None: for link in links: - if link.bytesNum > self._max_link_bytes: - self._max_link_bytes = link.bytesNum + # 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 diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py index 98dfd9e81..d0d74df1d 100644 --- a/tests/unit/test_fetches.py +++ b/tests/unit/test_fetches.py @@ -499,6 +499,17 @@ def test_caps_at_four_buffers(self): 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],