Skip to content

Commit 307c48d

Browse files
vuanhphungIsaac
andauthored
perf(thrift): buffer reply body, use accelerated protocol, log payloads lazily (#1020)
* 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> * 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 <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * style: black-format new thrift transport test Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * 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 <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * 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 <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * 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 <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * 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 <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * fix(thrift): keep pure-Python request encoding, decode replies in C TBinaryProtocolAccelerated also moved request encoding to fastbinary, whose recursion limit rejects a native array parameter nested 62+ levels once it is wrapped in ExecuteStatement_args. Main's pure-Python encoder accepts it, and a warehouse runs such queries at depths 62 and 63. Clear the protocol's _fast_encode so requests serialize as before, while replies still decode in C. Requests are small, so this has no measurable cost. Replace the C-vs-Python request byte comparison with a test that sends a 63-level nested parameter through make_request and checks the wire bytes. Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * refactor(thrift): split decode/encode protocols via Client(iprot, oprot) Instead of clearing thrift's private _fast_encode attribute, pass the generated client a TBinaryProtocolAccelerated for replies and a plain TBinaryProtocol for requests. The nested-parameter test now also checks the protocol types, so it guards the split where fastbinary isn't built. Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> --------- Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 4394c11 commit 307c48d

4 files changed

Lines changed: 367 additions & 7 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: 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.
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: 34 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,38 @@ 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()
159+
160+
# Thrift decodes the reply one field at a time, and one urllib3 read per field
161+
# costs far more than the decode, so the body is read in chunks of this size.
162+
_READ_CHUNK_SIZE = 1024 * 1024
156163

157164
def read(self, sz):
158-
return self.__resp.read(sz)
165+
ret = self.__rbuf.read(sz)
166+
if ret:
167+
return ret
168+
self.__rbuf = BytesIO(self.__resp.read(max(sz, self._READ_CHUNK_SIZE)))
169+
return self.__rbuf.read(sz)
170+
171+
# CReadableTransport: lets TBinaryProtocolAccelerated decode the chunks in C.
172+
@property
173+
def cstringio_buf(self):
174+
return self.__rbuf
175+
176+
def cstringio_refill(self, partialread, reqlen):
177+
# As in TBufferedTransport: read one more chunk, or all of a larger field.
178+
parts, have = [partialread], len(partialread)
179+
if reqlen < self._READ_CHUNK_SIZE:
180+
parts.append(self.__resp.read(self._READ_CHUNK_SIZE))
181+
have += len(parts[-1])
182+
while have < reqlen:
183+
chunk = self.__resp.read(reqlen - have)
184+
if not chunk:
185+
raise EOFError()
186+
parts.append(chunk)
187+
have += len(chunk)
188+
self.__rbuf = BytesIO(b"".join(parts))
189+
return self.__rbuf
159190

160191
def isOpen(self):
161192
return self.__resp is not None
@@ -206,6 +237,7 @@ def flush(self):
206237
self.code = self.__resp.status
207238
self.message = self.__resp.reason
208239
self.headers = self.__resp.headers
240+
self.__rbuf = BytesIO()
209241

210242
logger.info(
211243
"HTTP Response with status code {}, message: {}".format(

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

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -253,12 +253,18 @@ 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+
# Replies decode with thrift's C extension when available, else in pure Python.
257+
input_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,
260261
)
261-
self._client = TCLIService.Client(protocol)
262+
# Requests use the pure-Python encoder: the C one rejects structs nested
263+
# deeper than ~64 levels, e.g. native array parameters nested 62+ deep.
264+
output_protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(
265+
self._transport
266+
)
267+
self._client = TCLIService.Client(input_protocol, output_protocol)
262268

263269
try:
264270
self._transport.open()
@@ -443,7 +449,9 @@ def attempt_request(attempt):
443449
this_method_name = getattr(method, "__name__")
444450

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

448456
# These three lines are no-ops if the v3 retry policy is not in use
449457
if self.enable_v3_retries:
@@ -457,7 +465,7 @@ def attempt_request(attempt):
457465
logger.debug(
458466
"Received response: {}(<REDACTED>)".format(type(response).__name__)
459467
)
460-
unsafe_logger.debug("Received response: {}".format(response))
468+
unsafe_logger.debug("Received response: %s", response)
461469
return response
462470

463471
except urllib3.exceptions.HTTPError as err:

0 commit comments

Comments
 (0)