From 408907592caa31636238b5fc65fffc9a44cfe318 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Thu, 12 Nov 2026 10:08:50 +0000
Subject: [PATCH 1/6] perf(thrift): return inline Arrow result memory to the OS
 when freed

Each inline FetchResults row set was decompressed into a bytearray grown
on the malloc heap. The Arrow table reads it zero-copy, so ~3.7 MB per
round trip stayed live there, mixed with short-lived reply buffers; after
a large result was released glibc kept almost all of it resident (a
1.7 GB inline fetchall_arrow left ~1.6 GB). Copy each set's stream into
a private anonymous mmap instead (pre-faulted with MAP_POPULATE where
available), which goes back to the OS when the table is freed. Sets under
128 KiB and mmap failures keep the heap path.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
---
 CHANGELOG.md                |  1 +
 src/databricks/sql/utils.py | 45 ++++++++++++++++--
 tests/unit/test_util.py     | 92 +++++++++++++++++++++++++++++++++++++
 3 files changed, 134 insertions(+), 4 deletions(-)

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..16cdb2295 100644
--- a/src/databricks/sql/utils.py
+++ b/src/databricks/sql/utils.py
@@ -20,6 +20,7 @@
 from collections.abc import Mapping, Set as AbstractSet
 from decimal import Decimal
 from enum import Enum
+import mmap
 import re
 import secrets
 
@@ -1012,20 +1013,56 @@ 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
+    parts = [schema_bytes]
     n_rows = 0
     for arrow_batch in arrow_batches:
         n_rows += arrow_batch.rowCount
