perf(thrift): return inline Arrow result memory to the OS when freed

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.

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 = []

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)