diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..a90e4ba40 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. +- Return memory to the OS when large inline Arrow results on Thrift (for example with `use_cloud_fetch=False`) are released, instead of keeping it resident in the process. # 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/utils.py b/src/databricks/sql/utils.py index 356a96f0f..17d9dc995 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -20,8 +20,10 @@ from collections.abc import Mapping, Set as AbstractSet from decimal import Decimal from enum import Enum +import mmap import re import secrets +import sys import lz4.frame @@ -1012,20 +1014,89 @@ def convert_arrow_based_file_to_arrow_table(file_bytes: bytes): def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema_bytes): - ba = bytearray() - ba += schema_bytes + compressed_size = sum(len(arrow_batch.batch) for arrow_batch in arrow_batches) + stream = _ArrowStreamBuffer( + len(schema_bytes) + + compressed_size * (_LZ4_RATIO_ESTIMATE if lz4_compressed else 1) + ) + stream.append(schema_bytes) n_rows = 0 for arrow_batch in arrow_batches: n_rows += arrow_batch.rowCount - ba += ( + stream.append( lz4.frame.decompress(arrow_batch.batch) if lz4_compressed else arrow_batch.batch ) - arrow_table = pyarrow.ipc.open_stream(ba).read_all() + arrow_table = pyarrow.ipc.open_stream(stream.getbuffer()).read_all() return arrow_table, n_rows +# Each live map is one VMA (vm.max_map_count, 65530 by default), so smaller streams +# stay on the heap: reaching that limit takes over 64 GB of inline results held at once. +_INLINE_ARROW_MMAP_MIN_BYTES = 1024 * 1024 +# Untouched pages of the reservation cost no memory; overflow grows the map. +_LZ4_RATIO_ESTIMATE = 4 +# The heap retention this works around was measured with glibc. +_INLINE_ARROW_USE_MMAP = sys.platform.startswith("linux") + + +class _ArrowStreamBuffer: + """Accumulates an inline Arrow IPC stream in a private anonymous memory map. + + The table reads the stream zero-copy, so this buffer lives as long as the result. + On the glibc heap such buffers often stayed resident after the result was freed + (~1 GB per 1 GB of inline results); a mapping is returned to the OS. Small streams, + mmap failures and non-Linux platforms use a bytearray, as before. + """ + + def __init__(self, size_estimate: int): + self._map: Optional[mmap.mmap] = None + self._bytes = bytearray() + self._size = 0 + if ( + _INLINE_ARROW_USE_MMAP + and size_estimate >= _INLINE_ARROW_MMAP_MIN_BYTES + and hasattr(mmap, "MAP_PRIVATE") + ): + try: + # Python defaults anonymous maps to MAP_SHARED (shmem, shared on fork). + self._map = mmap.mmap(-1, size_estimate, flags=mmap.MAP_PRIVATE) + except (OSError, ValueError): + pass + + def append(self, data: bytes) -> None: + end = self._size + len(data) + if self._map is not None and end > len(self._map): + for new_size in (max(end, 2 * len(self._map)), end): + try: + self._map.resize(new_size) + break + except (OSError, SystemError, ValueError): + pass + else: + # The map can't grow: continue on the heap + with memoryview(self._map) as view: + self._bytes += view[: self._size] + self._map = None + if self._map is not None: + self._map[self._size : end] = data + else: + self._bytes += data + self._size = end + + def getbuffer(self): + if self._map is None: + return self._bytes + if len(self._map) > self._size: + # Release the unused reservation (commit charge under strict overcommit) + try: + self._map.resize(self._size) + except (OSError, SystemError, ValueError): + pass + return pyarrow.py_buffer(self._map).slice(0, self._size) + + def convert_decimals_in_arrow_table(table, description) -> "pyarrow.Table": new_columns = [] new_fields = [] diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py index efed8be53..416065bcf 100644 --- a/tests/unit/test_util.py +++ b/tests/unit/test_util.py @@ -1,8 +1,19 @@ import decimal import datetime +import gc +import mmap +import sys +import weakref from datetime import timezone, timedelta +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import lz4.frame import pytest + +import databricks.sql.utils as utils from databricks.sql.utils import ( + convert_arrow_based_set_to_arrow_table, convert_to_assigned_datatypes_in_column_table, ColumnTable, concat_table_chunks, @@ -263,3 +274,259 @@ def test_serialize_query_tags_all_none_values(self): query_tags = {"key1": None, "key2": None, "key3": None} result = serialize_query_tags(query_tags) assert result == "key1,key2,key3" + + +def _inline_arrow_set(num_batches, rows_per_batch, lz4_compressed): + table = pyarrow.table( + { + "id": list(range(num_batches * rows_per_batch)), + "s": [f"row-{i}" for i in range(num_batches * rows_per_batch)], + } + ) + schema_bytes = table.schema.serialize().to_pybytes() + batches = [] + for batch in table.to_batches(max_chunksize=rows_per_batch): + data = batch.serialize().to_pybytes() + batches.append( + SimpleNamespace( + batch=lz4.frame.compress(data) if lz4_compressed else data, + rowCount=batch.num_rows, + ) + ) + return table, schema_bytes, batches + + +def _stream_size(schema_bytes, batches): + return len(schema_bytes) + sum(len(lz4.frame.decompress(b.batch)) for b in batches) + + +linux_only = pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="inline Arrow maps are Linux-only" +) + + +class _RecordingMap(mmap.mmap): + resizes = [] + + def resize(self, newsize): + type(self).resizes.append((len(self), newsize)) + return super().resize(newsize) + + +@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") +class TestConvertArrowBasedSet: + @linux_only + @pytest.mark.parametrize("lz4_compressed", [True, False]) + def test_large_set_is_read_from_a_private_memory_map(self, lz4_compressed): + table, schema_bytes, batches = _inline_arrow_set(4, 50000, lz4_compressed) + + with patch.object(utils.mmap, "mmap", wraps=mmap.mmap) as mmap_spy: + result, n_rows = convert_arrow_based_set_to_arrow_table( + batches, lz4_compressed, schema_bytes + ) + + assert n_rows == table.num_rows + assert result.equals(table) + mmap_spy.assert_called_once() + size = mmap_spy.call_args.args[1] + assert size >= utils._INLINE_ARROW_MMAP_MIN_BYTES + if hasattr(mmap, "MAP_PRIVATE"): + flags = mmap_spy.call_args.kwargs["flags"] + assert flags & mmap.MAP_PRIVATE + assert not flags & mmap.MAP_SHARED + + @linux_only + def test_map_grows_when_the_size_estimate_is_too_small(self): + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + _RecordingMap.resizes = [] + + with patch.object(utils, "_LZ4_RATIO_ESTIMATE", 1), patch.object( + utils.mmap, "mmap", _RecordingMap + ): + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + assert result.equals(table) + assert any(new > old for old, new in _RecordingMap.resizes) + assert _RecordingMap.resizes[-1][1] == _stream_size(schema_bytes, batches) + + @linux_only + def test_map_is_shrunk_to_the_stream_size(self): + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + _RecordingMap.resizes = [] + + with patch.object(utils.mmap, "mmap", _RecordingMap): + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + assert result.equals(table) + assert len(_RecordingMap.resizes) == 1 + old, new = _RecordingMap.resizes[0] + assert new == _stream_size(schema_bytes, batches) < old + + @linux_only + def test_unshrinkable_map_is_sliced_to_the_stream(self): + class NoShrink(mmap.mmap): + def resize(self, newsize): + if newsize < len(self): + raise SystemError("mmap: resizing not available--no mremap()") + return super().resize(newsize) + + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + buffers = [] + getbuffer = utils._ArrowStreamBuffer.getbuffer + + def recording_getbuffer(stream): + buffer = getbuffer(stream) + buffers.append((len(stream._map), len(buffer))) + return buffer + + with patch.object(utils.mmap, "mmap", NoShrink), patch.object( + utils._ArrowStreamBuffer, "getbuffer", recording_getbuffer + ): + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + map_size, buffer_size = buffers[0] + assert map_size > buffer_size == _stream_size(schema_bytes, batches) + assert result.equals(table) + + @linux_only + def test_falls_back_to_the_heap_when_the_map_cannot_grow(self): + attempts = [] + + class NoResize(mmap.mmap): + def resize(self, newsize): + attempts.append((len(self), newsize)) + raise SystemError("mmap: resizing not available--no mremap()") + + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + states = [] + getbuffer = utils._ArrowStreamBuffer.getbuffer + + def recording_getbuffer(stream): + states.append((stream._map, len(stream._bytes))) + return getbuffer(stream) + + with patch.object(utils, "_LZ4_RATIO_ESTIMATE", 1), patch.object( + utils.mmap, "mmap", NoResize + ), patch.object(utils._ArrowStreamBuffer, "getbuffer", recording_getbuffer): + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + assert result.equals(table) + assert attempts and attempts[0][1] > attempts[0][0] + assert states == [(None, _stream_size(schema_bytes, batches))] + + def test_non_posix_platform_stays_on_the_heap(self): + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + no_map_private = SimpleNamespace(mmap=Mock(side_effect=AssertionError)) + + with patch.object(utils, "mmap", no_map_private): + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + no_map_private.mmap.assert_not_called() + assert result.equals(table) + + def test_non_linux_platform_stays_on_the_heap(self): + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + + with patch.object(utils, "_INLINE_ARROW_USE_MMAP", False), patch.object( + utils.mmap, "mmap" + ) as mmap_mock: + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + mmap_mock.assert_not_called() + assert result.equals(table) + + @linux_only + def test_map_grows_to_the_exact_size_when_doubling_fails(self): + sizes = [] + + class NoDoubling(mmap.mmap): + def resize(self, newsize): + sizes.append(newsize) + if len(sizes) == 1: + raise OSError("ENOMEM") + return super().resize(newsize) + + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + states = [] + getbuffer = utils._ArrowStreamBuffer.getbuffer + + def recording_getbuffer(stream): + states.append(stream._map is not None) + return getbuffer(stream) + + with patch.object(utils, "_LZ4_RATIO_ESTIMATE", 1), patch.object( + utils.mmap, "mmap", NoDoubling + ), patch.object(utils._ArrowStreamBuffer, "getbuffer", recording_getbuffer): + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + assert result.equals(table) + assert sizes[1] < sizes[0] + assert states == [True] + + @linux_only + def test_table_keeps_the_map_alive_and_frees_it_with_the_table(self): + maps = [] + + class TrackedMap(mmap.mmap): + def __init__(self, *args, **kwargs): + maps.append(weakref.ref(self)) + + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + + with patch.object(utils.mmap, "mmap", TrackedMap): + result, _ = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + gc.collect() + + assert len(maps) == 1 and maps[0]() is not None + assert result.equals(table) + del result + gc.collect() + assert maps[0]() is None + + def test_small_set_stays_on_the_heap(self): + table, schema_bytes, batches = _inline_arrow_set(2, 10, True) + + with patch.object(utils.mmap, "mmap") as mmap_mock: + result, n_rows = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + mmap_mock.assert_not_called() + assert n_rows == 20 + assert result.equals(table) + + def test_mmap_failure_falls_back_to_bytes(self): + table, schema_bytes, batches = _inline_arrow_set(4, 50000, True) + + with patch.object(utils.mmap, "mmap", side_effect=OSError("ENOMEM")): + result, n_rows = convert_arrow_based_set_to_arrow_table( + batches, True, schema_bytes + ) + + assert n_rows == table.num_rows + assert result.equals(table) + + def test_schema_only_set(self): + table, schema_bytes, _ = _inline_arrow_set(1, 1, True) + + result, n_rows = convert_arrow_based_set_to_arrow_table([], True, schema_bytes) + + assert n_rows == 0 + assert result.num_rows == 0 + assert result.schema.equals(table.schema)