-        ba += (
+        parts.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(_join_arrow_stream(parts)).read_all()
     return arrow_table, n_rows
 
 
+# Below glibc's default mmap threshold a dedicated mapping only adds a syscall.
+_INLINE_ARROW_MMAP_MIN_BYTES = 128 * 1024
+
+
+def _join_arrow_stream(parts: List[bytes]):
+    """Concatenate an inline Arrow IPC stream into its own memory map.
+
+    The table reads the stream zero-copy, so this buffer lives as long as the result.
+    On the malloc heap such buffers often stayed resident after the result was freed
+    (~1 GB per 1 GB of inline results); a private mapping is returned to the OS.
+    """
+    size = sum(len(part) for part in parts)
+    if size >= _INLINE_ARROW_MMAP_MIN_BYTES:
+        try:
+            out = _populated_private_mmap(size)
+        except (OSError, ValueError):
+            pass
+        else:
+            view = memoryview(out)
+            pos = 0
+            for part in parts:
+                view[pos : pos + len(part)] = part
+                pos += len(part)
+            view.release()
+            return pyarrow.py_buffer(out)
+    return b"".join(parts)
+
+
+def _populated_private_mmap(size: int) -> mmap.mmap:
+    # Python defaults anonymous maps to MAP_SHARED (shmem, shared with forked children).
+    # The whole map is written right away, so pre-faulting it is cheaper.
+    if hasattr(mmap, "MAP_PRIVATE"):
+        flags = mmap.MAP_PRIVATE | getattr(mmap, "MAP_POPULATE", 0)
+        return mmap.mmap(-1, size, flags=flags)
+    return mmap.mmap(-1, 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..a8fe247d6 100644
--- a/tests/unit/test_util.py
+++ b/tests/unit/test_util.py
@@ -1,8 +1,17 @@
 import decimal
 import datetime
+import gc
+import mmap
 from datetime import timezone, timedelta
+from types import SimpleNamespace
+from unittest.mock import 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 +272,86 @@ 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
+
+
+@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
+class TestConvertArrowBasedSet:
+    @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, 5000, 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
+
+    def test_table_keeps_the_mapping_alive(self):
+        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+
+        result, _ = convert_arrow_based_set_to_arrow_table(batches, True, schema_bytes)
+        del batches, schema_bytes
+        gc.collect()
+
+        assert result.equals(table)
+
+    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, 5000, 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)

From 98c65bb0bf35bb285d6f14d9a92752d567e7fd5d Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Thu, 12 Nov 2026 10:31:09 +0000
Subject: [PATCH 2/6] fix(thrift): fill the inline Arrow map batch by batch

Holding every decompressed batch until the copy needed about twice the
row set's memory, and releasing the parts did not help because glibc
keeps freed chunks resident. Reserve the map from the compressed size
instead (untouched pages cost nothing), copy each batch in as it is
decompressed, and grow the map with mremap on overflow. Platforms
without mremap continue on the heap. A 51 MB row set now converts with
no transient, as on main.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
---
 src/databricks/sql/utils.py | 69 ++++++++++++++++++++++++-------------
 tests/unit/test_util.py     | 29 ++++++++++++++++
 2 files changed, 74 insertions(+), 24 deletions(-)

diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
index 16cdb2295..34a921826 100644
--- a/src/databricks/sql/utils.py
+++ b/src/databricks/sql/utils.py
@@ -1013,53 +1013,74 @@ 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):
-    parts = [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
-        parts.append(
+        stream.append(
             lz4.frame.decompress(arrow_batch.batch)
             if lz4_compressed
             else arrow_batch.batch
         )
-    arrow_table = pyarrow.ipc.open_stream(_join_arrow_stream(parts)).read_all()
+    arrow_table = pyarrow.ipc.open_stream(stream.getbuffer()).read_all()
     return arrow_table, n_rows
 
 
 # Below glibc's default mmap threshold a dedicated mapping only adds a syscall.
 _INLINE_ARROW_MMAP_MIN_BYTES = 128 * 1024
+# Untouched pages of the reservation cost no memory; overflow grows the map.
+_LZ4_RATIO_ESTIMATE = 4
 
 
-def _join_arrow_stream(parts: List[bytes]):
-    """Concatenate an inline Arrow IPC stream into its own memory map.
+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 malloc heap such buffers often stayed resident after the result was freed
-    (~1 GB per 1 GB of inline results); a private mapping is returned to the OS.
+    (~1 GB per 1 GB of inline results); a mapping is returned to the OS. Small streams
+    and mmap failures use a bytearray, as before.
     """
-    size = sum(len(part) for part in parts)
-    if size >= _INLINE_ARROW_MMAP_MIN_BYTES:
-        try:
-            out = _populated_private_mmap(size)
-        except (OSError, ValueError):
-            pass
+
+    def __init__(self, size_estimate: int):
+        self._map: Optional[mmap.mmap] = None
+        self._bytes = bytearray()
+        self._size = 0
+        if size_estimate >= _INLINE_ARROW_MMAP_MIN_BYTES:
+            try:
+                self._map = _private_anonymous_mmap(size_estimate)
+            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):
+            try:
+                self._map.resize(max(end, 2 * len(self._map)))
+            except (OSError, SystemError, ValueError):
+                # No mremap (e.g. macOS): continue on the heap
+                self._bytes += self._map[: self._size]
+                self._map = None
+        if self._map is not None:
+            self._map[self._size : end] = data
         else:
-            view = memoryview(out)
-            pos = 0
-            for part in parts:
-                view[pos : pos + len(part)] = part
-                pos += len(part)
-            view.release()
-            return pyarrow.py_buffer(out)
-    return b"".join(parts)
+            self._bytes += data
+        self._size = end
+
+    def getbuffer(self):
+        if self._map is None:
+            return self._bytes
+        return pyarrow.py_buffer(self._map).slice(0, self._size)
 
 
-def _populated_private_mmap(size: int) -> mmap.mmap:
+def _private_anonymous_mmap(size: int) -> mmap.mmap:
     # Python defaults anonymous maps to MAP_SHARED (shmem, shared with forked children).
-    # The whole map is written right away, so pre-faulting it is cheaper.
     if hasattr(mmap, "MAP_PRIVATE"):
-        flags = mmap.MAP_PRIVATE | getattr(mmap, "MAP_POPULATE", 0)
-        return mmap.mmap(-1, size, flags=flags)
+        return mmap.mmap(-1, size, flags=mmap.MAP_PRIVATE)
     return mmap.mmap(-1, size)
 
 
diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py
index a8fe247d6..785171c4a 100644
--- a/tests/unit/test_util.py
+++ b/tests/unit/test_util.py
@@ -315,6 +315,35 @@ def test_large_set_is_read_from_a_private_memory_map(self, lz4_compressed):
             assert flags & mmap.MAP_PRIVATE
             assert not flags & mmap.MAP_SHARED
 
+    def test_map_grows_when_the_size_estimate_is_too_small(self):
+        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+
+        with patch.object(utils, "_LZ4_RATIO_ESTIMATE", 1), patch.object(
+            utils.mmap, "mmap", wraps=mmap.mmap
+        ) as mmap_spy:
+            result, _ = convert_arrow_based_set_to_arrow_table(
+                batches, True, schema_bytes
+            )
+
+        mmap_spy.assert_called_once()
+        assert result.equals(table)
+
+    def test_falls_back_to_the_heap_when_the_map_cannot_grow(self):
+        class NoResize(mmap.mmap):
+            def resize(self, newsize):
+                raise SystemError("mmap: resizing not available--no mremap()")
+
+        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+
+        with patch.object(utils, "_LZ4_RATIO_ESTIMATE", 1), patch.object(
+            utils.mmap, "mmap", NoResize
+        ):
+            result, _ = convert_arrow_based_set_to_arrow_table(
+                batches, True, schema_bytes
+            )
+
+        assert result.equals(table)
+
     def test_table_keeps_the_mapping_alive(self):
         table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
 

From ac8fb875a700b9a1cd2c04f7f9ccb7a0632b3059 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Thu, 12 Nov 2026 11:14:50 +0000
Subject: [PATCH 3/6] fix(thrift): shrink the inline Arrow map; keep non-POSIX
 on the heap

Release the unused part of the size estimate once the stream is
complete, so it no longer counts against a strict overcommit limit
while the result is alive. Use the map only where MAP_PRIVATE exists:
the heap retention is a glibc behaviour, and Windows commits anonymous
maps up front and copies on resize. Tests now assert the growth,
shrink, and heap-fallback paths are actually taken.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
---
 src/databricks/sql/utils.py | 26 ++++++++--------
 tests/unit/test_util.py     | 61 ++++++++++++++++++++++++++++++++++---
 2 files changed, 70 insertions(+), 17 deletions(-)

diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
index 34a921826..061eeb21f 100644
--- a/src/databricks/sql/utils.py
+++ b/src/databricks/sql/utils.py
@@ -1041,18 +1041,21 @@ 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 malloc 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
-    and mmap failures use a bytearray, as before.
+    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-POSIX 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 size_estimate >= _INLINE_ARROW_MMAP_MIN_BYTES:
+        if size_estimate >= _INLINE_ARROW_MMAP_MIN_BYTES and hasattr(
+            mmap, "MAP_PRIVATE"
+        ):
             try:
-                self._map = _private_anonymous_mmap(size_estimate)
+                # 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
 
@@ -1074,16 +1077,15 @@ def append(self, data: bytes) -> None:
     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 _private_anonymous_mmap(size: int) -> mmap.mmap:
-    # Python defaults anonymous maps to MAP_SHARED (shmem, shared with forked children).
-    if hasattr(mmap, "MAP_PRIVATE"):
-        return mmap.mmap(-1, size, flags=mmap.MAP_PRIVATE)
-    return mmap.mmap(-1, 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 785171c4a..0f7b3dc3e 100644
--- a/tests/unit/test_util.py
+++ b/tests/unit/test_util.py
@@ -4,7 +4,7 @@
 import mmap
 from datetime import timezone, timedelta
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import lz4.frame
 import pytest
@@ -294,6 +294,18 @@ def _inline_arrow_set(num_batches, rows_per_batch, lz4_compressed):
     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)
+
+
+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:
     @pytest.mark.parametrize("lz4_compressed", [True, False])
@@ -317,31 +329,70 @@ def test_large_set_is_read_from_a_private_memory_map(self, lz4_compressed):
 
     def test_map_grows_when_the_size_estimate_is_too_small(self):
         table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        _RecordingMap.resizes = []
 
         with patch.object(utils, "_LZ4_RATIO_ESTIMATE", 1), patch.object(
-            utils.mmap, "mmap", wraps=mmap.mmap
-        ) as mmap_spy:
+            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)
+
+    def test_map_is_shrunk_to_the_stream_size(self):
+        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        _RecordingMap.resizes = []
+
+        with patch.object(utils.mmap, "mmap", _RecordingMap):
             result, _ = convert_arrow_based_set_to_arrow_table(
                 batches, True, schema_bytes
             )
 
