Skip to content

Commit 5d10365

Browse files
vuanhphungIsaac
andcommitted
perf(thrift): fetch more inline rows per round trip for fetchall and large fetchmany
Every FetchResults asked for arraysize (100k) rows, so a 10M-row boolean result took 97 round trips although each reply was tiny. fetchall*() and fetchmany*(size > arraysize) now ask for the rows they still need, capped at 2M rows (the JDBC driver's limit). fetchone(), iteration and fetchmany(size <= arraysize) still send arraysize, and replies stay bounded by maxBytes. Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 8d6db43 commit 5d10365

3 files changed

Lines changed: 135 additions & 10 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 `fetchall()` and large `fetchmany()` calls on narrow inline Thrift results: they now ask the server for up to 2 million rows per round trip instead of `arraysize` rows. `fetchone()`, iteration and `fetchmany(size)` with `size <= arraysize` still fetch `arraysize` rows at a time.
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: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from __future__ import annotations
22

33
from abc import ABC, abstractmethod
4-
from typing import List, Optional, TYPE_CHECKING, Tuple
4+
from typing import List, Optional, TYPE_CHECKING, Tuple, cast
55

66
import logging
77
import pandas
@@ -27,6 +27,10 @@
2727

2828
logger = logging.getLogger(__name__)
2929

30+
# Most rows a Thrift FetchResults asks for when the caller wants more than arraysize rows
31+
# (the JDBC driver's per-fetch row limit).
32+
_MAX_ROWS_PER_FETCH = 2_000_000
33+
3034

3135
class ResultSet(ABC):
3236
"""
@@ -274,10 +278,18 @@ def __init__(
274278
if not self.results:
275279
self._fill_results_buffer()
276280

277-
def _fill_results_buffer(self):
278-
results, has_more_rows, result_links_count = self.backend.fetch_results(
281+
def _fill_results_buffer(self, rows_wanted: Optional[int] = None):
282+
# Streaming callers (fetchone, iteration, fetchmany(size <= arraysize)) keep
283+
# arraysize-row batches. fetchall and larger fetchmany calls ask for the rows they
284+
# still need, so narrow inline results take fewer round trips; replies remain
285+
# bounded by maxBytes.
286+
max_rows = self.arraysize
287+
if rows_wanted is not None:
288+
max_rows = max(max_rows, min(rows_wanted, _MAX_ROWS_PER_FETCH))
289+
backend = cast("ThriftDatabricksClient", self.backend)
290+
results, has_more_rows, result_links_count = backend.fetch_results(
279291
command_id=self.command_id,
280-
max_rows=self.arraysize,
292+
max_rows=max_rows,
281293
max_bytes=self.buffer_size_bytes,
282294
expected_row_start_offset=self._next_row_index,
283295
lz4_compressed=self.lz4_compressed,
@@ -331,7 +343,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
331343
and not self.has_been_closed_server_side
332344
and self.has_more_rows
333345
):
334-
self._fill_results_buffer()
346+
self._fill_results_buffer(rows_wanted=n_remaining_rows)
335347
partial_results = self.results.next_n_rows(n_remaining_rows)
336348
if partial_results.num_rows == 0:
337349
continue
@@ -360,7 +372,7 @@ def fetchmany_columnar(self, size: int):
360372
and not self.has_been_closed_server_side
361373
and self.has_more_rows
362374
):
363-
self._fill_results_buffer()
375+
self._fill_results_buffer(rows_wanted=n_remaining_rows)
364376
partial_results = self.results.next_n_rows(n_remaining_rows)
365377
partial_result_chunks.append(partial_results)
366378
n_remaining_rows -= partial_results.num_rows
@@ -384,7 +396,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
384396
self._next_row_index += results.num_rows
385397

386398
while not self.has_been_closed_server_side and self.has_more_rows:
387-
self._fill_results_buffer()
399+
self._fill_results_buffer(rows_wanted=_MAX_ROWS_PER_FETCH)
388400
partial_results = self.results.remaining_rows()
389401
if partial_results.num_rows == 0:
390402
continue
@@ -413,7 +425,7 @@ def fetchall_columnar(self):
413425
self._next_row_index += results.num_rows
414426
partial_result_chunks = [results]
415427
while not self.has_been_closed_server_side and self.has_more_rows:
416-
self._fill_results_buffer()
428+
self._fill_results_buffer(rows_wanted=_MAX_ROWS_PER_FETCH)
417429
partial_results = self.results.remaining_rows()
418430
partial_result_chunks.append(partial_results)
419431
self._next_row_index += partial_results.num_rows

‎tests/unit/test_fetches.py‎

Lines changed: 114 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,7 +9,7 @@
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, ColumnQueue, ColumnTable
1313
from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
1414
from databricks.sql.result_set import ThriftResultSet
1515

@@ -391,5 +391,117 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self):
391391
self.assertEqual(result.num_rows, 0)
392392

393393

394+
class FetchRowsPerRequestTests(unittest.TestCase):
395+
"""maxRows sent by ThriftResultSet._fill_results_buffer for inline results."""
396+
397+
@staticmethod
398+
def _make_result_set(n_rows, arraysize, columnar=False):
399+
"""A result set over ids 0..n_rows-1 from a fake server that returns exactly
400+
maxRows rows per FetchResults. Returns the result set and the maxRows sent."""
401+
calls = []
402+
next_row = 0
403+
404+
def fetch_results(**kwargs):
405+
nonlocal next_row
406+
calls.append(kwargs["max_rows"])
407+
values = list(range(next_row, min(n_rows, next_row + kwargs["max_rows"])))
408+
next_row += len(values)
409+
if columnar:
410+
queue = ColumnQueue(ColumnTable([values], ["col0"]))
411+
else:
412+
queue = ArrowQueue(pa.table({"col0": values}), len(values))
413+
return queue, next_row < n_rows, 0
414+
415+
backend = Mock(spec=ThriftDatabricksClient)
416+
backend.fetch_results = fetch_results
417+
rs = ThriftResultSet(
418+
connection=Mock(),
419+
execute_response=ExecuteResponse(
420+
command_id=None,
421+
status=None,
422+
has_been_closed_server_side=False,
423+
description=[("col0", "integer", None, None, None, None, None)],
424+
lz4_compressed=True,
425+
is_staging_operation=False,
426+
),
427+
thrift_client=backend,
428+
arraysize=arraysize,
429+
)
430+
return rs, calls
431+
432+
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
433+
def test_fetchone_and_iteration_keep_arraysize(self):
434+
rs, calls = self._make_result_set(25, arraysize=10)
435+
436+
rows = list(rs)
437+
438+
self.assertEqual([r[0] for r in rows], list(range(25)))
439+
self.assertEqual(calls, [10, 10, 10])
440+
441+
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
442+
def test_fetchmany_up_to_arraysize_keeps_arraysize(self):
443+
rs, calls = self._make_result_set(25, arraysize=10)
444+
445+
ids = [r[0] for _ in range(4) for r in rs.fetchmany(7)]
446+
447+
self.assertEqual(ids, list(range(25)))
448+
self.assertEqual(calls, [10, 10, 10])
449+
450+
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
451+
def test_fetchall_requests_up_to_the_cap(self):
452+
rs, calls = self._make_result_set(25, arraysize=10)
453+
454+
table = rs.fetchall_arrow()
455+
456+
self.assertEqual(table.column(0).to_pylist(), list(range(25)))
457+
self.assertEqual(calls, [10, 2_000_000])
458+
459+
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
460+
def test_large_fetchmany_requests_the_remaining_rows(self):
461+
rs, calls = self._make_result_set(100, arraysize=10)
462+
463+
first = rs.fetchmany_arrow(35)
464+
second = rs.fetchmany_arrow(35)
465+
466+
self.assertEqual(first.column(0).to_pylist(), list(range(35)))
467+
self.assertEqual(second.column(0).to_pylist(), list(range(35, 70)))
468+
self.assertEqual(calls, [10, 25, 35])
469+
470+
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
471+
def test_rows_per_request_are_capped(self):
472+
with patch("databricks.sql.result_set._MAX_ROWS_PER_FETCH", 50):
473+
rs, calls = self._make_result_set(200, arraysize=10)
474+
table = rs.fetchmany_arrow(120)
475+
476+
self.assertEqual(table.column(0).to_pylist(), list(range(120)))
477+
self.assertEqual(calls, [10, 50, 50, 10])
478+
479+
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
480+
def test_arraysize_above_the_cap_is_kept(self):
481+
with patch("databricks.sql.result_set._MAX_ROWS_PER_FETCH", 50):
482+
rs, calls = self._make_result_set(200, arraysize=80)
483+
table = rs.fetchall_arrow()
484+
485+
self.assertEqual(table.num_rows, 200)
486+
self.assertEqual(calls, [80, 80, 80])
487+
488+
def test_columnar_fetchall_requests_up_to_the_cap(self):
489+
rs, calls = self._make_result_set(25, arraysize=10, columnar=True)
490+
491+
rows = rs.fetchall()
492+
493+
self.assertEqual([r[0] for r in rows], list(range(25)))
494+
self.assertEqual(calls, [10, 2_000_000])
495+
496+
def test_columnar_fetchone_and_fetchmany(self):
497+
rs, calls = self._make_result_set(40, arraysize=10, columnar=True)
498+
499+
ids = [rs.fetchone()[0]] + [r[0] for r in rs.fetchmany(30)]
500+
ids += [r[0] for r in rs.fetchmany(5)]
501+
502+
self.assertEqual(ids, list(range(36)))
503+
self.assertEqual(calls, [10, 21, 10])
504+
505+
394506
if __name__ == "__main__":
395507
unittest.main()

0 commit comments

Comments
 (0)