perf(cloudfetch): cut peak memory of large CloudFetch results

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 peak memory for large Thrift CloudFetch results: downloaded files are decompressed with pyarrow into memory that is returned to the OS when the result is released.
 
 # 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/cloudfetch/downloader.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
 import logging
+import mmap
 from dataclasses import dataclass
-from typing import Optional, TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING, Union
 
 import lz4.frame
 import time
+
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
 from databricks.sql.common.http import HttpMethod
 from databricks.sql.exc import Error
 from databricks.sql.types import SSLOptions
@@ -30,12 +37,12 @@ class DownloadedFile:
     Class for the result file and metadata.
 
     Attributes:
-        file_bytes (bytes): Downloaded file in bytes.
+        file_bytes (bytes | pyarrow.Buffer): Downloaded (decompressed) Arrow IPC stream.
         start_row_offset (int): The offset of the starting row in relation to the full result.
         row_count (int): Number of rows the file represents in the result.
     """
 
-    file_bytes: bytes
+    file_bytes: Union[bytes, "pyarrow.Buffer"]
     start_row_offset: int
     row_count: int
 
@@ -119,7 +126,9 @@ def run(self) -> DownloadedFile:
         )
 
         decompressed_data = (
-            ResultSetDownloadHandler._decompress_data(compressed_data)
+            ResultSetDownloadHandler._decompress_data(
+                compressed_data, self.link.bytesNum
+            )
             if self.settings.is_lz4_compressed
             else compressed_data
         )
@@ -186,12 +195,28 @@ def _validate_link(link: TSparkArrowResultLink, expiry_buffer_secs: int):
             raise Error("CloudFetch link has expired")
 
     @staticmethod
-    def _decompress_data(compressed_data: bytes) -> bytes:
+    def _decompress_data(
+        compressed_data: bytes, expected_size: Optional[int] = None
+    ) -> Union[bytes, "pyarrow.Buffer"]:
         """
         Decompress lz4 frame compressed data.
 
         Decompresses data that has been lz4 compressed, either via the whole frame or by series of chunks.
