perf(kernel): cut per-query Python work for small results

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.

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:

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.

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()``;