diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..fa588da2c 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 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. diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index 295f147cd..815bae7d2 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/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) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 356a96f0f..7b4de80ae 100644 --- a/src/databricks/sql/utils.py +++ b/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: diff --git a/tests/unit/test_downloader.py b/tests/unit/test_downloader.py index 00b1b849a..5334dd83a 100644 --- a/tests/unit/test_downloader.py +++ b/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)