+        When pyarrow is available and ``expected_size`` (the link's ``bytesNum``) is the decompressed
+        size, the result is a pyarrow Buffer backed by its own memory map.
         """
+        # SEA reports the compressed size in bytesNum; only a larger value can be the output size.
+        if (
+            _ARROW_LZ4_AVAILABLE
+            and expected_size is not None
+            and expected_size > len(compressed_data)
+        ):
+            mapped = ResultSetDownloadHandler._decompress_to_mmap(
+                compressed_data, expected_size
+            )
+            if mapped is not None:
+                return mapped
+
         uncompressed_data, bytes_read = lz4.frame.decompress(
             compressed_data, return_bytes_read=True
         )
@@ -208,3 +233,46 @@ def _decompress_data(compressed_data: bytes) -> bytes:
                 uncompressed_data += data
                 start += num_bytes
         return uncompressed_data
+
+    @staticmethod
+    def _decompress_to_mmap(
+        compressed_data: bytes, expected_size: int
+    ) -> Optional["pyarrow.Buffer"]:
+        """Decompress into a private anonymous memory map; return None to fall back to python-lz4.
+
+        The buffer lives as long as the Arrow table built on it. Allocated by python-lz4
+        in the download threads, such buffers fragmented the threads' malloc arenas (up to
+        ~40% extra peak RSS on large results); a dedicated mapping is returned to the OS
+        when the table is released.
+        """
+        try:
+            out = _anonymous_mmap(expected_size)
+            stream = pyarrow.CompressedInputStream(
+                pyarrow.BufferReader(compressed_data), "lz4"
+            )
+            size = stream.readinto(out)
+            if stream.read(1):
+                # expected_size was smaller than the decompressed stream
+                return None
+            return pyarrow.py_buffer(out).slice(0, size)
+        except (OSError, ValueError, pyarrow.ArrowException) as e:
+            # Let python-lz4 decode it, or raise its usual error for corrupt data
+            logger.debug("Arrow lz4 decompression failed, falling back: %s", e)
+            return None
+
+
+def _arrow_lz4_available() -> bool:
+    try:
+        return pyarrow is not None and pyarrow.Codec.is_available("lz4")
+    except Exception:
+        return False
+
+
+_ARROW_LZ4_AVAILABLE = _arrow_lz4_available()
+
+
+def _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)

src/databricks/sql/utils.py

@@ -998,13 +998,15 @@ def transform_paramstyle(
 
 
 def create_arrow_table_from_arrow_file(
-    file_bytes: bytes, description
+    file_bytes: Union[bytes, "pyarrow.Buffer"], description
 ) -> "pyarrow.Table":
     arrow_table = convert_arrow_based_file_to_arrow_table(file_bytes)
     return convert_decimals_in_arrow_table(arrow_table, description)
 
 
-def convert_arrow_based_file_to_arrow_table(file_bytes: bytes):
+def convert_arrow_based_file_to_arrow_table(
+    file_bytes: Union[bytes, "pyarrow.Buffer"],
+):
     try:
         return pyarrow.ipc.open_stream(file_bytes).read_all()
     except Exception as e:

tests/unit/test_downloader.py

@@ -1,8 +1,20 @@
 import unittest
 from unittest.mock import patch, MagicMock, Mock
+
+import lz4.frame
+import pytest
 import requests
 
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
 import databricks.sql.cloudfetch.downloader as downloader
+from databricks.sql.utils import (
+    convert_arrow_based_file_to_arrow_table,
+    create_arrow_table_from_arrow_file,
+)
 from databricks.sql.exc import Error
 from databricks.sql.types import SSLOptions
 
@@ -160,11 +172,16 @@ def test_run_compressed_successful(self, mock_time):
         result_link.fileLink = "https://s3.amazonaws.com/bucket/file.arrow?token=xyz789"
 
         # Setup mock HTTP response using helper method
-        self._setup_mock_http_response(mock_http_client, status=200, data=compressed_bytes)
+        self._setup_mock_http_response(
+            mock_http_client, status=200, data=compressed_bytes
+        )
 
         # Mock the decompression method and log metrics to avoid issues
-        with patch.object(downloader.ResultSetDownloadHandler, '_decompress_data', return_value=file_bytes), \
-             patch.object(downloader.ResultSetDownloadHandler, '_log_download_metrics'):
+        with patch.object(
+            downloader.ResultSetDownloadHandler,
+            "_decompress_data",
+            return_value=file_bytes,
+        ), patch.object(downloader.ResultSetDownloadHandler, "_log_download_metrics"):
             d = downloader.ResultSetDownloadHandler(
                 settings,
                 result_link,
@@ -178,6 +195,9 @@ def test_run_compressed_successful(self, mock_time):
             self.assertEqual(file.file_bytes, file_bytes)
             self.assertEqual(file.start_row_offset, result_link.startRowOffset)
             self.assertEqual(file.row_count, result_link.rowCount)
+            downloader.ResultSetDownloadHandler._decompress_data.assert_called_once_with(
+                compressed_bytes, len(file_bytes)
+            )
 
     @patch("time.time", return_value=1000)
     def test_download_connection_error(self, mock_time):
@@ -222,3 +242,129 @@ def test_download_timeout(self, mock_time):
         )
         with self.assertRaises(TimeoutError):
             d.run()
+
+
+def _arrow_ipc_stream(num_rows=5000):
+    table = pyarrow.table(
+        {"id": list(range(num_rows)), "s": [f"row-{i}" for i in range(num_rows)]}
+    )
+    sink = pyarrow.BufferOutputStream()
+    with pyarrow.ipc.new_stream(sink, table.schema) as writer:
+        writer.write_table(table, max_chunksize=1000)
+    return table, sink.getvalue().to_pybytes()
+
+
+@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
+class DecompressDataTests(unittest.TestCase):
+    def setUp(self):
+        self.table, self.raw = _arrow_ipc_stream()
+        self.compressed = lz4.frame.compress(self.raw)
+        self.decompress = downloader.ResultSetDownloadHandler._decompress_data
+
+    def test_expected_size_decompresses_into_private_memory_map(self):
+        with patch.object(
+            downloader.mmap, "mmap", wraps=downloader.mmap.mmap
+        ) as mmap_spy:
+            result = self.decompress(self.compressed, len(self.raw))
+
+        if hasattr(downloader.mmap, "MAP_PRIVATE"):
+            mmap_spy.assert_called_once_with(
+                -1, len(self.raw), flags=downloader.mmap.MAP_PRIVATE
+            )
+        else:
+            mmap_spy.assert_called_once_with(-1, len(self.raw))
+        self.assertIsInstance(result, pyarrow.Buffer)
+        self.assertEqual(result.to_pybytes(), self.raw)
+        self.assertTrue(
+            convert_arrow_based_file_to_arrow_table(result).equals(self.table)
+        )
+
+    @patch("time.time", return_value=1000)
+    def test_run_returns_arrow_buffer_for_lz4_file(self, mock_time):
+        http_client = MagicMock()
+        response = MagicMock(status=200, data=self.compressed)
+        http_client.request_context.return_value.__enter__.return_value = response
+        link = Mock(
+            expiryTime=1001,
+            bytesNum=len(self.raw),
+            startRowOffset=0,
+            rowCount=self.table.num_rows,
+            fileLink="https://example.com/file?sig=x",
+            httpHeaders={},
+        )
+        handler = downloader.ResultSetDownloadHandler(
+            downloader.DownloadableResultSettings(is_lz4_compressed=True),
+            link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=None,
+            statement_id="s",
+            http_client=http_client,
+        )
+        with patch.object(downloader.ResultSetDownloadHandler, "_log_download_metrics"):
+            file = handler.run()
+
+        self.assertIsInstance(file.file_bytes, pyarrow.Buffer)
+        description = [
+            ("id", "bigint", None, None, None, None, None),
+            ("s", "string", None, None, None, None, None),
+        ]
+        table = create_arrow_table_from_arrow_file(file.file_bytes, description)
+        self.assertTrue(table.equals(self.table))
+
+    def test_mmap_failure_falls_back_to_python_lz4(self):
+        with patch.object(downloader.mmap, "mmap", side_effect=OSError("ENOMEM")):
+            result = self.decompress(self.compressed, len(self.raw))
+
+        self.assertIsInstance(result, bytes)
+        self.assertEqual(result, self.raw)
+
+    def test_oversized_expected_size_returns_exact_length(self):
+        result = self.decompress(self.compressed, len(self.raw) + 4096)
+
+        self.assertIsInstance(result, pyarrow.Buffer)
+        self.assertEqual(result.to_pybytes(), self.raw)
+
+    def test_undersized_expected_size_falls_back_to_python_lz4(self):
+        result = self.decompress(self.compressed, len(self.raw) - 1)
+
+        self.assertIsInstance(result, bytes)
+        self.assertEqual(result, self.raw)
+
+    def test_unknown_expected_size_uses_python_lz4(self):
+        # SEA links carry the compressed size, which is not a usable output size
+        for expected_size in (None, 0, len(self.compressed)):
+            with patch.object(downloader.mmap, "mmap") as mmap_mock:
+                result = self.decompress(self.compressed, expected_size)
+            mmap_mock.assert_not_called()
+            self.assertIsInstance(result, bytes)
+            self.assertEqual(result, self.raw)
+
+    def test_concatenated_frames(self):
+        middle = len(self.raw) // 2
+        compressed = lz4.frame.compress(self.raw[:middle]) + lz4.frame.compress(
+            self.raw[middle:]
+        )
+        for expected_size in (len(self.raw), None):
+            result = self.decompress(compressed, expected_size)
+            self.assertEqual(bytes(result), self.raw)
+
+    def test_corrupt_data_raises_python_lz4_error(self):
+        # Databricks frames carry block and content checksums
+        compressed = lz4.frame.compress(
+            self.raw, block_checksum=True, content_checksum=True
+        )
+        corrupt = bytearray(compressed)
+        corrupt[len(corrupt) // 2] ^= 0xFF
+        truncated = compressed[:-10]
+        for data in (bytes(corrupt), truncated):
+            for expected_size in (len(self.raw), None):
+                with self.assertRaises(RuntimeError):
+                    self.decompress(data, expected_size)
+
+    def test_without_arrow_lz4_codec_uses_python_lz4(self):
+        with patch.object(downloader, "_ARROW_LZ4_AVAILABLE", False):
+            result = self.decompress(self.compressed, len(self.raw))
+
+        self.assertIsInstance(result, bytes)
+        self.assertEqual(result, self.raw)