perf(thrift): fetch more inline rows per round trip for fetchall and large fetchmany

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.

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

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()