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..b2b53cf13 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,6 +179,17 @@ def execute(self, query): def decorator(func): @functools.wraps(func) def wrapper(self, *args, **kwargs): + # 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) + telemetry_enabled = bool(connection) and getattr( + connection, "telemetry_enabled", False + ) + if not telemetry_enabled and not logger.isEnabledFor(logging.DEBUG): + return func(self, *args, **kwargs) + start_time = time.monotonic() try: return func(self, *args, **kwargs) @@ -186,10 +199,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..30065c7f8 --- /dev/null +++ b/tests/unit/test_latency_logger.py @@ -0,0 +1,120 @@ +import itertools +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 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) + level = logger.level + logger.setLevel(logging.INFO) + yield + logger.setLevel(level) + + +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, monotonic, caplog +): + cursor = Cursor(telemetry_enabled=False) + with caplog.at_level(logging.DEBUG, logger=LOGGER): + assert cursor.execute(1) == 1 + 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, 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, 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() + + +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)