From 5d1036556eeef058ce2830d9d36693ad5c38bc42 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 09:22:15 +0000 Subject: [PATCH] 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 Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/result_set.py | 28 +++++--- tests/unit/test_fetches.py | 116 ++++++++++++++++++++++++++++++- 3 files changed, 135 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..c7351a300 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 `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. # 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..904817101 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 @@ -27,6 +27,10 @@ logger = logging.getLogger(__name__) +# Most rows a Thrift FetchResults asks for when the caller wants more than arraysize rows +# (the JDBC driver's per-fetch row limit). +_MAX_ROWS_PER_FETCH = 2_000_000 + class ResultSet(ABC): """ @@ -274,10 +278,18 @@ def __init__( if not self.results: self._fill_results_buffer() - def _fill_results_buffer(self): - results, has_more_rows, result_links_count = self.backend.fetch_results( + def _fill_results_buffer(self, rows_wanted: Optional[int] = None): + # Streaming callers (fetchone, iteration, fetchmany(size <= arraysize)) keep + # arraysize-row batches. fetchall and larger fetchmany calls ask for the rows they + # still need, so narrow inline results take fewer round trips; replies remain + # bounded by maxBytes. + max_rows = self.arraysize + if rows_wanted is not None: + max_rows = max(max_rows, min(rows_wanted, _MAX_ROWS_PER_FETCH)) + backend = cast("ThriftDatabricksClient", self.backend) + results, has_more_rows, result_links_count = backend.fetch_results( command_id=self.command_id, - max_rows=self.arraysize, + max_rows=max_rows, max_bytes=self.buffer_size_bytes, expected_row_start_offset=self._next_row_index, lz4_compressed=self.lz4_compressed, @@ -331,7 +343,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table": and not self.has_been_closed_server_side and self.has_more_rows ): - self._fill_results_buffer() + self._fill_results_buffer(rows_wanted=n_remaining_rows) partial_results = self.results.next_n_rows(n_remaining_rows) if partial_results.num_rows == 0: continue @@ -360,7 +372,7 @@ def fetchmany_columnar(self, size: int): and not self.has_been_closed_server_side and self.has_more_rows ): - self._fill_results_buffer() + self._fill_results_buffer(rows_wanted=n_remaining_rows) partial_results = self.results.next_n_rows(n_remaining_rows) partial_result_chunks.append(partial_results) n_remaining_rows -= partial_results.num_rows @@ -384,7 +396,7 @@ def fetchall_arrow(self) -> "pyarrow.Table": self._next_row_index += results.num_rows while not self.has_been_closed_server_side and self.has_more_rows: - self._fill_results_buffer() + self._fill_results_buffer(rows_wanted=_MAX_ROWS_PER_FETCH) partial_results = self.results.remaining_rows() if partial_results.num_rows == 0: continue @@ -413,7 +425,7 @@ def fetchall_columnar(self): self._next_row_index += results.num_rows partial_result_chunks = [results] while not self.has_been_closed_server_side and self.has_more_rows: - self._fill_results_buffer() + self._fill_results_buffer(rows_wanted=_MAX_ROWS_PER_FETCH) partial_results = self.results.remaining_rows() partial_result_chunks.append(partial_results) self._next_row_index += partial_results.num_rows diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py index e3a963355..55fbcda9d 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,7 +9,7 @@ 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, ColumnQueue, ColumnTable from databricks.sql.backend.thrift_backend import ThriftDatabricksClient from databricks.sql.result_set import ThriftResultSet @@ -391,5 +391,117 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self): self.assertEqual(result.num_rows, 0) +class FetchRowsPerRequestTests(unittest.TestCase): + """maxRows sent by ThriftResultSet._fill_results_buffer for inline results.""" + + @staticmethod + def _make_result_set(n_rows, arraysize, columnar=False): + """A result set over ids 0..n_rows-1 from a fake server that returns exactly + maxRows rows per FetchResults. Returns the result set and the maxRows sent.""" + calls = [] + next_row = 0 + + def fetch_results(**kwargs): + nonlocal next_row + calls.append(kwargs["max_rows"]) + values = list(range(next_row, min(n_rows, next_row + kwargs["max_rows"]))) + next_row += len(values) + if columnar: + queue = ColumnQueue(ColumnTable([values], ["col0"])) + else: + queue = ArrowQueue(pa.table({"col0": values}), len(values)) + return queue, next_row < n_rows, 0 + + backend = Mock(spec=ThriftDatabricksClient) + backend.fetch_results = fetch_results + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=None, + status=None, + has_been_closed_server_side=False, + description=[("col0", "integer", None, None, None, None, None)], + lz4_compressed=True, + is_staging_operation=False, + ), + thrift_client=backend, + arraysize=arraysize, + ) + return rs, calls + + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") + def test_fetchone_and_iteration_keep_arraysize(self): + rs, calls = self._make_result_set(25, arraysize=10) + + rows = list(rs) + + self.assertEqual([r[0] for r in rows], list(range(25))) + self.assertEqual(calls, [10, 10, 10]) + + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") + def test_fetchmany_up_to_arraysize_keeps_arraysize(self): + rs, calls = self._make_result_set(25, arraysize=10) + + ids = [r[0] for _ in range(4) for r in rs.fetchmany(7)] + + self.assertEqual(ids, list(range(25))) + self.assertEqual(calls, [10, 10, 10]) + + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") + def test_fetchall_requests_up_to_the_cap(self): + rs, calls = self._make_result_set(25, arraysize=10) + + table = rs.fetchall_arrow() + + self.assertEqual(table.column(0).to_pylist(), list(range(25))) + self.assertEqual(calls, [10, 2_000_000]) + + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") + def test_large_fetchmany_requests_the_remaining_rows(self): + rs, calls = self._make_result_set(100, arraysize=10) + + first = rs.fetchmany_arrow(35) + second = rs.fetchmany_arrow(35) + + self.assertEqual(first.column(0).to_pylist(), list(range(35))) + self.assertEqual(second.column(0).to_pylist(), list(range(35, 70))) + self.assertEqual(calls, [10, 25, 35]) + + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") + def test_rows_per_request_are_capped(self): + with patch("databricks.sql.result_set._MAX_ROWS_PER_FETCH", 50): + rs, calls = self._make_result_set(200, arraysize=10) + table = rs.fetchmany_arrow(120) + + self.assertEqual(table.column(0).to_pylist(), list(range(120))) + self.assertEqual(calls, [10, 50, 50, 10]) + + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") + def test_arraysize_above_the_cap_is_kept(self): + with patch("databricks.sql.result_set._MAX_ROWS_PER_FETCH", 50): + rs, calls = self._make_result_set(200, arraysize=80) + table = rs.fetchall_arrow() + + self.assertEqual(table.num_rows, 200) + self.assertEqual(calls, [80, 80, 80]) + + def test_columnar_fetchall_requests_up_to_the_cap(self): + rs, calls = self._make_result_set(25, arraysize=10, columnar=True) + + rows = rs.fetchall() + + self.assertEqual([r[0] for r in rows], list(range(25))) + self.assertEqual(calls, [10, 2_000_000]) + + def test_columnar_fetchone_and_fetchmany(self): + rs, calls = self._make_result_set(40, arraysize=10, columnar=True) + + ids = [rs.fetchone()[0]] + [r[0] for r in rs.fetchmany(30)] + ids += [r[0] for r in rs.fetchmany(5)] + + self.assertEqual(ids, list(range(36))) + self.assertEqual(calls, [10, 21, 10]) + + if __name__ == "__main__": unittest.main()