From 124cc5d0f77893a706f957f3ccec70e71c14f69c Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 00:49:20 +0000 Subject: [PATCH 1/7] perf(thrift): buffer reply body, use accelerated protocol, log payloads lazily THttpClient.read() forwarded every primitive read of the pure-Python thrift decoder to urllib3, which dominated decode time for wide result schemas and column-based metadata results. Read the reply body once, expose it as a CReadableTransport, and decode with TBinaryProtocolAccelerated (falls back to pure Python when thrift's C extension is unavailable). unsafe_logger.debug() pre-formatted every request and response with str.format, stringifying whole payloads (including inline Arrow batches) even though the logger only has a NullHandler by default. Pass them as %s args. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/auth/thrift_http_client.py | 20 +++- src/databricks/sql/backend/thrift_backend.py | 9 +- tests/unit/test_thrift_backend.py | 105 +++++++++++++++++- 4 files changed, 129 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fbb64b68..1e36c1942 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,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: decode replies from a buffered body 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. # 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/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 0108cc88d..1441a6d14 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -11,6 +11,7 @@ from http.client import HTTPResponse from io import BytesIO +from thrift.transport.TTransport import CReadableTransport from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager from urllib3.util import make_headers from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy @@ -22,7 +23,7 @@ logger = logging.getLogger(__name__) -class THttpClient(thrift.transport.THttpClient.THttpClient): +class THttpClient(thrift.transport.THttpClient.THttpClient, CReadableTransport): realhost: Optional[str] realport: Optional[int] proxy_uri: Optional[str] @@ -94,6 +95,7 @@ def __init__( self.retry_policy = retry_policy self.__wbuf = BytesIO() + self.__rbuf = BytesIO() self.__resp: Union[None, HTTPResponse] = None self.__timeout = None self.__custom_headers = None @@ -153,9 +155,19 @@ def close(self): self.__resp and self.__resp.drain_conn() self.__resp and self.__resp.release_conn() self.__resp = None + self.__rbuf = BytesIO() def read(self, sz): - return self.__resp.read(sz) + return self.__rbuf.read(sz) + + # CReadableTransport: lets TBinaryProtocolAccelerated decode the buffered body in C. + @property + def cstringio_buf(self): + return self.__rbuf + + def cstringio_refill(self, partialread, reqlen): + # The whole response body is already buffered, so a refill means it was truncated. + raise EOFError() def isOpen(self): return self.__resp is not None @@ -207,6 +219,10 @@ def flush(self): self.message = self.__resp.reason self.headers = self.__resp.headers + # Thrift decodes the reply one field at a time; reading each field + # straight from urllib3 costs far more than the decode itself. + self.__rbuf = BytesIO(self.__resp.read()) + logger.info( "HTTP Response with status code {}, message: {}".format( self.code, self.message diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 19e378292..7fefd655f 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -253,7 +253,8 @@ def __init__( self._transport.setTimeout(timeout and (float(timeout) * 1000.0)) self._transport.setCustomHeaders(dict(http_headers)) - protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol( + # Uses thrift's C extension when available, else the pure-Python protocol. + protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocolAccelerated( self._transport, string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, @@ -443,7 +444,9 @@ def attempt_request(attempt): this_method_name = getattr(method, "__name__") logger.debug("Sending request: {}()".format(this_method_name)) - unsafe_logger.debug("Sending request: {}".format(request)) + # Lazy %s args: unsafe_logger is DEBUG-level with a NullHandler by + # default, so eager formatting would stringify every payload. + unsafe_logger.debug("Sending request: %s", request) # These three lines are no-ops if the v3 retry policy is not in use if self.enable_v3_retries: @@ -457,7 +460,7 @@ def attempt_request(attempt): logger.debug( "Received response: {}()".format(type(response).__name__) ) - unsafe_logger.debug("Received response: {}".format(response)) + unsafe_logger.debug("Received response: %s", response) return response except urllib3.exceptions.HTTPError as err: diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index dedadaeb2..4e08774dc 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -1,6 +1,7 @@ from collections import OrderedDict from decimal import Decimal import itertools +import logging import unittest import pytest from unittest.mock import patch, MagicMock, Mock @@ -320,7 +321,7 @@ def test_binary_protocol_preserves_unbounded_string_and_container_limits( import thrift.protocol.TBinaryProtocol with patch( - "thrift.protocol.TBinaryProtocol.TBinaryProtocol" + "thrift.protocol.TBinaryProtocol.TBinaryProtocolAccelerated" ) as mock_protocol_class: ThriftDatabricksClient( "foo", @@ -367,6 +368,108 @@ def test_binary_protocol_reads_result_larger_than_thrift_default_limit(self): self.assertEqual(protocol.readBinary(), oversized_payload) + @staticmethod + def _http_client_with_reply(reply_bytes): + """A real THttpClient whose connection pool returns ``reply_bytes``.""" + from databricks.sql.auth.thrift_http_client import THttpClient + + http_client = THttpClient( + auth_provider=Mock(), + uri_or_host="https://example.com/cliservice", + ssl_options=SSLOptions(), + ) + http_client.setCustomHeaders({"User-Agent": "test"}) + http_client.open() + mock_resp = Mock(status=200, reason="OK", headers={}) + mock_resp.read.return_value = reply_bytes + http_client._THttpClient__pool = Mock() + http_client._THttpClient__pool.request.return_value = mock_resp + return http_client, mock_resp + + def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self): + """The reply body is read from urllib3 once, then decoded from memory, + including by the C-accelerated protocol (the transport is a + CReadableTransport) for payloads above thrift's ~15.6 MiB default limit.""" + from thrift.protocol import TBinaryProtocol + from thrift.transport import TTransport + from databricks.sql.thrift_api.TCLIService import TCLIService + from databricks.sql.backend.thrift_backend import ( + THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, + THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, + ) + + batch = b"x" * (20 * 1024 * 1024) + resp = ttypes.TFetchResultsResp( + status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), + hasMoreRows=False, + results=ttypes.TRowSet( + startRowOffset=0, + rows=[], + arrowBatches=[ttypes.TSparkArrowBatch(batch=batch, rowCount=1)], + ), + ) + reply = TTransport.TMemoryBuffer() + reply_proto = TBinaryProtocol.TBinaryProtocol(reply) + reply_proto.writeMessageBegin("FetchResults", 2, 0) # REPLY + TCLIService.FetchResults_result(success=resp).write(reply_proto) + reply_proto.writeMessageEnd() + + http_client, mock_resp = self._http_client_with_reply(reply.getvalue()) + self.assertIsInstance(http_client, TTransport.CReadableTransport) + protocol = TBinaryProtocol.TBinaryProtocolAccelerated( + http_client, + string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, + container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, + ) + client = TCLIService.Client(protocol) + + op_handle = ttypes.TOperationHandle( + operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16), + operationType=ttypes.TOperationType.EXECUTE_STATEMENT, + hasResultSet=True, + ) + result = client.FetchResults( + ttypes.TFetchResultsReq(operationHandle=op_handle, orientation=0, maxRows=1) + ) + + self.assertEqual(result.results.arrowBatches[0].batch, batch) + mock_resp.read.assert_called_once_with() + + http_client.close() + self.assertEqual(http_client.read(1), b"") + + def test_http_client_cstringio_refill_on_truncated_reply_raises(self): + http_client, _ = self._http_client_with_reply(b"") + with self.assertRaises(EOFError): + http_client.cstringio_refill(b"", 4) + + @patch("databricks.sql.backend.thrift_backend.unsafe_logger") + @patch("databricks.sql.auth.thrift_http_client.THttpClient") + def test_unsafe_logger_formats_payloads_lazily( + self, t_http_client_class, mock_unsafe_logger + ): + """unsafe_logger is DEBUG-level with only a NullHandler by default, so + payloads must be passed as %s args rather than pre-formatted strings.""" + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + request = Mock() + response = Mock() + response.status.statusCode = ttypes.TStatusCode.SUCCESS_STATUS + method = Mock(return_value=response) + method.__name__ = "method name" + + thrift_backend.make_request(method, request) + + mock_unsafe_logger.debug.assert_any_call("Sending request: %s", request) + mock_unsafe_logger.debug.assert_any_call("Received response: %s", response) + def test_proxy_headers_are_set(self): from databricks.sql.common.http_utils import create_basic_proxy_auth_headers From a05dc82508bd1368b8684e2be355fcd87a9ebc65 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 01:04:33 +0000 Subject: [PATCH 2/7] test(thrift): assert the buffered reply is released after make_request Locks in that the transport's reply buffer is cleared once a successful RPC returns (make_request closes the transport in its finally block). Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- tests/unit/test_thrift_backend.py | 50 +++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 4e08774dc..6dac47741 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -443,6 +443,56 @@ def test_http_client_cstringio_refill_on_truncated_reply_raises(self): with self.assertRaises(EOFError): http_client.cstringio_refill(b"", 4) + def test_reply_buffer_is_released_after_successful_make_request(self): + """make_request closes the transport on success too, so the buffered + reply body is not kept on the connection after it has been decoded.""" + from thrift.protocol import TBinaryProtocol + from thrift.transport import TTransport + from databricks.sql.thrift_api.TCLIService import TCLIService + + batch = b"x" * (1024 * 1024) + resp = ttypes.TFetchResultsResp( + status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), + hasMoreRows=False, + results=ttypes.TRowSet( + startRowOffset=0, + rows=[], + arrowBatches=[ttypes.TSparkArrowBatch(batch=batch, rowCount=1)], + ), + ) + reply = TTransport.TMemoryBuffer() + reply_proto = TBinaryProtocol.TBinaryProtocol(reply) + reply_proto.writeMessageBegin("FetchResults", 2, 0) # REPLY + TCLIService.FetchResults_result(success=resp).write(reply_proto) + reply_proto.writeMessageEnd() + + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [("User-Agent", "test")], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + mock_resp = Mock(status=200, reason="OK", headers={}) + mock_resp.read.return_value = reply.getvalue() + thrift_backend._transport._THttpClient__pool = Mock() + thrift_backend._transport._THttpClient__pool.request.return_value = mock_resp + + op_handle = ttypes.TOperationHandle( + operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16), + operationType=ttypes.TOperationType.EXECUTE_STATEMENT, + hasResultSet=True, + ) + result = thrift_backend.make_request( + thrift_backend._client.FetchResults, + ttypes.TFetchResultsReq(operationHandle=op_handle, orientation=0, maxRows=1), + ) + + self.assertEqual(result.results.arrowBatches[0].batch, batch) + self.assertEqual(thrift_backend._transport.cstringio_buf.getvalue(), b"") + @patch("databricks.sql.backend.thrift_backend.unsafe_logger") @patch("databricks.sql.auth.thrift_http_client.THttpClient") def test_unsafe_logger_formats_payloads_lazily( From a1dd416d1a1d5cab7e417c00d2b808f6c3892b73 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 01:05:07 +0000 Subject: [PATCH 3/7] style: black-format new thrift transport test Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- tests/unit/test_thrift_backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 6dac47741..dac50e4cd 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -487,7 +487,9 @@ def test_reply_buffer_is_released_after_successful_make_request(self): ) result = thrift_backend.make_request( thrift_backend._client.FetchResults, - ttypes.TFetchResultsReq(operationHandle=op_handle, orientation=0, maxRows=1), + ttypes.TFetchResultsReq( + operationHandle=op_handle, orientation=0, maxRows=1 + ), ) self.assertEqual(result.results.arrowBatches[0].batch, batch) From 4ab0c5407df55ab39dfe9b423fb41570b42729ed Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 01:12:23 +0000 Subject: [PATCH 4/7] test(thrift): assert the C decoder runs; document reply-buffer trade-off TBinaryProtocolAccelerated silently falls back to pure Python when thrift's C extension isn't built, so add a test (skipped without fastbinary) that asserts the C decoder actually ran on the buffered body. Share reply/request builders across the transport tests, and note in the transport that the raw reply is held alongside the decoded objects until close(). Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/auth/thrift_http_client.py | 4 +- tests/unit/test_thrift_backend.py | 110 ++++++++++-------- 2 files changed, 65 insertions(+), 49 deletions(-) diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 1441a6d14..535e94f55 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -220,7 +220,9 @@ def flush(self): self.headers = self.__resp.headers # Thrift decodes the reply one field at a time; reading each field - # straight from urllib3 costs far more than the decode itself. + # straight from urllib3 costs far more than the decode itself. The + # trade-off: the raw body (bounded by the request's maxBytes) is held + # alongside the decoded objects until close() at the end of the RPC. self.__rbuf = BytesIO(self.__resp.read()) logger.info( diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index dac50e4cd..713036cf6 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -22,6 +22,13 @@ from databricks.sql.result_set import ResultSet, ThriftResultSet from databricks.sql.backend.types import CommandId, CommandState, SessionId, BackendType +try: + from thrift.protocol import fastbinary # noqa: F401 + + _HAS_THRIFT_FASTBINARY = True +except ImportError: + _HAS_THRIFT_FASTBINARY = False + def retry_policy_factory(): return { # (type, default, min, max) @@ -386,19 +393,13 @@ def _http_client_with_reply(reply_bytes): http_client._THttpClient__pool.request.return_value = mock_resp return http_client, mock_resp - def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self): - """The reply body is read from urllib3 once, then decoded from memory, - including by the C-accelerated protocol (the transport is a - CReadableTransport) for payloads above thrift's ~15.6 MiB default limit.""" + @staticmethod + def _fetch_results_reply(batch): + """Wire bytes of a FetchResults REPLY carrying one inline Arrow batch.""" from thrift.protocol import TBinaryProtocol from thrift.transport import TTransport from databricks.sql.thrift_api.TCLIService import TCLIService - from databricks.sql.backend.thrift_backend import ( - THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, - THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, - ) - batch = b"x" * (20 * 1024 * 1024) resp = ttypes.TFetchResultsResp( status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), hasMoreRows=False, @@ -413,24 +414,48 @@ def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self): reply_proto.writeMessageBegin("FetchResults", 2, 0) # REPLY TCLIService.FetchResults_result(success=resp).write(reply_proto) reply_proto.writeMessageEnd() + return reply.getvalue() + + @staticmethod + def _fetch_results_req(): + op_handle = ttypes.TOperationHandle( + operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16), + operationType=ttypes.TOperationType.EXECUTE_STATEMENT, + hasResultSet=True, + ) + return ttypes.TFetchResultsReq( + operationHandle=op_handle, orientation=0, maxRows=1 + ) + + def _accelerated_client(self, http_client): + from thrift.protocol import TBinaryProtocol + from databricks.sql.thrift_api.TCLIService import TCLIService + from databricks.sql.backend.thrift_backend import ( + THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, + THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, + ) - http_client, mock_resp = self._http_client_with_reply(reply.getvalue()) - self.assertIsInstance(http_client, TTransport.CReadableTransport) protocol = TBinaryProtocol.TBinaryProtocolAccelerated( http_client, string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, ) - client = TCLIService.Client(protocol) + return protocol, TCLIService.Client(protocol) - op_handle = ttypes.TOperationHandle( - operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16), - operationType=ttypes.TOperationType.EXECUTE_STATEMENT, - hasResultSet=True, - ) - result = client.FetchResults( - ttypes.TFetchResultsReq(operationHandle=op_handle, orientation=0, maxRows=1) + def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self): + """The reply body is read from urllib3 once and decoded from memory via + TBinaryProtocolAccelerated, for payloads above thrift's ~15.6 MiB + default limit. Runs on both the C and the pure-Python fallback path.""" + from thrift.transport import TTransport + + batch = b"x" * (20 * 1024 * 1024) + http_client, mock_resp = self._http_client_with_reply( + self._fetch_results_reply(batch) ) + self.assertIsInstance(http_client, TTransport.CReadableTransport) + _, client = self._accelerated_client(http_client) + + result = client.FetchResults(self._fetch_results_req()) self.assertEqual(result.results.arrowBatches[0].batch, batch) mock_resp.read.assert_called_once_with() @@ -438,6 +463,22 @@ def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self): http_client.close() self.assertEqual(http_client.read(1), b"") + @unittest.skipUnless( + _HAS_THRIFT_FASTBINARY, "thrift C extension (fastbinary) not built" + ) + def test_http_client_reply_is_decoded_by_thrift_c_extension(self): + """TBinaryProtocolAccelerated silently falls back to pure Python without + fastbinary, so assert the C decoder actually ran on the buffered body.""" + batch = b"x" * (20 * 1024 * 1024) + http_client, _ = self._http_client_with_reply(self._fetch_results_reply(batch)) + protocol, client = self._accelerated_client(http_client) + protocol._fast_decode = Mock(wraps=protocol._fast_decode) + + result = client.FetchResults(self._fetch_results_req()) + + self.assertEqual(result.results.arrowBatches[0].batch, batch) + protocol._fast_decode.assert_called() + def test_http_client_cstringio_refill_on_truncated_reply_raises(self): http_client, _ = self._http_client_with_reply(b"") with self.assertRaises(EOFError): @@ -446,26 +487,7 @@ def test_http_client_cstringio_refill_on_truncated_reply_raises(self): def test_reply_buffer_is_released_after_successful_make_request(self): """make_request closes the transport on success too, so the buffered reply body is not kept on the connection after it has been decoded.""" - from thrift.protocol import TBinaryProtocol - from thrift.transport import TTransport - from databricks.sql.thrift_api.TCLIService import TCLIService - batch = b"x" * (1024 * 1024) - resp = ttypes.TFetchResultsResp( - status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), - hasMoreRows=False, - results=ttypes.TRowSet( - startRowOffset=0, - rows=[], - arrowBatches=[ttypes.TSparkArrowBatch(batch=batch, rowCount=1)], - ), - ) - reply = TTransport.TMemoryBuffer() - reply_proto = TBinaryProtocol.TBinaryProtocol(reply) - reply_proto.writeMessageBegin("FetchResults", 2, 0) # REPLY - TCLIService.FetchResults_result(success=resp).write(reply_proto) - reply_proto.writeMessageEnd() - thrift_backend = ThriftDatabricksClient( "foobar", 443, @@ -476,20 +498,12 @@ def test_reply_buffer_is_released_after_successful_make_request(self): http_client=MagicMock(), ) mock_resp = Mock(status=200, reason="OK", headers={}) - mock_resp.read.return_value = reply.getvalue() + mock_resp.read.return_value = self._fetch_results_reply(batch) thrift_backend._transport._THttpClient__pool = Mock() thrift_backend._transport._THttpClient__pool.request.return_value = mock_resp - op_handle = ttypes.TOperationHandle( - operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16), - operationType=ttypes.TOperationType.EXECUTE_STATEMENT, - hasResultSet=True, - ) result = thrift_backend.make_request( - thrift_backend._client.FetchResults, - ttypes.TFetchResultsReq( - operationHandle=op_handle, orientation=0, maxRows=1 - ), + thrift_backend._client.FetchResults, self._fetch_results_req() ) self.assertEqual(result.results.arrowBatches[0].batch, batch) From 4a8e3772969c915a3f12ba81b9a8a3c8650883fd Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 01:19:14 +0000 Subject: [PATCH 5/7] test(thrift): check C-encoded requests match pure Python; qualify buffer bound TBinaryProtocolAccelerated also moves request encoding to C. Add a test (skipped without fastbinary) that builds OpenSession, ExecuteStatement (with native parameters and query tags), GetTables and GetColumns requests through the backend and checks both encoders produce identical bytes. Clarify that the buffered reply is bounded by maxBytes only for fetches. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/auth/thrift_http_client.py | 4 +- tests/unit/test_thrift_backend.py | 90 +++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 535e94f55..2e98eaab2 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -221,8 +221,8 @@ def flush(self): # Thrift decodes the reply one field at a time; reading each field # straight from urllib3 costs far more than the decode itself. The - # trade-off: the raw body (bounded by the request's maxBytes) is held - # alongside the decoded objects until close() at the end of the RPC. + # trade-off: the raw body is held alongside the decoded objects until + # close() at the end of the RPC; for fetches it is bounded by maxBytes. self.__rbuf = BytesIO(self.__resp.read()) logger.info( diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 713036cf6..c50bd550d 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -509,6 +509,96 @@ def test_reply_buffer_is_released_after_successful_make_request(self): self.assertEqual(result.results.arrowBatches[0].batch, batch) self.assertEqual(thrift_backend._transport.cstringio_buf.getvalue(), b"") + @unittest.skipUnless( + _HAS_THRIFT_FASTBINARY, "thrift C extension (fastbinary) not built" + ) + @patch("databricks.sql.auth.thrift_http_client.THttpClient") + def test_accelerated_protocol_encodes_backend_requests_identically( + self, t_http_client_class + ): + """TBinaryProtocolAccelerated also moves request encoding to C. Requests + built by the backend must serialize byte-for-byte as with pure Python.""" + import datetime + from thrift.protocol import TBinaryProtocol + from thrift.transport import TTransport + from databricks.sql.parameters import native + + class Captured(Exception): + pass + + captured = [] + + def capture(method, request, *args, **kwargs): + captured.append(request) + raise Captured() + + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + thrift_backend.make_request = capture + session_id = SessionId.from_thrift_handle( + ttypes.TSessionHandle( + sessionId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16) + ) + ) + params = [ + native.IntegerParameter(value=42, name="i"), + native.BigIntegerParameter(value=2**40, name="b"), + native.StringParameter(value="héllo", name="s"), + native.DecimalParameter(value=Decimal("12.345"), name="d"), + native.TimestampParameter( + value=datetime.datetime(2026, 11, 11, 1, 2, 3), name="t" + ), + native.DateParameter(value=datetime.date(2026, 11, 11), name="dt"), + native.BooleanParameter(value=True, name="bo"), + native.VoidParameter(value=None, name="v"), + native.ArrayParameter(value=[1, 2, 3], name="a"), + native.MapParameter(value={"k": "v"}, name="m"), + ] + calls = [ + lambda: thrift_backend.open_session( + {"use_cached_result": False, "ansi_mode": "true"}, "main", "default" + ), + lambda: thrift_backend.execute_command( + operation="SELECT :i, :s", + session_id=session_id, + max_rows=100000, + max_bytes=104857600, + lz4_compression=True, + cursor=Mock(), + use_cloud_fetch=True, + parameters=[p.as_tspark_param(named=True) for p in params], + async_op=False, + enforce_embedded_schema_correctness=False, + row_limit=10, + query_tags={"team": "peco", "flag": None}, + ), + lambda: thrift_backend.get_tables( + session_id, 100000, 104857600, Mock(), "main", "%", "t%", ["TABLE"] + ), + lambda: thrift_backend.get_columns( + session_id, 100000, 104857600, Mock(), "main", "%", "%", "%" + ), + ] + for call in calls: + with self.assertRaises(Captured): + call() + + self.assertEqual(len(captured), len(calls)) + for request in captured: + pure, accelerated = TTransport.TMemoryBuffer(), TTransport.TMemoryBuffer() + request.write(TBinaryProtocol.TBinaryProtocol(pure)) + request.write(TBinaryProtocol.TBinaryProtocolAccelerated(accelerated)) + self.assertEqual( + accelerated.getvalue(), pure.getvalue(), type(request).__name__ + ) + @patch("databricks.sql.backend.thrift_backend.unsafe_logger") @patch("databricks.sql.auth.thrift_http_client.THttpClient") def test_unsafe_logger_formats_payloads_lazily( From 44077a0cd647eafab5fab0457781924d64d2e00d Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 06:14:04 +0000 Subject: [PATCH 6/7] perf(thrift): read the reply in 1 MiB chunks instead of buffering it whole Reading the whole body in flush() held a second copy of it during decode and, on urllib3 1.26, turned a truncated 503/429 body into a ProtocolError that skipped the status-based retry. The transport now refills from urllib3 in 1 MiB chunks (or a whole field if larger), like TBufferedTransport, so truncated bodies fail inside decode as on main. Peak heap for a 100 MiB multi-batch reply drops from ~200 to ~103 MiB at the same speed. Tests now run the decode paths under both the C and the pure-Python protocol, and cover the truncated-503 retry. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 2 +- src/databricks/sql/auth/thrift_http_client.py | 30 ++-- tests/unit/test_thrift_backend.py | 132 +++++++++++++----- 3 files changed, 121 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e36c1942..d9c58ebf3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,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: decode replies from a buffered body 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 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. # 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/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 2e98eaab2..292009d71 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -157,17 +157,34 @@ def close(self): self.__resp = None self.__rbuf = BytesIO() + # Thrift decodes the reply one field at a time, and one urllib3 read per field + # costs far more than the decode, so the body is read in chunks of this size. + _READ_CHUNK_SIZE = 1024 * 1024 + def read(self, sz): + ret = self.__rbuf.read(sz) + if ret: + return ret + self.__rbuf = BytesIO(self.__resp.read(max(sz, self._READ_CHUNK_SIZE))) return self.__rbuf.read(sz) - # CReadableTransport: lets TBinaryProtocolAccelerated decode the buffered body in C. + # CReadableTransport: lets TBinaryProtocolAccelerated decode the chunks in C. @property def cstringio_buf(self): return self.__rbuf def cstringio_refill(self, partialread, reqlen): - # The whole response body is already buffered, so a refill means it was truncated. - raise EOFError() + # As in TBufferedTransport: read one more chunk, or all of a larger field. + data = partialread + if reqlen < self._READ_CHUNK_SIZE: + data += self.__resp.read(self._READ_CHUNK_SIZE) + while len(data) < reqlen: + chunk = self.__resp.read(reqlen - len(data)) + if not chunk: + raise EOFError() + data += chunk + self.__rbuf = BytesIO(data) + return self.__rbuf def isOpen(self): return self.__resp is not None @@ -218,12 +235,7 @@ def flush(self): self.code = self.__resp.status self.message = self.__resp.reason self.headers = self.__resp.headers - - # Thrift decodes the reply one field at a time; reading each field - # straight from urllib3 costs far more than the decode itself. The - # trade-off: the raw body is held alongside the decoded objects until - # close() at the end of the RPC; for fetches it is bounded by maxBytes. - self.__rbuf = BytesIO(self.__resp.read()) + self.__rbuf = BytesIO() logger.info( "HTTP Response with status code {}, message: {}".format( diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index c50bd550d..e69b63941 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -1,5 +1,6 @@ from collections import OrderedDict from decimal import Decimal +from io import BytesIO import itertools import logging import unittest @@ -377,7 +378,7 @@ def test_binary_protocol_reads_result_larger_than_thrift_default_limit(self): @staticmethod def _http_client_with_reply(reply_bytes): - """A real THttpClient whose connection pool returns ``reply_bytes``.""" + """A real THttpClient whose connection pool streams ``reply_bytes``.""" from databricks.sql.auth.thrift_http_client import THttpClient http_client = THttpClient( @@ -388,14 +389,14 @@ def _http_client_with_reply(reply_bytes): http_client.setCustomHeaders({"User-Agent": "test"}) http_client.open() mock_resp = Mock(status=200, reason="OK", headers={}) - mock_resp.read.return_value = reply_bytes + mock_resp.read.side_effect = BytesIO(reply_bytes).read http_client._THttpClient__pool = Mock() http_client._THttpClient__pool.request.return_value = mock_resp return http_client, mock_resp @staticmethod - def _fetch_results_reply(batch): - """Wire bytes of a FetchResults REPLY carrying one inline Arrow batch.""" + def _fetch_results_reply(batches): + """Wire bytes of a FetchResults REPLY carrying inline Arrow ``batches``.""" from thrift.protocol import TBinaryProtocol from thrift.transport import TTransport from databricks.sql.thrift_api.TCLIService import TCLIService @@ -406,7 +407,10 @@ def _fetch_results_reply(batch): results=ttypes.TRowSet( startRowOffset=0, rows=[], - arrowBatches=[ttypes.TSparkArrowBatch(batch=batch, rowCount=1)], + arrowBatches=[ + ttypes.TSparkArrowBatch(batch=batch, rowCount=1) + for batch in batches + ], ), ) reply = TTransport.TMemoryBuffer() @@ -427,51 +431,71 @@ def _fetch_results_req(): operationHandle=op_handle, orientation=0, maxRows=1 ) - def _accelerated_client(self, http_client): - from thrift.protocol import TBinaryProtocol + @staticmethod + def _protocol_client(http_client, protocol_class): from databricks.sql.thrift_api.TCLIService import TCLIService from databricks.sql.backend.thrift_backend import ( THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, ) - protocol = TBinaryProtocol.TBinaryProtocolAccelerated( + protocol = protocol_class( http_client, string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT, container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT, ) return protocol, TCLIService.Client(protocol) - def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self): - """The reply body is read from urllib3 once and decoded from memory via - TBinaryProtocolAccelerated, for payloads above thrift's ~15.6 MiB - default limit. Runs on both the C and the pure-Python fallback path.""" - from thrift.transport import TTransport + @staticmethod + def _protocol_classes(): + """The C-accelerated protocol the backend uses, and the pure-Python + protocol it falls back to when thrift's C extension isn't built.""" + from thrift.protocol import TBinaryProtocol - batch = b"x" * (20 * 1024 * 1024) - http_client, mock_resp = self._http_client_with_reply( - self._fetch_results_reply(batch) + return ( + TBinaryProtocol.TBinaryProtocolAccelerated, + TBinaryProtocol.TBinaryProtocol, ) - self.assertIsInstance(http_client, TTransport.CReadableTransport) - _, client = self._accelerated_client(http_client) - result = client.FetchResults(self._fetch_results_req()) + def test_http_client_reads_reply_in_chunks(self): + """The reply is pulled from urllib3 in a few chunk-sized reads: not one + read per field, and never the whole body at once. Includes a field + larger than a chunk and than thrift's ~15.6 MiB default limit.""" + from thrift.transport import TTransport - self.assertEqual(result.results.arrowBatches[0].batch, batch) - mock_resp.read.assert_called_once_with() + small = [bytes([i % 256]) * 1000 for i in range(3000)] + big = b"x" * (20 * 1024 * 1024) + reply = self._fetch_results_reply(small + [big]) + for protocol_class in self._protocol_classes(): + with self.subTest(protocol=protocol_class.__name__): + http_client, mock_resp = self._http_client_with_reply(reply) + self.assertIsInstance(http_client, TTransport.CReadableTransport) + _, client = self._protocol_client(http_client, protocol_class) + + result = client.FetchResults(self._fetch_results_req()) - http_client.close() - self.assertEqual(http_client.read(1), b"") + self.assertEqual( + [b.batch for b in result.results.arrowBatches], small + [big] + ) + read_sizes = [c.args[0] for c in mock_resp.read.call_args_list] + self.assertLess(len(read_sizes), 10) + self.assertLessEqual(max(read_sizes), len(big)) @unittest.skipUnless( _HAS_THRIFT_FASTBINARY, "thrift C extension (fastbinary) not built" ) def test_http_client_reply_is_decoded_by_thrift_c_extension(self): """TBinaryProtocolAccelerated silently falls back to pure Python without - fastbinary, so assert the C decoder actually ran on the buffered body.""" + fastbinary, so assert the C decoder actually ran on the streamed body.""" + from thrift.protocol import TBinaryProtocol + batch = b"x" * (20 * 1024 * 1024) - http_client, _ = self._http_client_with_reply(self._fetch_results_reply(batch)) - protocol, client = self._accelerated_client(http_client) + http_client, _ = self._http_client_with_reply( + self._fetch_results_reply([batch]) + ) + protocol, client = self._protocol_client( + http_client, TBinaryProtocol.TBinaryProtocolAccelerated + ) protocol._fast_decode = Mock(wraps=protocol._fast_decode) result = client.FetchResults(self._fetch_results_req()) @@ -479,14 +503,56 @@ def test_http_client_reply_is_decoded_by_thrift_c_extension(self): self.assertEqual(result.results.arrowBatches[0].batch, batch) protocol._fast_decode.assert_called() - def test_http_client_cstringio_refill_on_truncated_reply_raises(self): - http_client, _ = self._http_client_with_reply(b"") - with self.assertRaises(EOFError): - http_client.cstringio_refill(b"", 4) + def test_http_client_truncated_reply_raises_eof(self): + reply = self._fetch_results_reply([b"x" * 100]) + for protocol_class in self._protocol_classes(): + with self.subTest(protocol=protocol_class.__name__): + http_client, _ = self._http_client_with_reply(reply[:40]) + _, client = self._protocol_client(http_client, protocol_class) + with self.assertRaises(EOFError): + client.FetchResults(self._fetch_results_req()) + + @patch("time.sleep") + def test_truncated_503_reply_is_retried(self, mock_sleep): + """A 503 whose body is cut short must still reach the status-based retry. + urllib3 1.26 raises on an unbounded read() of a truncated body, but not on + read(amt), so the body must only be read in bounded chunks.""" + import urllib3 + + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [("User-Agent", "test")], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + _enable_v3_retries=False, + ) + truncated_body = BytesIO(b"Service Unavail") + + def urllib3_1_26_read(amt=None): + if amt is None: + raise urllib3.exceptions.ProtocolError("IncompleteRead") + return truncated_body.read(amt) + + unavailable = Mock(status=503, reason="", headers={"Retry-After": "1"}) + unavailable.read.side_effect = urllib3_1_26_read + ok = Mock(status=200, reason="OK", headers={}) + ok.read.side_effect = BytesIO(self._fetch_results_reply([b"x" * 100])).read + pool = thrift_backend._transport._THttpClient__pool = Mock() + pool.request.side_effect = [unavailable, ok] + + result = thrift_backend.make_request( + thrift_backend._client.FetchResults, self._fetch_results_req() + ) + + self.assertEqual(result.results.arrowBatches[0].batch, b"x" * 100) + self.assertEqual(pool.request.call_count, 2) def test_reply_buffer_is_released_after_successful_make_request(self): - """make_request closes the transport on success too, so the buffered - reply body is not kept on the connection after it has been decoded.""" + """make_request closes the transport on success too, so the last chunk + of the reply is not kept on the connection after it has been decoded.""" batch = b"x" * (1024 * 1024) thrift_backend = ThriftDatabricksClient( "foobar", @@ -498,7 +564,7 @@ def test_reply_buffer_is_released_after_successful_make_request(self): http_client=MagicMock(), ) mock_resp = Mock(status=200, reason="OK", headers={}) - mock_resp.read.return_value = self._fetch_results_reply(batch) + mock_resp.read.side_effect = BytesIO(self._fetch_results_reply([batch])).read thrift_backend._transport._THttpClient__pool = Mock() thrift_backend._transport._THttpClient__pool.request.return_value = mock_resp From 2163819f4569e6191ec4a8458d3adbef91c58921 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 06:22:29 +0000 Subject: [PATCH 7/7] perf(thrift): join refill pieces once instead of concatenating per read If urllib3 returns short reads for a field larger than a chunk, bytes += recopied everything read so far on every read. Collect the pieces and join them once. Adds a test that decodes a reply from short reads on both protocols. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/auth/thrift_http_client.py | 14 ++++++++------ tests/unit/test_thrift_backend.py | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 292009d71..72a20bbf6 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -175,15 +175,17 @@ def cstringio_buf(self): def cstringio_refill(self, partialread, reqlen): # As in TBufferedTransport: read one more chunk, or all of a larger field. - data = partialread + parts, have = [partialread], len(partialread) if reqlen < self._READ_CHUNK_SIZE: - data += self.__resp.read(self._READ_CHUNK_SIZE) - while len(data) < reqlen: - chunk = self.__resp.read(reqlen - len(data)) + parts.append(self.__resp.read(self._READ_CHUNK_SIZE)) + have += len(parts[-1]) + while have < reqlen: + chunk = self.__resp.read(reqlen - have) if not chunk: raise EOFError() - data += chunk - self.__rbuf = BytesIO(data) + parts.append(chunk) + have += len(chunk) + self.__rbuf = BytesIO(b"".join(parts)) return self.__rbuf def isOpen(self): diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index e69b63941..eb4c26d82 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -503,6 +503,25 @@ def test_http_client_reply_is_decoded_by_thrift_c_extension(self): self.assertEqual(result.results.arrowBatches[0].batch, batch) protocol._fast_decode.assert_called() + def test_http_client_decodes_reply_from_short_urllib3_reads(self): + """urllib3 may return fewer bytes than asked for; a field spanning many + short reads must still be assembled correctly.""" + small = [bytes([i % 256]) * 1000 for i in range(300)] + big = b"y" * (4 * 1024 * 1024) + reply = self._fetch_results_reply(small + [big]) + for protocol_class in self._protocol_classes(): + with self.subTest(protocol=protocol_class.__name__): + http_client, mock_resp = self._http_client_with_reply(b"") + stream = BytesIO(reply) + mock_resp.read.side_effect = lambda amt: stream.read(min(amt, 65536)) + _, client = self._protocol_client(http_client, protocol_class) + + result = client.FetchResults(self._fetch_results_req()) + + self.assertEqual( + [b.batch for b in result.results.arrowBatches], small + [big] + ) + def test_http_client_truncated_reply_raises_eof(self): reply = self._fetch_results_reply([b"x" * 100]) for protocol_class in self._protocol_classes():