From 53893d3fb4c015e708c492a36dc09ee66ed834f0 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Sat, 14 Nov 2026 21:01:08 +0000 Subject: [PATCH 1/4] perf: skip latency timing nothing records; build rows in fewer calls log_latency timed every decorated call (execute, execute_async, metadata calls) and then dropped the result when telemetry was off and DEBUG logging was disabled, which is always the case on kernel connections. It now checks first and calls the method directly. Row factories built each row through Row.__call__, _create_row, Row.__new__ and Row.__setattr__; they now build the tuple and set __fields__ directly, producing the same object. Single-thread benchmark of one-row kernel queries with the kernel faked: 17.2 -> 16.2 us of Python per query; a 5-column row factory call drops from 402 to 288 ns. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + .../sql/telemetry/latency_logger.py | 16 ++- src/databricks/sql/types.py | 14 ++- tests/unit/test_latency_logger.py | 102 ++++++++++++++++++ tests/unit/test_row.py | 50 +++++++++ 5 files changed, 176 insertions(+), 7 deletions(-) create mode 100644 tests/unit/test_latency_logger.py create mode 100644 tests/unit/test_row.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 867be71ec..7ae79b06c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - 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 fetching small results on the kernel backend (`use_kernel=True`): they are converted to rows without building a pandas DataFrame when that is cheaper and gives the same values, which cut client CPU per query by 47% in a local benchmark of 200 concurrent sessions running one-row queries against a mock server. - Cut more per-query work on the kernel backend: reuse the cursor description while a connection's statements return the same result schema, and build small results' rows straight from the record batches instead of first combining them into a table. +- Cut per-call and per-row work on all backends: don't time `execute()` and metadata calls when neither telemetry nor debug logging would record the latency, and build each result `Row` with fewer Python calls. # 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/telemetry/latency_logger.py b/src/databricks/sql/telemetry/latency_logger.py index 2445c25c2..2bb86cdb9 100644 --- a/src/databricks/sql/telemetry/latency_logger.py +++ b/src/databricks/sql/telemetry/latency_logger.py @@ -177,6 +177,17 @@ def execute(self, query): def decorator(func): @functools.wraps(func) def wrapper(self, *args, **kwargs): + # Fast check: use cached telemetry_enabled flag from connection + # Avoids dictionary lookup + instance check on every operation + connection = getattr(self, "connection", None) + telemetry_enabled = bool(connection) and getattr( + connection, "telemetry_enabled", False + ) + if not telemetry_enabled and not logger.isEnabledFor(logging.DEBUG): + # Nothing would record the latency (always the case on kernel + # connections, which export telemetry from the kernel). + return func(self, *args, **kwargs) + start_time = time.monotonic() try: return func(self, *args, **kwargs) @@ -186,10 +197,7 @@ def wrapper(self, *args, **kwargs): # Always log for debugging logger.debug("%s completed in %dms", func.__name__, duration_ms) - # Fast check: use cached telemetry_enabled flag from connection - # Avoids dictionary lookup + instance check on every operation - connection = getattr(self, "connection", None) - if connection and getattr(connection, "telemetry_enabled", False): + if telemetry_enabled: session_id_hex = connection.get_session_id_hex() if session_id_hex: # Telemetry enabled - extract and send diff --git a/src/databricks/sql/types.py b/src/databricks/sql/types.py index 2d1ddc42a..9d250160a 100644 --- a/src/databricks/sql/types.py +++ b/src/databricks/sql/types.py @@ -257,7 +257,10 @@ def __call__(self, *args: Any) -> "Row": "Can not create Row with fields %s, expected %d values " "but got %s" % (self, len(self), args) ) - return _create_row(self, args) + # Same as _create_row(self, args), inlined: this runs once per fetched row. + row = _tuple_new(Row, args) + row.__dict__["__fields__"] = self + return row def __getitem__(self, item: Any) -> Any: if isinstance(item, (int, slice)): @@ -309,9 +312,14 @@ def __repr__(self) -> str: return "" % ", ".join("%r" % field for field in self) +_tuple_new = tuple.__new__ + + def _create_row( fields: Union["Row", List[str]], values: Union[Tuple[Any, ...], List[Any]] ) -> "Row": - row = Row(*values) - row.__fields__ = fields + # Equivalent to ``row = Row(*values); row.__fields__ = fields`` without the + # Row.__new__ and Row.__setattr__ calls. + row = _tuple_new(Row, values) + row.__dict__["__fields__"] = fields return row diff --git a/tests/unit/test_latency_logger.py b/tests/unit/test_latency_logger.py new file mode 100644 index 000000000..c6b8f8a58 --- /dev/null +++ b/tests/unit/test_latency_logger.py @@ -0,0 +1,102 @@ +import logging +from unittest.mock import MagicMock, patch + +import pytest + +from databricks.sql.telemetry.latency_logger import log_latency +from databricks.sql.telemetry.models.enums import StatementType + +LOGGER = "databricks.sql.telemetry.latency_logger" + + +class Cursor: + """Stand-in for databricks.sql.client.Cursor (telemetry keys on the class name).""" + + def __init__(self, telemetry_enabled): + self.connection = MagicMock() + self.connection.telemetry_enabled = telemetry_enabled + self.connection.get_session_id_hex.return_value = "session-hex" + self.connection.session.host = "host" + self.connection.lz4_compression = False + self.query_id = "statement-id" + self.active_result_set = None + self.backend = MagicMock(retry_policy=None) + + @log_latency(StatementType.QUERY) + def execute(self, value): + return value + + @log_latency(StatementType.QUERY) + def fail(self): + raise ValueError("boom") + + +@pytest.fixture +def telemetry_client(): + client = MagicMock() + with patch( + "databricks.sql.telemetry.latency_logger.TelemetryClientFactory.get_telemetry_client", + return_value=client, + ): + yield client + + +@pytest.fixture +def no_debug(): + logger = logging.getLogger(LOGGER) + level = logger.level + logger.setLevel(logging.INFO) + yield + logger.setLevel(level) + + +def test_disabled_telemetry_without_debug_only_calls_the_method( + telemetry_client, no_debug, caplog +): + cursor = Cursor(telemetry_enabled=False) + with caplog.at_level(logging.INFO, logger=LOGGER): + assert cursor.execute(42) == 42 + with pytest.raises(ValueError, match="boom"): + cursor.fail() + telemetry_client.export_latency_log.assert_not_called() + assert not caplog.records + + +def test_debug_logging_records_latency_without_telemetry(telemetry_client, caplog): + cursor = Cursor(telemetry_enabled=False) + with caplog.at_level(logging.DEBUG, logger=LOGGER): + assert cursor.execute(1) == 1 + assert [r.getMessage() for r in caplog.records] == ["execute completed in 0ms"] + telemetry_client.export_latency_log.assert_not_called() + + +@pytest.mark.parametrize("debug", [False, True]) +def test_enabled_telemetry_exports_latency_even_on_error( + telemetry_client, no_debug, debug, caplog +): + cursor = Cursor(telemetry_enabled=True) + with caplog.at_level(logging.DEBUG if debug else logging.INFO, logger=LOGGER): + assert cursor.execute("ok") == "ok" + with pytest.raises(ValueError, match="boom"): + cursor.fail() + assert telemetry_client.export_latency_log.call_count == 2 + kwargs = telemetry_client.export_latency_log.call_args.kwargs + assert kwargs["session_id"] == "session-hex" + assert kwargs["sql_statement_id"] == "statement-id" + assert kwargs["sql_execution_event"].statement_type == StatementType.QUERY + assert len(caplog.records) == (2 if debug else 0) + + +def test_object_without_connection_is_not_exported(telemetry_client, no_debug): + class Handler: + @log_latency(StatementType.QUERY) + def run(self): + return "done" + + assert Handler().run() == "done" + telemetry_client.export_latency_log.assert_not_called() + + +def test_wrapper_keeps_the_method_metadata(): + assert Cursor.execute.__name__ == "execute" + assert Cursor.execute.__wrapped__(Cursor(False), 3) == 3 diff --git a/tests/unit/test_row.py b/tests/unit/test_row.py new file mode 100644 index 000000000..fc71ac4f9 --- /dev/null +++ b/tests/unit/test_row.py @@ -0,0 +1,50 @@ +import pickle + +import pytest + +from databricks.sql.types import Row, _create_row + + +def test_row_factory_builds_named_rows(): + Person = Row("name", "age") + row = Person("Alice", 11) + assert type(row) is Row + assert row == ("Alice", 11) + assert row.__fields__ is Person + assert row.__dict__ == {"__fields__": Person} + assert row.name == "Alice" and row["age"] == 11 and row[0] == "Alice" + assert row.asDict() == {"name": "Alice", "age": 11} + assert repr(row) == "Row(name='Alice', age=11)" + assert "name" in row and "wrong" not in row + assert row == Row(name="Alice", age=11) + + +def test_row_factory_accepts_fewer_values_and_rejects_more(): + Person = Row("name", "age") + assert Person("Alice") == ("Alice",) + assert Person() == () + with pytest.raises(ValueError, match="expected 2 values but got"): + Person("Alice", 11, "extra") + + +def test_rows_stay_read_only(): + row = Row("a")(1) + with pytest.raises(RuntimeError, match="read-only"): + row.a = 2 + + +def test_create_row_matches_the_factory(): + fields = ["a", "b"] + row = _create_row(fields, [1, 2]) + assert type(row) is Row + assert row == (1, 2) + assert row.__fields__ is fields + assert row.asDict() == {"a": 1, "b": 2} + + +def test_pickle_round_trip(): + row = Row("a", "b")(1, None) + copy = pickle.loads(pickle.dumps(row)) + assert copy == row and type(copy) is Row + assert copy.asDict() == {"a": 1, "b": None} + assert repr(copy) == repr(row) From 8fedb5f2acbbea3c76fd5d01f4f054d018493347 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Sat, 14 Nov 2026 21:29:45 +0000 Subject: [PATCH 2/4] test: don't assert the measured latency in the debug-log test Match "execute completed in ms" instead of 0ms, so a slow CI runner can't fail it, and say in the log_latency docstring and comment that nothing is timed when telemetry and DEBUG logging are both off. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/telemetry/latency_logger.py | 11 ++++++----- tests/unit/test_latency_logger.py | 4 +++- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/telemetry/latency_logger.py b/src/databricks/sql/telemetry/latency_logger.py index 2bb86cdb9..4baf478ea 100644 --- a/src/databricks/sql/telemetry/latency_logger.py +++ b/src/databricks/sql/telemetry/latency_logger.py @@ -155,7 +155,9 @@ def log_latency(statement_type: StatementType = StatementType.NONE): This decorator measures the execution time of a method and sends telemetry data about the operation, including latency, statement information, and - execution context. + execution context. When the connection's telemetry is off and DEBUG logging + is disabled, nothing would record the latency, so the method is called + without timing it. Args: statement_type (StatementType): The type of SQL statement being executed. @@ -177,15 +179,14 @@ def execute(self, query): def decorator(func): @functools.wraps(func) def wrapper(self, *args, **kwargs): - # Fast check: use cached telemetry_enabled flag from connection - # Avoids dictionary lookup + instance check on every operation + # Use the connection's cached telemetry_enabled flag. When it is off + # and DEBUG logging is disabled (always the case on kernel + # connections, which export telemetry from the kernel), skip timing. connection = getattr(self, "connection", None) telemetry_enabled = bool(connection) and getattr( connection, "telemetry_enabled", False ) if not telemetry_enabled and not logger.isEnabledFor(logging.DEBUG): - # Nothing would record the latency (always the case on kernel - # connections, which export telemetry from the kernel). return func(self, *args, **kwargs) start_time = time.monotonic() diff --git a/tests/unit/test_latency_logger.py b/tests/unit/test_latency_logger.py index c6b8f8a58..ba02ea567 100644 --- a/tests/unit/test_latency_logger.py +++ b/tests/unit/test_latency_logger.py @@ -1,4 +1,5 @@ import logging +import re from unittest.mock import MagicMock, patch import pytest @@ -66,7 +67,8 @@ def test_debug_logging_records_latency_without_telemetry(telemetry_client, caplo cursor = Cursor(telemetry_enabled=False) with caplog.at_level(logging.DEBUG, logger=LOGGER): assert cursor.execute(1) == 1 - assert [r.getMessage() for r in caplog.records] == ["execute completed in 0ms"] + assert len(caplog.records) == 1 + assert re.fullmatch(r"execute completed in \d+ms", caplog.records[0].getMessage()) telemetry_client.export_latency_log.assert_not_called() From c503ab0d3f1373f8027464314d3491f74f6fa6c9 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Sat, 14 Nov 2026 21:32:36 +0000 Subject: [PATCH 3/4] docs: note that log_latency reads telemetry_enabled once, before the call Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/telemetry/latency_logger.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/telemetry/latency_logger.py b/src/databricks/sql/telemetry/latency_logger.py index 4baf478ea..b2b53cf13 100644 --- a/src/databricks/sql/telemetry/latency_logger.py +++ b/src/databricks/sql/telemetry/latency_logger.py @@ -179,7 +179,8 @@ def execute(self, query): def decorator(func): @functools.wraps(func) def wrapper(self, *args, **kwargs): - # Use the connection's cached telemetry_enabled flag. When it is off + # Use the connection's cached telemetry_enabled flag (set in + # Connection.__init__), read once before the call. When it is off # and DEBUG logging is disabled (always the case on kernel # connections, which export telemetry from the kernel), skip timing. connection = getattr(self, "connection", None) From d8212f32369f28d59a3a5ef9ccabf5473635b940 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Sat, 14 Nov 2026 21:36:55 +0000 Subject: [PATCH 4/4] test: pin that log_latency doesn't time calls nothing records Patch the module's clock and assert monotonic() is not called with telemetry and DEBUG logging off (or without a connection), and is called when either is on. The fake clock also makes the logged and exported latency exact. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- tests/unit/test_latency_logger.py | 32 +++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_latency_logger.py b/tests/unit/test_latency_logger.py index ba02ea567..30065c7f8 100644 --- a/tests/unit/test_latency_logger.py +++ b/tests/unit/test_latency_logger.py @@ -1,5 +1,5 @@ +import itertools import logging -import re from unittest.mock import MagicMock, patch import pytest @@ -42,6 +42,14 @@ def telemetry_client(): yield client +@pytest.fixture +def monotonic(): + # Each call advances one second, so every timed call measures exactly 1000 ms. + with patch("databricks.sql.telemetry.latency_logger.time") as fake_time: + fake_time.monotonic.side_effect = itertools.count(0.0, 1.0) + yield fake_time.monotonic + + @pytest.fixture def no_debug(): logger = logging.getLogger(LOGGER) @@ -51,51 +59,59 @@ def no_debug(): logger.setLevel(level) -def test_disabled_telemetry_without_debug_only_calls_the_method( - telemetry_client, no_debug, caplog +def test_disabled_telemetry_without_debug_does_not_time_the_call( + telemetry_client, monotonic, no_debug, caplog ): cursor = Cursor(telemetry_enabled=False) with caplog.at_level(logging.INFO, logger=LOGGER): assert cursor.execute(42) == 42 with pytest.raises(ValueError, match="boom"): cursor.fail() + monotonic.assert_not_called() telemetry_client.export_latency_log.assert_not_called() assert not caplog.records -def test_debug_logging_records_latency_without_telemetry(telemetry_client, caplog): +def test_debug_logging_records_latency_without_telemetry( + telemetry_client, monotonic, caplog +): cursor = Cursor(telemetry_enabled=False) with caplog.at_level(logging.DEBUG, logger=LOGGER): assert cursor.execute(1) == 1 - assert len(caplog.records) == 1 - assert re.fullmatch(r"execute completed in \d+ms", caplog.records[0].getMessage()) + assert monotonic.call_count == 2 + assert [r.getMessage() for r in caplog.records] == ["execute completed in 1000ms"] telemetry_client.export_latency_log.assert_not_called() @pytest.mark.parametrize("debug", [False, True]) def test_enabled_telemetry_exports_latency_even_on_error( - telemetry_client, no_debug, debug, caplog + telemetry_client, monotonic, no_debug, debug, caplog ): cursor = Cursor(telemetry_enabled=True) with caplog.at_level(logging.DEBUG if debug else logging.INFO, logger=LOGGER): assert cursor.execute("ok") == "ok" with pytest.raises(ValueError, match="boom"): cursor.fail() + assert monotonic.call_count == 4 assert telemetry_client.export_latency_log.call_count == 2 kwargs = telemetry_client.export_latency_log.call_args.kwargs + assert kwargs["latency_ms"] == 1000 assert kwargs["session_id"] == "session-hex" assert kwargs["sql_statement_id"] == "statement-id" assert kwargs["sql_execution_event"].statement_type == StatementType.QUERY assert len(caplog.records) == (2 if debug else 0) -def test_object_without_connection_is_not_exported(telemetry_client, no_debug): +def test_object_without_connection_is_not_exported( + telemetry_client, monotonic, no_debug +): class Handler: @log_latency(StatementType.QUERY) def run(self): return "done" assert Handler().run() == "done" + monotonic.assert_not_called() telemetry_client.export_latency_log.assert_not_called()