From 60cf99d2bfcda4355173be8a5b5ba0ba70cef936 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 09:13:27 +0000 Subject: [PATCH] perf: convert large Arrow results to rows in slices fetchall() and large fetchmany() converted the whole Arrow table with one to_pandas() call and one object ndarray, so a DataFrame, the object array and the final Row list were all alive at the peak. Convert 100,000 rows at a time instead. Columns whose pandas conversion depends on the column's other values (list/map/struct: pyarrow picks the element dtype from the whole column) are still converted over the whole table, so every value and type is unchanged. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/result_set.py | 58 ++++++++++++-- tests/unit/test_pandas_compatibility.py | 101 +++++++++++++++++++++++- 3 files changed, 152 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..e54a0a71c 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. +- 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. diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index f24a07505..114179336 100644 --- a/src/databricks/sql/result_set.py +++ b/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 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): diff --git a/tests/unit/test_pandas_compatibility.py b/tests/unit/test_pandas_compatibility.py index 5434559b2..82bdba237 100644 --- a/tests/unit/test_pandas_compatibility.py +++ b/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()