From bf8752c09164d4ef30885cf9dd884f593c3ab58c Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 08:45:12 +0000 Subject: [PATCH 1/4] perf(thrift): decompress many-batch inline Arrow results in one call The warehouse returns small inline results as dozens of LZ4 frames (72 for both 100 and 1,000 rows of range()). Decompressing each frame with lz4.frame.decompress releases and reacquires the GIL per frame, so with many busy threads a 1,000-row statement made ~600 context switches and spent most of its time waiting for the GIL. For three or more compressed batches, decode the concatenated frames with one pyarrow CompressedInputStream read. One- and two-batch results keep the python-lz4 path, as does any pyarrow build without the lz4 codec. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + src/databricks/sql/utils.py | 36 +++++++++-- tests/unit/test_thrift_backend.py | 103 ++++++++++++++++++++---------- 3 files changed, 100 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..9c66f4ded 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. +- Speed up multi-threaded clients that fetch small inline results on the Thrift backend: LZ4-compressed results split into many Arrow batches are now decompressed in one call instead of one GIL-releasing call per batch, roughly doubling statements per second for 1,000-row results at 64 threads. # 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..65601557d 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -1011,17 +1011,39 @@ def convert_arrow_based_file_to_arrow_table(file_bytes: bytes): raise RuntimeError("Failure to convert arrow based file to arrow table", e) +_PYARROW_HAS_LZ4 = pyarrow is not None and pyarrow.Codec.is_available("lz4") +_LZ4_MAX_READ_BYTES = 16 * 1024 * 1024 + + +def _decompress_lz4_frames_into(out: bytearray, frames: List[bytes]) -> None: + # One native call for all frames. Every lz4.frame.decompress call releases the + # GIL, and with many busy threads each release can wait a full switch interval. + compressed = b"".join(frames) + stream = pyarrow.CompressedInputStream(pyarrow.BufferReader(compressed), "lz4") + nbytes = min(4 * len(compressed) + 65536, _LZ4_MAX_READ_BYTES) + while True: + chunk = stream.read(nbytes) + if not chunk: + return + out += chunk + + def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema_bytes): ba = bytearray() ba += schema_bytes n_rows = 0 - for arrow_batch in arrow_batches: - n_rows += arrow_batch.rowCount - ba += ( - lz4.frame.decompress(arrow_batch.batch) - if lz4_compressed - else arrow_batch.batch - ) + # The Arrow path releases the GIL about twice, so it only helps from 3 batches. + if lz4_compressed and len(arrow_batches) > 2 and _PYARROW_HAS_LZ4: + n_rows = sum(arrow_batch.rowCount for arrow_batch in arrow_batches) + _decompress_lz4_frames_into(ba, [b.batch for b in arrow_batches]) + else: + for arrow_batch in arrow_batches: + n_rows += arrow_batch.rowCount + ba += ( + lz4.frame.decompress(arrow_batch.batch) + if lz4_compressed + else arrow_batch.batch + ) arrow_table = pyarrow.ipc.open_stream(ba).read_all() return arrow_table, n_rows diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index f127ad8e2..d1fa01d36 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -8,6 +8,7 @@ from unittest.mock import patch, MagicMock, Mock from ssl import CERT_NONE, CERT_REQUIRED from urllib3 import HTTPSConnectionPool +import lz4.frame try: import pyarrow @@ -2200,42 +2201,78 @@ def test_create_arrow_table_calls_correct_conversion_method( arrow_batches, lz4_compressed, schema ) - @patch("lz4.frame.decompress") - @patch("pyarrow.ipc.open_stream") - def test_convert_arrow_based_set_to_arrow_table( - self, open_stream_mock, lz4_decompress_mock - ): - thrift_backend = ThriftDatabricksClient( - "foobar", - 443, - "path", - [], - auth_provider=AuthProvider(), - ssl_options=SSLOptions(), - http_client=MagicMock(), + @staticmethod + def _arrow_set(n_batches, rows_per_batch, lz4_compressed, zeros=False): + n = n_batches * rows_per_batch + table = pyarrow.table( + { + "id": pyarrow.array([0] * n if zeros else range(n), pyarrow.int64()), + "s": pyarrow.array([None if zeros else str(i) for i in range(n)]), + } ) - - lz4_decompress_mock.return_value = bytearray("Testing", "utf-8") - - schema = ( - pyarrow.schema( - [ - pyarrow.field("column1", pyarrow.int32()), - ] + arrow_batches = [] + for record_batch in table.to_batches(max_chunksize=rows_per_batch): + data = record_batch.serialize().to_pybytes() + arrow_batches.append( + ttypes.TSparkArrowBatch( + batch=lz4.frame.compress(data) if lz4_compressed else data, + rowCount=record_batch.num_rows, + ) ) - .serialize() - .to_pybytes() - ) - - arrow_batches = [ - ttypes.TSparkArrowBatch(batch=bytearray("Testing", "utf-8"), rowCount=1) - for _ in range(10) - ] - utils.convert_arrow_based_set_to_arrow_table(arrow_batches, False, schema) - lz4_decompress_mock.assert_not_called() + return table, arrow_batches, table.schema.serialize().to_pybytes() - utils.convert_arrow_based_set_to_arrow_table(arrow_batches, True, schema) - lz4_decompress_mock.assert_called() + def test_convert_arrow_based_set_to_arrow_table(self): + for n_batches in (1, 2, 3, 72): + for lz4_compressed in (False, True): + table, arrow_batches, schema = self._arrow_set( + n_batches, 14, lz4_compressed + ) + result, n_rows = utils.convert_arrow_based_set_to_arrow_table( + arrow_batches, lz4_compressed, schema + ) + self.assertEqual(n_rows, n_batches * 14) + self.assertTrue(result.equals(table), (n_batches, lz4_compressed)) + + def test_convert_arrow_based_set_decodes_many_lz4_batches_in_one_call(self): + # Each lz4.frame.decompress call releases the GIL; many-batch results must + # not pay that per batch. + real_decompress = lz4.frame.decompress + for n_batches, expected_calls in ((2, 2), (72, 0)): + table, arrow_batches, schema = self._arrow_set(n_batches, 14, True) + with patch( + "lz4.frame.decompress", side_effect=real_decompress + ) as decompress: + result, _ = utils.convert_arrow_based_set_to_arrow_table( + arrow_batches, True, schema + ) + self.assertEqual(decompress.call_count, expected_calls) + self.assertTrue(result.equals(table)) + + def test_convert_arrow_based_set_without_pyarrow_lz4_codec(self): + table, arrow_batches, schema = self._arrow_set(72, 14, True) + real_decompress = lz4.frame.decompress + with patch.object(utils, "_PYARROW_HAS_LZ4", False): + with patch( + "lz4.frame.decompress", side_effect=real_decompress + ) as decompress: + result, n_rows = utils.convert_arrow_based_set_to_arrow_table( + arrow_batches, True, schema + ) + self.assertEqual(decompress.call_count, 72) + self.assertEqual(n_rows, 72 * 14) + self.assertTrue(result.equals(table)) + + def test_convert_arrow_based_set_with_highly_compressible_batches(self): + # Decompressed output far larger than the first read size. + table, arrow_batches, schema = self._arrow_set(4, 50000, True, zeros=True) + decompressed = sum(len(lz4.frame.decompress(b.batch)) for b in arrow_batches) + compressed = sum(len(b.batch) for b in arrow_batches) + self.assertGreater(decompressed, 4 * compressed + 65536) + result, n_rows = utils.convert_arrow_based_set_to_arrow_table( + arrow_batches, True, schema + ) + self.assertEqual(n_rows, 200000) + self.assertTrue(result.equals(table)) def test_convert_column_based_set_to_arrow_table_without_nulls(self): # Deliberately duplicate the column name to check that dups work From 2140f9d4fad2e8e8999b4b01d1b5ef698fd36669 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 10:43:41 +0000 Subject: [PATCH 2/4] Raise RuntimeError for corrupt LZ4 input on the multi-batch path pyarrow raises OSError where python-lz4 raises RuntimeError, so the error type depended on how many batches a result had. Also skip the one-call test on pyarrow builds without the lz4 codec. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- src/databricks/sql/utils.py | 16 ++++++++++------ tests/unit/test_thrift_backend.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 65601557d..ebfaf55b0 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -1019,13 +1019,17 @@ def _decompress_lz4_frames_into(out: bytearray, frames: List[bytes]) -> None: # One native call for all frames. Every lz4.frame.decompress call releases the # GIL, and with many busy threads each release can wait a full switch interval. compressed = b"".join(frames) - stream = pyarrow.CompressedInputStream(pyarrow.BufferReader(compressed), "lz4") nbytes = min(4 * len(compressed) + 65536, _LZ4_MAX_READ_BYTES) - while True: - chunk = stream.read(nbytes) - if not chunk: - return - out += chunk + try: + stream = pyarrow.CompressedInputStream(pyarrow.BufferReader(compressed), "lz4") + while True: + chunk = stream.read(nbytes) + if not chunk: + return + out += chunk + except (OSError, pyarrow.ArrowException) as e: + # Raise what python-lz4 raises, so the error type doesn't depend on batch count. + raise RuntimeError(f"LZ4 decompression failed: {e}") from e def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema_bytes): diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index d1fa01d36..eb5963305 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -2233,6 +2233,7 @@ def test_convert_arrow_based_set_to_arrow_table(self): self.assertEqual(n_rows, n_batches * 14) self.assertTrue(result.equals(table), (n_batches, lz4_compressed)) + @unittest.skipUnless(utils._PYARROW_HAS_LZ4, "pyarrow lacks the lz4 codec") def test_convert_arrow_based_set_decodes_many_lz4_batches_in_one_call(self): # Each lz4.frame.decompress call releases the GIL; many-batch results must # not pay that per batch. @@ -2262,6 +2263,15 @@ def test_convert_arrow_based_set_without_pyarrow_lz4_codec(self): self.assertEqual(n_rows, 72 * 14) self.assertTrue(result.equals(table)) + def test_convert_arrow_based_set_corrupt_lz4_raises_runtime_error(self): + for n_batches in (1, 72): + _, arrow_batches, schema = self._arrow_set(n_batches, 14, True) + arrow_batches[-1].batch = arrow_batches[-1].batch[:-20] + with self.assertRaises(RuntimeError, msg=n_batches): + utils.convert_arrow_based_set_to_arrow_table( + arrow_batches, True, schema + ) + def test_convert_arrow_based_set_with_highly_compressible_batches(self): # Decompressed output far larger than the first read size. table, arrow_batches, schema = self._arrow_set(4, 50000, True, zeros=True) From 38cbb9a6fefb71cb5b380280a59023ece320fc78 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 10:47:53 +0000 Subject: [PATCH 3/4] Let MemoryError through; test corruption in middle frames ArrowMemoryError subclasses both MemoryError and ArrowException, so an allocation failure was reported as an LZ4 decompression error. Also cover truncated and corrupted middle frames, with checksums as the warehouse sets them. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- src/databricks/sql/utils.py | 2 ++ tests/unit/test_thrift_backend.py | 37 ++++++++++++++++++++++++++----- 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index ebfaf55b0..2ca394f93 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -1027,6 +1027,8 @@ def _decompress_lz4_frames_into(out: bytearray, frames: List[bytes]) -> None: if not chunk: return out += chunk + except MemoryError: + raise except (OSError, pyarrow.ArrowException) as e: # Raise what python-lz4 raises, so the error type doesn't depend on batch count. raise RuntimeError(f"LZ4 decompression failed: {e}") from e diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index eb5963305..222c154fc 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -2202,7 +2202,7 @@ def test_create_arrow_table_calls_correct_conversion_method( ) @staticmethod - def _arrow_set(n_batches, rows_per_batch, lz4_compressed, zeros=False): + def _arrow_set(n_batches, rows_per_batch, lz4_compressed, zeros=False, **lz4_kw): n = n_batches * rows_per_batch table = pyarrow.table( { @@ -2215,7 +2215,9 @@ def _arrow_set(n_batches, rows_per_batch, lz4_compressed, zeros=False): data = record_batch.serialize().to_pybytes() arrow_batches.append( ttypes.TSparkArrowBatch( - batch=lz4.frame.compress(data) if lz4_compressed else data, + batch=( + lz4.frame.compress(data, **lz4_kw) if lz4_compressed else data + ), rowCount=record_batch.num_rows, ) ) @@ -2264,10 +2266,35 @@ def test_convert_arrow_based_set_without_pyarrow_lz4_codec(self): self.assertTrue(result.equals(table)) def test_convert_arrow_based_set_corrupt_lz4_raises_runtime_error(self): + def truncate(batch): + return batch[:-20] + + def flip_byte(batch): + corrupted = bytearray(batch) + corrupted[len(corrupted) // 2] ^= 0xFF + return bytes(corrupted) + + # Checksums as the warehouse sets them, so flipped bytes are detected. for n_batches in (1, 72): - _, arrow_batches, schema = self._arrow_set(n_batches, 14, True) - arrow_batches[-1].batch = arrow_batches[-1].batch[:-20] - with self.assertRaises(RuntimeError, msg=n_batches): + for which in (n_batches // 2, -1): + for corrupt in (truncate, flip_byte): + _, arrow_batches, schema = self._arrow_set( + n_batches, 14, True, block_checksum=True, content_checksum=True + ) + arrow_batches[which].batch = corrupt(arrow_batches[which].batch) + with self.assertRaises( + RuntimeError, msg=(n_batches, which, corrupt.__name__) + ): + utils.convert_arrow_based_set_to_arrow_table( + arrow_batches, True, schema + ) + + def test_convert_arrow_based_set_lz4_memory_error_is_not_wrapped(self): + _, arrow_batches, schema = self._arrow_set(72, 14, True) + with patch.object( + utils.pyarrow, "CompressedInputStream", side_effect=pyarrow.ArrowMemoryError + ): + with self.assertRaises(MemoryError): utils.convert_arrow_based_set_to_arrow_table( arrow_batches, True, schema ) From c6891ffeed50489284ecbccc1e030e99f5ef6879 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 10:53:40 +0000 Subject: [PATCH 4/4] Close the LZ4 stream; cover unchecksummed truncation Close the CompressedInputStream with a with block so its buffers are freed even when an exception traceback keeps the frame alive. Skip the memory error test without the lz4 codec, test truncated frames without checksums, and note that each batch is one LZ4 frame. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- src/databricks/sql/utils.py | 19 +++++++++++-------- tests/unit/test_thrift_backend.py | 12 ++++++++---- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 2ca394f93..20d449dfd 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -1016,17 +1016,20 @@ def convert_arrow_based_file_to_arrow_table(file_bytes: bytes): def _decompress_lz4_frames_into(out: bytearray, frames: List[bytes]) -> None: - # One native call for all frames. Every lz4.frame.decompress call releases the - # GIL, and with many busy threads each release can wait a full switch interval. + # Each batch is one LZ4 frame; decode them all in one native call. Every + # lz4.frame.decompress call releases the GIL, and with many busy threads each + # release can wait a full switch interval. compressed = b"".join(frames) nbytes = min(4 * len(compressed) + 65536, _LZ4_MAX_READ_BYTES) try: - stream = pyarrow.CompressedInputStream(pyarrow.BufferReader(compressed), "lz4") - while True: - chunk = stream.read(nbytes) - if not chunk: - return - out += chunk + with pyarrow.CompressedInputStream( + pyarrow.BufferReader(compressed), "lz4" + ) as stream: + while True: + chunk = stream.read(nbytes) + if not chunk: + return + out += chunk except MemoryError: raise except (OSError, pyarrow.ArrowException) as e: diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 222c154fc..de3715e6b 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -2274,21 +2274,25 @@ def flip_byte(batch): corrupted[len(corrupted) // 2] ^= 0xFF return bytes(corrupted) - # Checksums as the warehouse sets them, so flipped bytes are detected. + # Truncation is detected with or without checksums. A flipped byte is only + # detected with checksums, which the warehouse sets. + checksums = {"block_checksum": True, "content_checksum": True} + cases = [(truncate, {}), (truncate, checksums), (flip_byte, checksums)] for n_batches in (1, 72): for which in (n_batches // 2, -1): - for corrupt in (truncate, flip_byte): + for corrupt, lz4_kw in cases: _, arrow_batches, schema = self._arrow_set( - n_batches, 14, True, block_checksum=True, content_checksum=True + n_batches, 14, True, **lz4_kw ) arrow_batches[which].batch = corrupt(arrow_batches[which].batch) with self.assertRaises( - RuntimeError, msg=(n_batches, which, corrupt.__name__) + RuntimeError, msg=(n_batches, which, corrupt.__name__, lz4_kw) ): utils.convert_arrow_based_set_to_arrow_table( arrow_batches, True, schema ) + @unittest.skipUnless(utils._PYARROW_HAS_LZ4, "pyarrow lacks the lz4 codec") def test_convert_arrow_based_set_lz4_memory_error_is_not_wrapped(self): _, arrow_batches, schema = self._arrow_set(72, 14, True) with patch.object(