perf: convert large Arrow results to rows in slices

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.
+- Lower the peak memory of `fetchall()` and large `fetchmany()` calls by converting Arrow results to rows 100,000 rows at a time instead of building one DataFrame for the whole result.
 
 # 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

@@ -25,6 +25,25 @@
 from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse
 from databricks.sql.telemetry.models.event import StatementType
 
+# Rows per to_pandas() call when converting an Arrow table to rows.
+_ROWS_PER_CONVERSION = 100_000
+
+
+def _converts_per_value(arrow_type) -> bool:
+    """Whether to_pandas() converts each value of this type independently."""
+    types = pyarrow.types
+    return (
+        types.is_primitive(arrow_type)
+        or types.is_null(arrow_type)
+        or types.is_string(arrow_type)
+        or types.is_large_string(arrow_type)
+        or types.is_binary(arrow_type)
+        or types.is_large_binary(arrow_type)
+        or types.is_fixed_size_binary(arrow_type)
+        or types.is_decimal(arrow_type)
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -121,16 +140,41 @@ def _convert_arrow_table(self, table):
             pyarrow.string(): pandas.StringDtype(),
         }
 
+        def to_object_array(t):
+            df = t.to_pandas(
+                types_mapper=dtype_mapping.get,
+                date_as_object=True,
+                timestamp_as_object=True,
+            )
+            return df.to_numpy(na_value=None, dtype="object")
+
         # Need to rename columns, as the to_pandas function cannot handle duplicate column names
         table_renamed = table.rename_columns([str(c) for c in range(table.num_columns)])
-        df = table_renamed.to_pandas(
-            types_mapper=dtype_mapping.get,
-            date_as_object=True,
-            timestamp_as_object=True,
-        )
 
-        res = df.to_numpy(na_value=None, dtype="object")
-        return [ResultRow(*v) for v in res]
+        # Convert in slices so that only one slice's DataFrame and object array are
+        # alive next to the rows. Columns whose conversion depends on the column's
+        # other values (pyarrow makes every list<int> value a float array if any
+        # element is NULL) are converted whole, as before.
+        per_value = [_converts_per_value(f.type) for f in table.schema]
+        if table.num_rows <= _ROWS_PER_CONVERSION or not any(per_value):
+            return [ResultRow(*v) for v in to_object_array(table_renamed)]
+
+        whole = [i for i, ok in enumerate(per_value) if not ok]
+        whole_values = to_object_array(table_renamed.select(whole)) if whole else None
+        # Keep an object column where those columns were: pandas' to_numpy converts
+        # some dtypes (e.g. durations) differently for a single-block DataFrame.
+        sliced_table = table_renamed
+        for i in whole:
+            sliced_table = sliced_table.set_column(
+                i, str(i), pyarrow.nulls(table.num_rows)
+            )
+        rows = []
+        for offset in range(0, table.num_rows, _ROWS_PER_CONVERSION):
+            res = to_object_array(sliced_table.slice(offset, _ROWS_PER_CONVERSION))
+            if whole_values is not None:
+                res[:, whole] = whole_values[offset : offset + len(res)]
+            rows.extend([ResultRow(*v) for v in res])
+        return rows
 
     @property
     def rownumber(self):

tests/unit/test_pandas_compatibility.py

@@ -9,8 +9,9 @@
 import datetime
 import unittest
 from decimal import Decimal
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
+import numpy
 import pandas
 import pytest
 
@@ -19,6 +20,7 @@
 except ImportError:
     pa = None
 
+from databricks.sql import result_set
 from databricks.sql.result_set import ThriftResultSet
 
 
@@ -458,5 +460,102 @@ def test_disable_pandas_path(self):
         self.assertEqual(rows[1].name, "b")
 
 