-        mmap_spy.assert_called_once()
         assert result.equals(table)
+        assert len(_RecordingMap.resizes) == 1
+        old, new = _RecordingMap.resizes[0]
+        assert new == _stream_size(schema_bytes, batches) < old
 
     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, 5000, 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, 5000, 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_table_keeps_the_mapping_alive(self):

From 7f874a7e1364159e93e7219354fc5d2c7dd8b4ea Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Thu, 12 Nov 2026 11:19:48 +0000
Subject: [PATCH 4/6] fix(thrift): map inline Arrow streams on Linux only; copy
 fallback via memoryview

The heap retention this works around is glibc's, so other platforms keep
the previous bytearray. When a map can't grow, copy what was written via
a memoryview instead of a temporary bytes slice.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
---
 src/databricks/sql/utils.py | 16 +++++++++++-----
 tests/unit/test_util.py     | 13 +++++++++++++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
index 061eeb21f..574fb4f75 100644
--- a/src/databricks/sql/utils.py
+++ b/src/databricks/sql/utils.py
@@ -23,6 +23,7 @@
 import mmap
 import re
 import secrets
+import sys
 
 import lz4.frame
 
@@ -1035,6 +1036,8 @@ def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema
 _INLINE_ARROW_MMAP_MIN_BYTES = 128 * 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:
@@ -1043,15 +1046,17 @@ class _ArrowStreamBuffer:
     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-POSIX platforms use a bytearray, as before.
+    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 size_estimate >= _INLINE_ARROW_MMAP_MIN_BYTES and hasattr(
-            mmap, "MAP_PRIVATE"
+        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).
@@ -1065,8 +1070,9 @@ def append(self, data: bytes) -> None:
             try:
                 self._map.resize(max(end, 2 * len(self._map)))
             except (OSError, SystemError, ValueError):
-                # No mremap (e.g. macOS): continue on the heap
-                self._bytes += self._map[: self._size]
+                # 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
diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py
index 0f7b3dc3e..1ddeacd90 100644
--- a/tests/unit/test_util.py
+++ b/tests/unit/test_util.py
@@ -395,6 +395,19 @@ def test_non_posix_platform_stays_on_the_heap(self):
         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, 5000, 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)
+
     def test_table_keeps_the_mapping_alive(self):
         table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
 

From 2b694a52902512b4fcec6af8d0b4724190eb38a3 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Thu, 12 Nov 2026 12:02:59 +0000
Subject: [PATCH 5/6] fix(thrift): retry the inline map grow at the exact size;
 tighten tests

If doubling the map fails, retry with just the needed size before
falling back to the heap. Mark the map-path tests Linux-only, and make
the keep-alive test hold a weakref to the map so it checks that the
table owns it and that it is released with the table.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
---
 src/databricks/sql/utils.py | 10 +++++--
 tests/unit/test_util.py     | 60 +++++++++++++++++++++++++++++++++++--
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
index 574fb4f75..2833fb3c5 100644
--- a/src/databricks/sql/utils.py
+++ b/src/databricks/sql/utils.py
@@ -1067,9 +1067,13 @@ def __init__(self, size_estimate: int):
     def append(self, data: bytes) -> None:
         end = self._size + len(data)
         if self._map is not None and end > len(self._map):
-            try:
-                self._map.resize(max(end, 2 * len(self._map)))
-            except (OSError, SystemError, ValueError):
+            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]
diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py
index 1ddeacd90..e9b24a1d7 100644
--- a/tests/unit/test_util.py
+++ b/tests/unit/test_util.py
@@ -2,6 +2,8 @@
 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
@@ -298,6 +300,11 @@ 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 = []
 
@@ -308,6 +315,7 @@ def resize(self, 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, 5000, lz4_compressed)
@@ -327,6 +335,7 @@ def test_large_set_is_read_from_a_private_memory_map(self, lz4_compressed):
             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, 5000, True)
         _RecordingMap.resizes = []
@@ -342,6 +351,7 @@ def test_map_grows_when_the_size_estimate_is_too_small(self):
         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, 5000, True)
         _RecordingMap.resizes = []
@@ -356,6 +366,7 @@ def test_map_is_shrunk_to_the_stream_size(self):
         old, new = _RecordingMap.resizes[0]
         assert new == _stream_size(schema_bytes, batches) < old
 
+    @linux_only
     def test_falls_back_to_the_heap_when_the_map_cannot_grow(self):
         attempts = []
 
@@ -408,14 +419,57 @@ def test_non_linux_platform_stays_on_the_heap(self):
         mmap_mock.assert_not_called()
         assert result.equals(table)
 
