Skip to content

Commit 124cc5d

Browse files
vuanhphungIsaac
andcommitted
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 <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 559f1d3 commit 124cc5d

4 files changed

Lines changed: 129 additions & 6 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
# Unreleased
44
- 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))
5+
- 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.
56

67
# 4.9.0 (2026-11-04)
78
- 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/auth/thrift_http_client.py‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from http.client import HTTPResponse
1212
from io import BytesIO
1313

14+
from thrift.transport.TTransport import CReadableTransport
1415
from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager
1516
from urllib3.util import make_headers
1617
from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy
@@ -22,7 +23,7 @@
2223
logger = logging.getLogger(__name__)
2324

2425

25-
class THttpClient(thrift.transport.THttpClient.THttpClient):
26+
class THttpClient(thrift.transport.THttpClient.THttpClient, CReadableTransport):
2627
realhost: Optional[str]
2728
realport: Optional[int]
2829
proxy_uri: Optional[str]
@@ -94,6 +95,7 @@ def __init__(
9495
self.retry_policy = retry_policy
9596

9697
self.__wbuf = BytesIO()
98+
self.__rbuf = BytesIO()
9799
self.__resp: Union[None, HTTPResponse] = None
98100
self.__timeout = None
99101
self.__custom_headers = None
@@ -153,9 +155,19 @@ def close(self):
153155
self.__resp and self.__resp.drain_conn()
154156
self.__resp and self.__resp.release_conn()
155157
self.__resp = None
158+
self.__rbuf = BytesIO()
156159

157160
def read(self, sz):
158-
return self.__resp.read(sz)
161+
return self.__rbuf.read(sz)
162+
163+
# CReadableTransport: lets TBinaryProtocolAccelerated decode the buffered body in C.
164+
@property
165+
def cstringio_buf(self):
166+
return self.__rbuf
167+
168+
def cstringio_refill(self, partialread, reqlen):
169+
# The whole response body is already buffered, so a refill means it was truncated.
170+
raise EOFError()
159171

160172
def isOpen(self):
161173
return self.__resp is not None
@@ -207,6 +219,10 @@ def flush(self):
207219
self.message = self.__resp.reason
208220
self.headers = self.__resp.headers
209221

222+
# Thrift decodes the reply one field at a time; reading each field
223+
# straight from urllib3 costs far more than the decode itself.
224+
self.__rbuf = BytesIO(self.__resp.read())
225+
210226
logger.info(
211227
"HTTP Response with status code {}, message: {}".format(
212228
self.code, self.message

‎src/databricks/sql/backend/thrift_backend.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,8 @@ def __init__(
253253
self._transport.setTimeout(timeout and (float(timeout) * 1000.0))
254254

255255
self._transport.setCustomHeaders(dict(http_headers))
256-
protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(
256+
# Uses thrift's C extension when available, else the pure-Python protocol.
257+
protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocolAccelerated(
257258
self._transport,
258259
string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT,
259260
container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT,
@@ -443,7 +444,9 @@ def attempt_request(attempt):
443444
this_method_name = getattr(method, "__name__")
444445

445446
logger.debug("Sending request: {}(<REDACTED>)".format(this_method_name))
446-
unsafe_logger.debug("Sending request: {}".format(request))
447+
# Lazy %s args: unsafe_logger is DEBUG-level with a NullHandler by
448+
# default, so eager formatting would stringify every payload.
449+
unsafe_logger.debug("Sending request: %s", request)
447450

448451
# These three lines are no-ops if the v3 retry policy is not in use
449452
if self.enable_v3_retries:
@@ -457,7 +460,7 @@ def attempt_request(attempt):
457460
logger.debug(
458461
"Received response: {}(<REDACTED>)".format(type(response).__name__)
459462
)
460-
unsafe_logger.debug("Received response: {}".format(response))
463+
unsafe_logger.debug("Received response: %s", response)
461464
return response
462465

463466
except urllib3.exceptions.HTTPError as err:

‎tests/unit/test_thrift_backend.py‎

Lines changed: 104 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from collections import OrderedDict
22
from decimal import Decimal
33
import itertools
4+
import logging
45
import unittest
56
import pytest
67
from unittest.mock import patch, MagicMock, Mock
@@ -320,7 +321,7 @@ def test_binary_protocol_preserves_unbounded_string_and_container_limits(
320321
import thrift.protocol.TBinaryProtocol
321322

322323
with patch(
323-
"thrift.protocol.TBinaryProtocol.TBinaryProtocol"
324+
"thrift.protocol.TBinaryProtocol.TBinaryProtocolAccelerated"
324325
) as mock_protocol_class:
325326
ThriftDatabricksClient(
326327
"foo",
@@ -367,6 +368,108 @@ def test_binary_protocol_reads_result_larger_than_thrift_default_limit(self):
367368

368369
self.assertEqual(protocol.readBinary(), oversized_payload)
369370

371+
@staticmethod
372+
def _http_client_with_reply(reply_bytes):
373+
"""A real THttpClient whose connection pool returns ``reply_bytes``."""
374+
from databricks.sql.auth.thrift_http_client import THttpClient
375+
376+
http_client = THttpClient(
377+
auth_provider=Mock(),
378+
uri_or_host="https://example.com/cliservice",
379+
ssl_options=SSLOptions(),
380+
)
381+
http_client.setCustomHeaders({"User-Agent": "test"})
382+
http_client.open()
383+
mock_resp = Mock(status=200, reason="OK", headers={})
384+
mock_resp.read.return_value = reply_bytes
385+
http_client._THttpClient__pool = Mock()
386+
http_client._THttpClient__pool.request.return_value = mock_resp
387+
return http_client, mock_resp
388+
389+
def test_http_client_buffers_reply_and_decodes_with_accelerated_protocol(self):
390+
"""The reply body is read from urllib3 once, then decoded from memory,
391+
including by the C-accelerated protocol (the transport is a
392+
CReadableTransport) for payloads above thrift's ~15.6 MiB default limit."""
393+
from thrift.protocol import TBinaryProtocol
394+
from thrift.transport import TTransport
395+
from databricks.sql.thrift_api.TCLIService import TCLIService
396+
from databricks.sql.backend.thrift_backend import (
397+
THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT,
398+
THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT,
399+
)
400+
401+
batch = b"x" * (20 * 1024 * 1024)
402+
resp = ttypes.TFetchResultsResp(
403+
status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS),
404+
hasMoreRows=False,
405+
results=ttypes.TRowSet(
406+
startRowOffset=0,
407+
rows=[],
408+
arrowBatches=[ttypes.TSparkArrowBatch(batch=batch, rowCount=1)],
409+
),
410+
)
411+
reply = TTransport.TMemoryBuffer()
412+
reply_proto = TBinaryProtocol.TBinaryProtocol(reply)
413+
reply_proto.writeMessageBegin("FetchResults", 2, 0) # REPLY
414+
TCLIService.FetchResults_result(success=resp).write(reply_proto)
415+
reply_proto.writeMessageEnd()
416+
417+
http_client, mock_resp = self._http_client_with_reply(reply.getvalue())
418+
self.assertIsInstance(http_client, TTransport.CReadableTransport)
419+
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(
420+
http_client,
421+
string_length_limit=THRIFT_BINARY_PROTOCOL_STRING_LENGTH_LIMIT,
422+
container_length_limit=THRIFT_BINARY_PROTOCOL_CONTAINER_LENGTH_LIMIT,
423+
)
424+
client = TCLIService.Client(protocol)
425+
426+
op_handle = ttypes.TOperationHandle(
427+
operationId=ttypes.THandleIdentifier(guid=b"g" * 16, secret=b"s" * 16),
428+
operationType=ttypes.TOperationType.EXECUTE_STATEMENT,
429+
hasResultSet=True,
430+
)
431+
result = client.FetchResults(
432+
ttypes.TFetchResultsReq(operationHandle=op_handle, orientation=0, maxRows=1)
433+
)
434+
435+
self.assertEqual(result.results.arrowBatches[0].batch, batch)
436+
mock_resp.read.assert_called_once_with()
437+
438+
http_client.close()
439+
self.assertEqual(http_client.read(1), b"")
440+
441+
def test_http_client_cstringio_refill_on_truncated_reply_raises(self):
442+
http_client, _ = self._http_client_with_reply(b"")
443+
with self.assertRaises(EOFError):
444+
http_client.cstringio_refill(b"", 4)
445+
446+
@patch("databricks.sql.backend.thrift_backend.unsafe_logger")
447+
@patch("databricks.sql.auth.thrift_http_client.THttpClient")
448+
def test_unsafe_logger_formats_payloads_lazily(
449+
self, t_http_client_class, mock_unsafe_logger
450+
):
451+
"""unsafe_logger is DEBUG-level with only a NullHandler by default, so
452+
payloads must be passed as %s args rather than pre-formatted strings."""
453+
thrift_backend = ThriftDatabricksClient(
454+
"foobar",
455+
443,
456+
"path",
457+
[],
458+
auth_provider=AuthProvider(),
459+
ssl_options=SSLOptions(),
460+
http_client=MagicMock(),
461+
)
462+
request = Mock()
463+
response = Mock()
464+
response.status.statusCode = ttypes.TStatusCode.SUCCESS_STATUS
465+
method = Mock(return_value=response)
466+
method.__name__ = "method name"
467+
468+
thrift_backend.make_request(method, request)
469+
470+
mock_unsafe_logger.debug.assert_any_call("Sending request: %s", request)
471+
mock_unsafe_logger.debug.assert_any_call("Received response: %s", response)
472+
370473
def test_proxy_headers_are_set(self):
371474

372475
from databricks.sql.common.http_utils import create_basic_proxy_auth_headers

0 commit comments

Comments
 (0)