+@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
+class TestConvertArrowTableInSlices(unittest.TestCase):
+    """Tables larger than _ROWS_PER_CONVERSION are converted in slices."""
+
+    def _table(self):
+        tz = datetime.timezone.utc
+        return pa.table(
+            [
+                pa.chunked_array([[1, None, 3], [4, 5, 6, None]], pa.int32()),
+                pa.array([1.5, None, float("inf"), 4.0, 5.0, 6.0, -7.0]),
+                pa.array(["a", None, "c", "d", "e", "f", "g"]),
+                pa.array(
+                    [Decimal(i).scaleb(-2) if i != 4 else None for i in range(7)],
+                    pa.decimal128(5, 2),
+                ),
+                pa.array([datetime.date(2024, 1, i + 1) for i in range(7)]),
+                pa.array(
+                    [datetime.datetime(2024, 3, 10, i, tzinfo=tz) for i in range(7)],
+                    pa.timestamp("us", tz="America/Los_Angeles"),
+                ),
+                pa.array([True, False, None, True, False, None, True]),
+                # pyarrow picks the element dtype of every list from the whole
+                # column, so the NULL in the last row affects the first row too.
+                pa.array([[i, i] for i in range(6)] + [[None]], pa.list_(pa.int64())),
+                pa.array(
+                    [{"x": i if i != 2 else None} for i in range(7)],
+                    pa.struct([("x", pa.int64())]),
+                ),
+                # Older pandas returns durations from a single-block DataFrame
+                # as ints, so slices must keep the whole table's block layout.
+                pa.array(
+                    [datetime.timedelta(seconds=i) for i in range(7)],
+                    pa.duration("us"),
+                ),
+            ],
+            names=["a", "b", "a", "d", "e", "f", "g", "h", "i", "j"],
+        )
+
+    def _convert(self, table, rows_per_conversion):
+        description = [
+            (name, "string", None, None, None, None, None)
+            for name in table.column_names
+        ]
+        with patch.object(result_set, "_ROWS_PER_CONVERSION", rows_per_conversion):
+            return _make_result_set(description)._convert_arrow_table(table)
+
+    def _to_numpy_row_counts(self, table, rows_per_conversion):
+        counts = []
+        to_numpy = pandas.DataFrame.to_numpy
+
+        def spy(df, *args, **kwargs):
+            counts.append(len(df))
+            return to_numpy(df, *args, **kwargs)
+
+        with patch.object(pandas.DataFrame, "to_numpy", spy):
+            rows = self._convert(table, rows_per_conversion)
+        self.assertEqual(len(rows), table.num_rows)
+        return counts
+
+    def test_sliced_conversion_matches_whole_table(self):
+        table = self._table()
+        whole = self._convert(table, table.num_rows)
+        for rows_per_conversion in (1, 2, 3, 6):
+            rows = self._convert(table, rows_per_conversion)
+            self.assertEqual(len(rows), len(whole))
+            for row, expected in zip(rows, whole):
+                self.assertEqual(row.__fields__, expected.__fields__)
+                for value, expected_value in zip(row, expected):
+                    self.assertIs(type(value), type(expected_value))
+                    if isinstance(expected_value, numpy.ndarray):
+                        self.assertEqual(value.dtype, expected_value.dtype)
+                        numpy.testing.assert_array_equal(value, expected_value)
+                    else:
+                        self.assertEqual(value, expected_value)
+
+    def test_slices_are_bounded(self):
+        table = self._table()
+        # Nested columns in one call over all rows, the rest in slices.
+        self.assertEqual(self._to_numpy_row_counts(table, 3), [7, 3, 3, 1])
+        self.assertEqual(
+            self._to_numpy_row_counts(table.select([0, 1, 2]), 3), [3, 3, 1]
+        )
+
+    def test_small_or_nested_only_tables_are_converted_whole(self):
+        table = self._table()
+        self.assertEqual(self._to_numpy_row_counts(table, 7), [7])
+        self.assertEqual(self._to_numpy_row_counts(table.select([7, 8]), 3), [7])
+
+    def test_duration_next_to_nested_column(self):
+        table = self._table().select([9, 7])
+        whole = self._convert(table, table.num_rows)
+        rows = self._convert(table, 3)
+        for row, expected in zip(rows, whole):
+            self.assertIs(type(row[0]), type(expected[0]))
+            self.assertEqual(row[0], expected[0])
+
+
 if __name__ == "__main__":
     unittest.main()