-    def test_table_keeps_the_mapping_alive(self):
+    @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, 5000, True)
+        states = []
+        getbuffer = utils._ArrowStreamBuffer.getbuffer
+
+        def recording_getbuffer(stream):
+            states.append(stream._map is not None)
+            return getbuffer(stream)
 
-        result, _ = convert_arrow_based_set_to_arrow_table(batches, True, schema_bytes)
-        del batches, schema_bytes
+        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, 5000, 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)

From 46537ade40e1be543c08074d1962560d0fbe7af3 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Thu, 12 Nov 2026 12:10:50 +0000
Subject: [PATCH 6/6] fix(thrift): keep inline Arrow streams under 1 MiB on the
 heap

Each held row-set map is its own VMA (the final shrink leaves a gap, so
neighbours don't merge), counted against vm.max_map_count. Raising the
floor from 128 KiB to 1 MiB means reaching the default limit takes over
64 GB of inline results held at once. Add a test for an unshrinkable
map that checks the stream is sliced to its exact size.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
---
 src/databricks/sql/utils.py |  5 ++--
 tests/unit/test_util.py     | 46 +++++++++++++++++++++++++++++--------
 2 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
index 2833fb3c5..17d9dc995 100644
--- a/src/databricks/sql/utils.py
+++ b/src/databricks/sql/utils.py
@@ -1032,8 +1032,9 @@ def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema
     return arrow_table, n_rows
 
 
-# Below glibc's default mmap threshold a dedicated mapping only adds a syscall.
-_INLINE_ARROW_MMAP_MIN_BYTES = 128 * 1024
+# 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.
diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py
index e9b24a1d7..416065bcf 100644
--- a/tests/unit/test_util.py
+++ b/tests/unit/test_util.py
@@ -318,7 +318,7 @@ 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, 5000, 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(
@@ -337,7 +337,7 @@ def test_large_set_is_read_from_a_private_memory_map(self, lz4_compressed):
 
     @linux_only
     def test_map_grows_when_the_size_estimate_is_too_small(self):
-        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        table, schema_bytes, batches = _inline_arrow_set(4, 50000, True)
         _RecordingMap.resizes = []
 
         with patch.object(utils, "_LZ4_RATIO_ESTIMATE", 1), patch.object(
@@ -353,7 +353,7 @@ def test_map_grows_when_the_size_estimate_is_too_small(self):
 
     @linux_only
     def test_map_is_shrunk_to_the_stream_size(self):
-        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        table, schema_bytes, batches = _inline_arrow_set(4, 50000, True)
         _RecordingMap.resizes = []
 
         with patch.object(utils.mmap, "mmap", _RecordingMap):
@@ -366,6 +366,34 @@ def test_map_is_shrunk_to_the_stream_size(self):
         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 = []
@@ -375,7 +403,7 @@ 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, 5000, True)
+        table, schema_bytes, batches = _inline_arrow_set(4, 50000, True)
         states = []
         getbuffer = utils._ArrowStreamBuffer.getbuffer
 
@@ -395,7 +423,7 @@ def recording_getbuffer(stream):
         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, 5000, True)
+        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):
@@ -407,7 +435,7 @@ def test_non_posix_platform_stays_on_the_heap(self):
         assert result.equals(table)
 
     def test_non_linux_platform_stays_on_the_heap(self):
-        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        table, schema_bytes, batches = _inline_arrow_set(4, 50000, True)
 
         with patch.object(utils, "_INLINE_ARROW_USE_MMAP", False), patch.object(
             utils.mmap, "mmap"
@@ -430,7 +458,7 @@ def resize(self, newsize):
                     raise OSError("ENOMEM")
                 return super().resize(newsize)
 
-        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        table, schema_bytes, batches = _inline_arrow_set(4, 50000, True)
         states = []
         getbuffer = utils._ArrowStreamBuffer.getbuffer
 
@@ -457,7 +485,7 @@ class TrackedMap(mmap.mmap):
             def __init__(self, *args, **kwargs):
                 maps.append(weakref.ref(self))
 
-        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        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(
@@ -484,7 +512,7 @@ def test_small_set_stays_on_the_heap(self):
         assert result.equals(table)
 
     def test_mmap_failure_falls_back_to_bytes(self):
-        table, schema_bytes, batches = _inline_arrow_set(4, 5000, True)
+        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(