diff --git a/CHANGELOG.md b/CHANGELOG.md index e62784f4c..867be71ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - 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: 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. # 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/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index df52f34a5..a916e007c 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -45,7 +45,7 @@ wrap_kernel_exception as _wrap_kernel_exception, ) from databricks.sql.backend.kernel.auth_bridge import kernel_auth_kwargs -from databricks.sql.backend.kernel.result_set import KernelResultSet +from databricks.sql.backend.kernel.result_set import KernelResultSet, _SchemaCache from databricks.sql.backend.kernel.type_mapping import bind_tspark_params from databricks.sql.backend.types import ( BackendType, @@ -90,10 +90,9 @@ _STAGING_VERBS = ("PUT", "GET", "REMOVE") -def _strip_leading_sql_comments(sql: str) -> str: - """Strip leading whitespace and SQL comments (``-- …`` line and - ``/* … */`` block, possibly several) from ``sql``, returning the - remainder. +def _sql_start(sql: str) -> int: + """Index of the first character of ``sql`` after leading whitespace and + SQL comments (``-- …`` line and ``/* … */`` block, possibly several). Needed so staging detection sees the real leading verb: a comment-prefixed staging op (``-- upload\\nPUT …`` or @@ -117,7 +116,7 @@ def _strip_leading_sql_comments(sql: str) -> str: i = n if close == -1 else close + 2 else: break - return sql[i:] + return i def _is_not_found(exc: BaseException) -> bool: @@ -164,9 +163,11 @@ def _is_staging_statement(operation: str) -> bool: only — so a normal query that merely *contains* the word (e.g. ``SELECT 'GET' AS x``) isn't misflagged. """ - stripped = _strip_leading_sql_comments(operation) - # First whitespace-delimited token, uppercased. - verb = stripped.split(None, 1)[0].upper() if stripped.strip() else "" + start = _sql_start(operation) + # First whitespace-delimited token, uppercased. The verbs have at most 6 + # characters, so the first 7 decide it without copying the statement. + head = operation[start : start + 7] + verb = head.split(None, 1)[0].upper() if head.strip() else "" return verb in _STAGING_VERBS @@ -436,6 +437,7 @@ def __init__( # Guarded by its own lock — cancel can race execute teardown. self._sync_cancellers: Dict[int, Any] = {} self._sync_cancellers_lock = threading.RLock() + self._schema_cache = _SchemaCache() # ── Session lifecycle ────────────────────────────────────────── @@ -822,6 +824,11 @@ def cancel_running_cursor(self, cursor: "Cursor") -> bool: return True def close_command(self, command_id: CommandId) -> None: + guid = command_id.guid + if guid not in self._async_handles and guid not in self._async_statements: + # Sync and metadata results are never tracked, so there is nothing + # to close or clean up. + return with self._async_handles_lock: handle = self._async_handles.pop(command_id.guid, None) stmt = self._async_statements.pop(command_id.guid, None) @@ -1056,6 +1063,7 @@ def _make_result_set( command_id=command_id, arraysize=cursor.arraysize, buffer_size_bytes=cursor.buffer_size_bytes, + schema_cache=self._schema_cache, ) def _synthetic_command_id(self) -> CommandId: diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py index fcb798158..c008f6eb3 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/src/databricks/sql/backend/kernel/result_set.py @@ -32,7 +32,7 @@ import logging from collections import deque -from typing import Any, Deque, List, Optional, Tuple, TYPE_CHECKING, cast +from typing import Any, Deque, List, Optional, TYPE_CHECKING, cast import pyarrow @@ -119,36 +119,76 @@ def _max_rows_without_pandas(types: List[pyarrow.DataType]) -> int: return min(_MAX_ROWS_WITHOUT_PANDAS, _MAX_COST_WITHOUT_PANDAS // total) +class _ResultSchema: + """What a result set derives from its Arrow schema.""" + + __slots__ = ( + "schema", + "description", + "row_factory", + "max_rows_without_pandas", + "float_columns", + ) + + def __init__(self, schema: pyarrow.Schema): + self.schema = schema + self.description = tuple(description_from_arrow_schema(schema)) + self.row_factory = Row(*[c[0] for c in self.description]) + types = schema.types + self.max_rows_without_pandas = _max_rows_without_pandas(types) + self.float_columns = tuple( + i for i, t in enumerate(types) if t.id in _FLOAT_TYPE_IDS + ) + + +class _SchemaCache: + """Keeps the last ``_ResultSchema`` of a connection, since repeated + statements return equal schemas.""" + + __slots__ = ("_last",) + + def __init__(self) -> None: + self._last: Optional[_ResultSchema] = None + + def get(self, schema: pyarrow.Schema) -> _ResultSchema: + last = self._last + if last is not None and schema.equals(last.schema, check_metadata=True): + return last + last = self._last = _ResultSchema(schema) + return last + + def _rows_without_pandas( - table: pyarrow.Table, description: List[Tuple] + batches: List[pyarrow.RecordBatch], result_schema: _ResultSchema ) -> Optional[List[Row]]: - """The rows of ``table`` built from ``to_pylist()``, or None if the table - must go through pandas (too costly, or a column whose values could - differ).""" - num_rows = table.num_rows - # No weight is below 1, so this rules out larger results without - # looking at every column type. - if num_rows * table.num_columns > _MAX_COST_WITHOUT_PANDAS: - return None - if num_rows > _max_rows_without_pandas(table.schema.types): - return None - columns = [] - for column in table.itercolumns(): - try: - values = column.to_pylist() - except Exception: - # E.g. a date past 9999-12-31: pandas then raises its own error, - # as it does without this fast path. - logger.debug( - "to_pylist() failed for a %s column; using pandas", - column.type, - exc_info=True, - ) + """The rows of ``batches`` built from ``to_pylist()``, or None if they + must go through a Table and pandas (too costly, a column whose values + could differ, or a batch whose schema does not match).""" + num_rows = 0 + for batch in batches: + if not batch.schema.equals(result_schema.schema): return None - if column.type.id in _FLOAT_TYPE_IDS and any(v != v for v in values): + num_rows += batch.num_rows + if num_rows > result_schema.max_rows_without_pandas: + return None + columns: List[List[Any]] = [[] for _ in result_schema.description] + for batch in batches: + for values, column in zip(columns, batch.columns): + try: + values.extend(column.to_pylist()) + except Exception: + # E.g. a date past 9999-12-31: pandas then raises its own + # error, as it does without this fast path. + logger.debug( + "to_pylist() failed for a %s column; using pandas", + column.type, + exc_info=True, + ) + return None + for i in result_schema.float_columns: + if any(v != v for v in columns[i]): return None - columns.append(values) - row_factory = Row(*[c[0] for c in description]) + row_factory = result_schema.row_factory return [row_factory(*values) for values in zip(*columns)] @@ -170,11 +210,16 @@ def __init__( command_id: CommandId, arraysize: int, buffer_size_bytes: int, + schema_cache: Optional[_SchemaCache] = None, ): try: schema = kernel_handle.arrow_schema() except Exception as exc: raise wrap_kernel_exception("KernelResultSet.arrow_schema", exc) from exc + if schema_cache is None: + self._result_schema = _ResultSchema(schema) + else: + self._result_schema = schema_cache.get(schema) super().__init__( connection=connection, backend=backend, @@ -185,7 +230,7 @@ def __init__( has_been_closed_server_side=False, has_more_rows=True, results_queue=None, - description=description_from_arrow_schema(schema), + description=list(self._result_schema.description), is_staging_operation=False, lz4_compressed=False, arrow_schema_bytes=None, @@ -237,8 +282,14 @@ def _ensure_buffered(self, n_rows: int) -> int: def _buffered_rows(self) -> int: return self._buffered_count + def _to_table(self, batches: List[pyarrow.RecordBatch]) -> pyarrow.Table: + return pyarrow.Table.from_batches(batches, schema=self._schema) + def _take_buffered(self, n: int) -> pyarrow.Table: """Slice up to ``n`` rows out of the buffer; advances state.""" + return self._to_table(self._take_buffered_batches(n)) + + def _take_buffered_batches(self, n: int) -> List[pyarrow.RecordBatch]: slices: List[pyarrow.RecordBatch] = [] remaining = n while remaining > 0 and self._buffer: @@ -254,13 +305,14 @@ def _take_buffered(self, n: int) -> pyarrow.Table: taken = n - remaining self._buffered_count -= taken self._next_row_index += taken - if not slices: - return pyarrow.Table.from_batches([], schema=self._schema) - return pyarrow.Table.from_batches(slices, schema=self._schema) + return slices def _drain(self) -> pyarrow.Table: """Consume everything left in the buffer + kernel stream and return as a single Table.""" + return self._to_table(self._drain_batches()) + + def _drain_batches(self) -> List[pyarrow.RecordBatch]: chunks: List[pyarrow.RecordBatch] = [] if self._buffer and self._buffer_offset > 0: head = self._buffer.popleft() @@ -286,9 +338,7 @@ def _drain(self) -> pyarrow.Table: rows = sum(c.num_rows for c in chunks) self._buffered_count = 0 self._next_row_index += rows - if not chunks: - return pyarrow.Table.from_batches([], schema=self._schema) - return pyarrow.Table.from_batches(chunks, schema=self._schema) + return chunks # ----- Arrow fetches ----- @@ -305,19 +355,20 @@ def fetchmany_arrow(self, size: int) -> pyarrow.Table: # ----- Row fetches ----- - def _convert_arrow_table(self, table): + def _rows(self, batches: List[pyarrow.RecordBatch]) -> List[Row]: + # Small results skip building a Table, which releases and re-takes + # the GIL, as well as pandas. if self.connection.disable_pandas is not True: - rows = _rows_without_pandas(table, self.description) + rows = _rows_without_pandas(batches, self._result_schema) if rows is not None: return rows - return super()._convert_arrow_table(table) + return self._convert_arrow_table(self._to_table(batches)) def fetchone(self) -> Optional[Row]: self._ensure_buffered(1) if self._buffered_rows() == 0: return None - table = self._take_buffered(1) - rows = self._convert_arrow_table(table) + rows = self._rows(self._take_buffered_batches(1)) return rows[0] if rows else None def fetchmany(self, size: int) -> List[Row]: @@ -326,11 +377,10 @@ def fetchmany(self, size: int) -> List[Row]: if size == 0: return [] self._ensure_buffered(size) - table = self._take_buffered(size) - return self._convert_arrow_table(table) + return self._rows(self._take_buffered_batches(size)) def fetchall(self) -> List[Row]: - return self._convert_arrow_table(self._drain()) + return self._rows(self._drain_batches()) def close(self) -> None: """Close the underlying kernel handle and notify the backend. diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index f94aef0c5..b96a35be4 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1066,6 +1066,15 @@ def test_staging_operation_raises_not_supported(operation): ("INSERT INTO t VALUES (1)", False), ("", False), ("-- just a comment", False), # comment only, no statement + # Only the leading token counts, however long it is. + ("PUT", True), + ("remove\t'/v'", True), + ("GET/* c */ '/v' TO '/f'", False), + ("PUTX '/f'", False), + ("REMOVED '/v'", False), + ("REMOVEX", False), + ("GETTING_STARTED", False), + (" -- c\n" + "x" * 10_000, False), ], ) def test_is_staging_statement(operation, is_staging): @@ -1285,6 +1294,22 @@ def test_close_command_tolerant_when_handle_missing(): c.close_command(fake_command_id) # must not raise +def test_close_command_closes_a_statement_tracked_without_its_handle(): + c = _make_client() + c._async_handles_lock = MagicMock() + untracked = CommandId.from_sea_statement_id("sync") + c.close_command(untracked) + # Sync results were never tracked, so closing them takes no lock. + c._async_handles_lock.__enter__.assert_not_called() + + stmt = MagicMock() + cid = CommandId.from_sea_statement_id("async") + c._async_statements[cid.guid] = stmt + c.close_command(cid) + stmt.close.assert_called_once() + assert cid.guid not in c._async_statements + + def _attach_returns(c, *, status=None, status_error=None, await_result=None): """Wire ``_kernel_session.attach_async_statement`` to return a fake handle. ``status`` is a ``(state, failure)`` tuple for ``handle.status()``; diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index aafb45679..beb98efad 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -49,7 +49,7 @@ def close(self): self.closed = True -def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet: +def _make_rs(handle, *, disable_pandas=False, schema_cache=None) -> KernelResultSet: # The base ResultSet __init__ takes a `connection` ref it never # actually dereferences during these buffer tests, so a Mock is # fine. @@ -63,6 +63,7 @@ def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet: command_id=CommandId.from_sea_statement_id("smoke-test"), arraysize=100, buffer_size_bytes=1024, + schema_cache=schema_cache, ) @@ -379,7 +380,7 @@ def _convert(table): with patch.object( ResultSet, "_convert_arrow_table", autospec=True, side_effect=pandas_path ) as spy: - rows = rs._convert_arrow_table(table) + rows = rs._rows(table.to_batches()) return rows, pandas_path(rs, table), spy.called @@ -474,7 +475,7 @@ def test_out_of_range_values_convert_as_on_the_pandas_path(column, expected): table = pa.table({"c": column}) rs = _make_rs(_FakeKernelHandle(table.schema, [])) for convert in ( - rs._convert_arrow_table, + lambda t: rs._rows(t.to_batches()), lambda t: ResultSet._convert_arrow_table(rs, t), ): if isinstance(expected, list): @@ -489,7 +490,7 @@ def test_conversion_errors_are_logged_before_falling_back(caplog): rs = _make_rs(_FakeKernelHandle(table.schema, [])) with caplog.at_level(logging.DEBUG, logger=kernel_result_set.__name__): with pytest.raises(ValueError): - rs._convert_arrow_table(table) + rs._rows(table.to_batches()) assert "to_pylist() failed for a date32[day] column" in caplog.text @@ -517,3 +518,138 @@ def test_fetch_methods_return_the_pandas_path_values(): _assert_same_rows(_make_rs(_FakeKernelHandle(schema, [batch])).fetchall(), expected) rs = _make_rs(_FakeKernelHandle(schema, [batch])) _assert_same_rows([rs.fetchone()] + rs.fetchmany(5), expected) + + +# ----- per-connection schema cache ----- + + +def _reimported(schema: pa.Schema) -> pa.Schema: + """An equal but distinct schema object, as each kernel statement returns.""" + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, schema): + pass + return pa.ipc.open_stream(sink.getvalue()).schema + + +def test_equal_schemas_reuse_what_is_derived_from_them(): + schema = pa.schema( + [ + pa.field("n", pa.int64()), + pa.field("v", pa.string(), metadata={b"databricks.type_name": b"VARIANT"}), + ] + ) + cache = kernel_result_set._SchemaCache() + first = _make_rs(_FakeKernelHandle(_reimported(schema), []), schema_cache=cache) + second = _make_rs(_FakeKernelHandle(_reimported(schema), []), schema_cache=cache) + assert first._result_schema is second._result_schema + assert ( + first.description + == second.description + == _make_rs(_FakeKernelHandle(schema, [])).description + ) + assert [c[1] for c in first.description] == ["bigint", "variant"] + # Each result set still gets its own description list. + first.description.append("changed") + assert len(second.description) == 2 + + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2]), pa.array(['{"a":1}', None])], schema=schema + ) + rows = _make_rs(_FakeKernelHandle(schema, [batch]), schema_cache=cache).fetchall() + assert [tuple(r) for r in rows] == [(1, '{"a":1}'), (2, None)] + assert rows[0].v == '{"a":1}' + + +@pytest.mark.parametrize( + "other", + [ + pa.schema([pa.field("n", pa.int32())]), + pa.schema([pa.field("m", pa.int64())]), + pa.schema( + [pa.field("n", pa.int64(), metadata={b"databricks.type_name": b"MAP"})] + ), + ], +) +def test_a_different_schema_is_derived_again(other): + cache = kernel_result_set._SchemaCache() + first = _make_rs( + _FakeKernelHandle(pa.schema([pa.field("n", pa.int64())]), []), + schema_cache=cache, + ) + second = _make_rs(_FakeKernelHandle(other, []), schema_cache=cache) + assert first._result_schema is not second._result_schema + assert second.description == _make_rs(_FakeKernelHandle(other, [])).description + assert first.description != second.description + + +# ----- rows built from record batches ----- + + +def _spy_to_table(): + return patch.object( + KernelResultSet, + "_to_table", + autospec=True, + side_effect=KernelResultSet._to_table, + ) + + +def test_small_results_are_converted_from_batches_without_a_table(): + schema = pa.schema([("n", pa.int64()), ("s", pa.string())]) + + def batch(lo, hi): + return pa.RecordBatch.from_arrays( + [pa.array(range(lo, hi)), pa.array([str(i) for i in range(lo, hi)])], + schema=schema, + ) + + expected = [(i, str(i)) for i in range(7)] + with _spy_to_table() as to_table: + rs = _make_rs( + _FakeKernelHandle(schema, [batch(0, 3), batch(3, 5), batch(5, 7)]) + ) + first = rs.fetchone() + rows = [first] + rs.fetchmany(3) + rs.fetchall() + to_table.assert_not_called() + assert [tuple(r) for r in rows] == expected + assert rows[-1].s == "6" and rs.rownumber == 7 + + limit = kernel_result_set._max_rows_without_pandas(schema.types) + with _spy_to_table() as to_table: + rs = _make_rs( + _FakeKernelHandle(schema, [batch(0, limit), batch(limit, limit + 1)]) + ) + assert len(rs.fetchall()) == limit + 1 + to_table.assert_called_once() + + +def test_a_batch_whose_schema_differs_still_fails_as_before(): + schema = pa.schema([("n", pa.int64())]) + other = pa.RecordBatch.from_arrays([pa.array(["x"])], names=["n"]) + with pytest.raises(pa.ArrowInvalid): + _make_rs(_FakeKernelHandle(schema, [other])).fetchall() + + +def test_the_cost_budget_is_computed_once_per_schema_not_per_fetch(): + schema = pa.schema([("n", pa.int64())]) + cache = kernel_result_set._SchemaCache() + with patch.object( + kernel_result_set, + "_max_rows_without_pandas", + wraps=kernel_result_set._max_rows_without_pandas, + ) as budget: + for _ in range(3): + rs = _make_rs( + _FakeKernelHandle( + _reimported(schema), + [_batch(schema, [1, 2]), _batch(schema, [3, 4])], + ), + schema_cache=cache, + ) + rows = [rs.fetchone()] + rs.fetchmany(2) + rs.fetchall() + assert [row[0] for row in rows] == [1, 2, 3, 4] + assert budget.call_count == 1 + _make_rs( + _FakeKernelHandle(pa.schema([("m", pa.int64())]), []), schema_cache=cache + ) + assert budget.call_count == 2