Skip to content

Commit 650fd5d

Browse files
vuanhphungIsaac
andcommitted
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 <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 8d6db43 commit 650fd5d

3 files changed

Lines changed: 151 additions & 3 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
# Unreleased
44
- 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))
55
- 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.
6+
- 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.
67

78
# 4.9.0 (2026-11-04)
89
- Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes.

‎src/databricks/sql/result_set.py‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from databricks.sql.utils import (
2121
ColumnTable,
2222
ColumnQueue,
23+
ThriftCloudFetchQueue,
2324
concat_table_chunks,
2425
)
2526
from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse
@@ -229,6 +230,8 @@ def __init__(
229230
# Initialize ThriftResultSet-specific attributes
230231
self._use_cloud_fetch = use_cloud_fetch
231232
self.has_more_rows = has_more_rows
233+
self._max_download_threads = max_download_threads
234+
self._max_link_bytes = 0
232235

233236
# Build the results queue if t_row_set is provided
234237
results_queue = None
@@ -251,6 +254,7 @@ def __init__(
251254
)
252255
if t_row_set.resultLinks:
253256
self.num_chunks += len(t_row_set.resultLinks)
257+
self._note_link_sizes(t_row_set.resultLinks)
254258

255259
# Call parent constructor with common attributes
256260
super().__init__(
@@ -274,11 +278,25 @@ def __init__(
274278
if not self.results:
275279
self._fill_results_buffer()
276280

281+
def _note_link_sizes(self, links) -> None:
282+
for link in links:
283+
if link.bytesNum > self._max_link_bytes:
284+
self._max_link_bytes = link.bytesNum
285+
286+
def _fetch_max_bytes(self) -> int:
287+
# The server packs CloudFetch links into a FetchResults reply up to maxBytes of
288+
# file data, so buffer_size_bytes alone yields only a few links per round trip.
289+
# Ask for about one link per download thread: every link in the batch then starts
290+
# downloading immediately, so none sits waiting while it ages towards expiry.
291+
return max(
292+
self.buffer_size_bytes, self._max_download_threads * self._max_link_bytes
293+
)
294+
277295
def _fill_results_buffer(self):
278296
results, has_more_rows, result_links_count = self.backend.fetch_results(
279297
command_id=self.command_id,
280298
max_rows=self.arraysize,
281-
max_bytes=self.buffer_size_bytes,
299+
max_bytes=self._fetch_max_bytes(),
282300
expected_row_start_offset=self._next_row_index,
283301
lz4_compressed=self.lz4_compressed,
284302
arrow_schema_bytes=self._arrow_schema_bytes,
@@ -289,6 +307,8 @@ def _fill_results_buffer(self):
289307
self.results = results
290308
self.has_more_rows = has_more_rows
291309
self.num_chunks += result_links_count
310+
if isinstance(results, ThriftCloudFetchQueue):
311+
self._note_link_sizes(results.result_links)
292312

293313
def _convert_columnar_table(self, table):
294314
column_names = [c[0] for c in self.description]

‎tests/unit/test_fetches.py‎

Lines changed: 129 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import unittest
22
import pytest
3-
from unittest.mock import Mock
3+
from unittest.mock import Mock, patch
44

55
try:
66
import pyarrow as pa
@@ -9,9 +9,10 @@
99

1010
import databricks.sql.client as client
1111
from databricks.sql.backend.types import ExecuteResponse
12-
from databricks.sql.utils import ArrowQueue
12+
from databricks.sql.utils import ArrowQueue, ThriftCloudFetchQueue
1313
from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
1414
from databricks.sql.result_set import ThriftResultSet
15+
from databricks.sql.thrift_api.TCLIService import ttypes
1516

1617

1718
class _StubArrowQueue:
@@ -391,5 +392,131 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self):
391392
self.assertEqual(result.num_rows, 0)
392393

393394

395+
class CloudFetchBatchSizeTests(unittest.TestCase):
396+
"""maxBytes sent by ThriftResultSet._fill_results_buffer."""
397+
398+
MIB = 2**20
399+
BUFFER = 100 * MIB
400+
401+
@staticmethod
402+
def _link(bytes_num):
403+
return ttypes.TSparkArrowResultLink(
404+
fileLink="https://example.com/file",
405+
expiryTime=0,
406+
startRowOffset=0,
407+
rowCount=1,
408+
bytesNum=bytes_num,
409+
)
410+
411+
@classmethod
412+
def _cloudfetch_queue(cls, *link_sizes):
413+
queue = Mock(spec=ThriftCloudFetchQueue)
414+
queue.result_links = [cls._link(b) for b in link_sizes]
415+
return queue
416+
417+
def _make_result_set(
418+
self, initial_link_sizes=None, fetched_queues=(), max_download_threads=10
419+
):
420+
"""Returns the result set and the kwargs of every fetch_results call."""
421+
calls = []
422+
queues = list(fetched_queues)
423+
424+
def fetch_results(**kwargs):
425+
calls.append(kwargs)
426+
queue = queues.pop(0)
427+
return queue, bool(queues), len(getattr(queue, "result_links", []))
428+
429+
backend = Mock(spec=ThriftDatabricksClient)
430+
backend.fetch_results = fetch_results
431+
t_row_set = None
432+
if initial_link_sizes is not None:
433+
t_row_set = ttypes.TRowSet(
434+
startRowOffset=0,
435+
rows=[],
436+
resultLinks=[self._link(b) for b in initial_link_sizes],
437+
)
438+
with patch(
439+
"databricks.sql.utils.ThriftResultSetQueueFactory.build_queue",
440+
return_value=Mock(),
441+
):
442+
rs = ThriftResultSet(
443+
connection=Mock(),
444+
execute_response=ExecuteResponse(
445+
command_id=Mock(),
446+
status=None,
447+
has_been_closed_server_side=False,
448+
description=[],
449+
result_format=ttypes.TSparkRowSetType.URL_BASED_SET,
450+
),
451+
thrift_client=backend,
452+
buffer_size_bytes=self.BUFFER,
453+
arraysize=1000,
454+
t_row_set=t_row_set,
455+
max_download_threads=max_download_threads,
456+
)
457+
return rs, calls
458+
459+
def test_requests_one_link_per_download_thread(self):
460+
rs, calls = self._make_result_set(
461+
initial_link_sizes=[20 * self.MIB, 20 * self.MIB, 5 * self.MIB],
462+
fetched_queues=[
463+
self._cloudfetch_queue(21 * self.MIB),
464+
self._cloudfetch_queue(),
465+
],
466+
)
467+
468+
rs._fill_results_buffer()
469+
rs._fill_results_buffer()
470+
471+
self.assertEqual(
472+
[c["max_bytes"] for c in calls], [200 * self.MIB, 210 * self.MIB]
473+
)
474+
self.assertEqual([c["max_rows"] for c in calls], [1000, 1000])
475+
476+
def test_scales_with_max_download_threads(self):
477+
rs, calls = self._make_result_set(
478+
initial_link_sizes=[20 * self.MIB],
479+
fetched_queues=[self._cloudfetch_queue()],
480+
max_download_threads=25,
481+
)
482+
483+
rs._fill_results_buffer()
484+
485+
self.assertEqual(calls[0]["max_bytes"], 500 * self.MIB)
486+
487+
def test_never_requests_less_than_buffer_size(self):
488+
rs, calls = self._make_result_set(
489+
initial_link_sizes=[20 * self.MIB],
490+
fetched_queues=[self._cloudfetch_queue()],
491+
max_download_threads=2,
492+
)
493+
494+
rs._fill_results_buffer()
495+
496+
self.assertEqual(calls[0]["max_bytes"], self.BUFFER)
497+
498+
def test_inline_results_keep_buffer_size(self):
499+
# No initial rows: the constructor fetches the first batch itself.
500+
rs, calls = self._make_result_set(
501+
fetched_queues=[Mock(spec=ArrowQueue), Mock(spec=ArrowQueue)]
502+
)
503+
504+
rs._fill_results_buffer()
505+
506+
self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, self.BUFFER])
507+
508+
def test_first_fetch_without_links_uses_buffer_size(self):
509+
rs, calls = self._make_result_set(
510+
fetched_queues=[
511+
self._cloudfetch_queue(20 * self.MIB),
512+
self._cloudfetch_queue(),
513+
]
514+
)
515+
516+
rs._fill_results_buffer()
517+
518+
self.assertEqual([c["max_bytes"] for c in calls], [self.BUFFER, 200 * self.MIB])
519+
520+
394521
if __name__ == "__main__":
395522
unittest.main()

0 commit comments

Comments
 (0)