chore(bench): kernel-vs-Thrift performance baseline harness + results by vikrantpuppala · Pull Request #790 · databricks/databricks-sql-python

vikrantpuppala · 2026-05-15T09:52:31Z

Stacks on #789 (kernel bind_param). Adds a benchmark script under `scripts/`; no production code change.

Summary

One-shot benchmark script that runs each (backend × SQL-shape) combination N+1 times against a live warehouse, drops the first run (cache warm-up), and reports min/median/max for session-open, time-to-first-row (TTFR), drain, and RSS delta.

Not a CI gate — single-machine, single-warehouse, high-variance. Meant to be re-run by hand when we want a baseline. Output is a Markdown table you can paste into a PR.

```
set -a && source ~/.databricks/pecotesting-creds && set +a
export DATABRICKS_SERVER_HOSTNAME=${DATABRICKS_HOST#https://}
.venv/bin/python scripts/bench_kernel_vs_thrift.py
```

Results (median of 5 samples, warm-up dropped, dogfood)

Shape	Thrift drain	Kernel drain	Ratio	Notes
`SELECT 1`	387 ms	1085 ms	2.8×	Fixed ~700ms kernel TTFR overhead
`range(10k)`	909 ms	1347 ms	1.48×	Overhead amortized
`wide-uuid(100k)`	6907 ms	9925 ms	1.44×	Overhead amortized
`metadata.catalogs`	413 ms	550 ms	1.33×	Overhead amortized
`range(1M)`	14.0 s	panic	—	kernel-side bug: issue #19

Detailed tables with min/max ranges and RSS-delta numbers are in the script output.

Three findings worth flagging

1. Fixed kernel TTFR overhead (~700ms)

`SELECT 1` is the cleanest signal because drain time is essentially zero. Kernel pays ~700ms more than Thrift on every query. On large queries (drain dominates) the relative cost shrinks to 1.3–1.5×.

Plausible causes (not investigated):

Per-Session tokio runtime construction (flagged by the kernel PR review).
SEA wait-for-result poll handshake doing an extra round-trip.

A flamegraph would distinguish these. Worth a follow-up.

2. CloudFetch panic on large results (kernel issue #19)

`range(1M)` crosses the CloudFetch threshold; the kernel's reader does `runtime_handle.block_on` from a sync trait method, which panics when called from inside our PyO3 `runtime.block_on`. `use_sea=True` is unusable in production for any large-result workload until this is fixed. The connector's e2e tests use `range(10000)` which is below the CloudFetch threshold, which is why it didn't surface earlier.

3. RSS overhead ~+1MB per kernel session

Consistent across every shape. Probably tokio worker thread stacks (default ~2MB × N workers, partially committed). Not a problem at small connection counts; maps to the "process-global `OnceLock`" follow-up the kernel reviewer flagged.

What this doesn't measure

Concurrent-connection scaling (single conn at a time here)
Anything past 100k rows on kernel (blocked by TLS version handshake issue #19)
Network-jitter sensitivity
Memory profile beyond coarse RSS

Useful next benchmarks once #19 lands: real CloudFetch shape (1M+ rows), concurrent sessions, and a memory-profile pass.

This pull request and its description were written by Isaac.

Phase 2 of the PySQL × kernel integration plan (databricks-sql-kernel/docs/designs/pysql-kernel-integration.md). Wires `use_sea=True` to a new `backend/kernel/` module that delegates to the Rust kernel via the `databricks_sql_kernel` PyO3 extension (kernel PR #13). New module: `src/databricks/sql/backend/kernel/` - `client.py` — `KernelDatabricksClient(DatabricksClient)`. Lazy- imports `databricks_sql_kernel` so a connector install without the kernel wheel doesn't `ImportError` at startup; only `use_sea=True` surfaces the missing-extra message. Implements open/close_session, sync + async execute_command (async_op=True goes through `Statement.submit()` and stashes the handle in a dict keyed on `CommandId`), cancel/close_command, get_query_state, get_execution_result, and the metadata calls (catalogs / schemas / tables / columns) via `Session.metadata().list_*`. Real server-issued session and statement IDs flow through (no synthetic UUIDs). - `auth_bridge.py` — translate the connector's `AuthProvider` into kernel `Session` kwargs. PAT (including federation-wrapped PAT — `get_python_sql_connector_auth_provider` always wraps the base in `TokenFederationProvider`, so a naive isinstance check never matches) routes through `auth_type="pat"`. Everything else routes through `auth_type="external"` with a callback that delegates to `auth_provider.add_headers({})`. (External today is rejected by the kernel at `build_auth_provider`; the separate kernel-side enablement PR will flip it on.) - `result_set.py` — `KernelResultSet(ResultSet)`. Duck-typed over `databricks_sql_kernel.ExecutedStatement` (sync execute) and `ResultStream` (metadata + async await_result) since both expose `arrow_schema()` / `fetch_next_batch()` / `fetch_all_arrow()` / `close()`. Same FIFO batch buffer the prior ADBC POC used, so `fetchmany(n)` for n smaller than the kernel's natural batch size doesn't re-fetch. - `type_mapping.py` — Arrow → PEP 249 description-string mapper. Lifted from the prior ADBC POC; centralised here so future kernel-result wrappers reuse the same mapping. Kernel errors → PEP 249 exceptions: `KernelError.code` is mapped in a single table to `ProgrammingError` / `OperationalError` / `DatabaseError`. The structured fields (`sql_state`, `error_code`, `query_id`, …) are copied onto the re-raised exception so callers can branch on them without reaching through `__cause__`. Routing: `Session._create_backend` flips the `use_sea=True` branch to instantiate `KernelDatabricksClient` instead of the native `SeaDatabricksClient`. The native `backend/sea/` module is left in place (no users on `use_sea=True` after this PR; its long- term fate is out of scope here). Packaging: `[tool.poetry.extras] kernel = ["databricks-sql-kernel"]`. `pip install 'databricks-sql-connector[kernel]'` pulls in the kernel wheel; `use_sea=True` without the extra raises a pointed ImportError telling the user how to install it. Known gaps (acknowledged, will be follow-ups): - Parameter binding (`execute_command(parameters=[...])`) raises NotSupportedError — PyO3 `Statement.bind_param` lands in a follow-up. - Statement-level `query_tags` raises NotSupportedError. - `get_tables(table_types=[...])` returns unfiltered rows (the native SEA backend's filter is keyed on `SeaResultSet`; needs a small port to operate on `KernelResultSet`). - External-auth end-to-end blocked on the kernel-side `AuthConfig::External` enablement PR. - Volume PUT/GET (staging operations): kernel has no Volume API. Test plan: - Unit: 37 new tests across `tests/unit/test_kernel_auth_bridge.py` (auth provider → kwargs mapping, including federation-wrapped PAT and the External trampoline call-counter check), `tests/unit/test_kernel_type_mapping.py` (Arrow type mapping + description shape), and `tests/unit/test_kernel_result_set.py` (buffer semantics, fetchmany across batch boundaries, idempotent close, close() swallowing handle-close failures). All pass. - Full unit suite: 600 pre-existing tests still pass; one pre-existing failure (`test_useragent_header` — agent detection adds `agent/claude-code` in this env) was already failing on main, unrelated to this change. - Live e2e against dogfood with `use_sea=True`: SELECT 1, `range(10000)`, `fetchmany` pacing, `fetchall_arrow`, all four metadata calls (returned 75 catalogs / 144 schemas in main / 47 tables in `system.information_schema` / 15 columns), `session_configuration={'ANSI_MODE': 'false'}` round-trips, bad SQL surfaces as DatabaseError with `code='SqlError'` and `sql_state='42P01'` on the exception. All checks pass. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

The earlier auth_bridge routed OAuth/MSAL/federation through the kernel's External token-provider trampoline (a Python callable the kernel invoked per HTTP request). Removing that for now. Why: routing OAuth into the kernel inherently requires per-request token resolution to keep refresh working during a long-running session. Two viable mechanisms (kernel-native OAuth, or the External callback); both have costs (duplicate OAuth flows vs GIL-per-request). Punting the decision until there's actual demand on use_sea=True. Today: the bridge accepts PAT (including TokenFederationProvider- wrapped PAT, which is how `get_python_sql_connector_auth_provider` always shapes it). Any non-PAT auth_provider raises a clear NotSupportedError pointing the user at use_sea=False (Thrift). This shrinks the auth_bridge to ~50 lines and means the kernel- side External enablement PR is no longer on the connector's critical path — there's no kernel-side prerequisite for shipping use_sea=True for PAT users. Unit tests updated: - TokenFederationProvider-wrapped PAT still routes to PAT (kept). - Generic OAuth provider raises NotSupportedError (new). - ExternalAuthProvider raises NotSupportedError (new). - Silent non-PAT provider raises NotSupportedError (new) — reject the type itself rather than trying to extract a token we already know we can't use. Live e2e against dogfood with use_sea=True (PAT): all checks still pass (SELECT 1, range(10000), fetchmany pacing, four metadata calls, session_configuration round-trip, structured DatabaseError on bad SQL). Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Moves the previously-ad-hoc /tmp/connector_smoke.py into the repo as a real pytest module under tests/e2e/ — same convention as the rest of the e2e suite. Uses the existing session-scoped `connection_details` fixture from the top-level conftest so it shares the credential surface with every other live test. 11 tests cover: - connect() with use_sea=True opens a session. - SELECT 1: rows + description shape (column name + dbapi type slug). - SELECT * FROM range(10000): multi-batch drain. - fetchmany() pacing across the buffer boundary. - fetchall_arrow() returns a pyarrow Table. - All four metadata methods (catalogs / schemas / tables / columns). - session_configuration={'ANSI_MODE': 'false'} round-trips. - Bad SQL surfaces as DatabaseError with `code='SqlError'` and `sql_state='42P01'` attached as exception attributes. Module-level skips: - `databricks_sql_kernel` not importable → whole module skipped via pytest.importorskip (the wheel hasn't been installed). - Live creds missing → fixture-level skip with a pointed message. Run: `pytest tests/e2e/test_kernel_backend.py -v`. All 11 pass against dogfood in ~20s. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

CI is failing across all jobs at \`poetry lock\` time: Because databricks-sql-connector depends on databricks-sql-kernel (^0.1.0) which doesn't match any versions, version solving failed. The kernel wheel isn't yet published to PyPI — we verified the name is available via the Databricks proxy, but the package itself hasn't been built and uploaded yet. Declaring it as a poetry dep (even an optional one inside an extra) requires the version to be resolvable, and \`poetry lock\` runs as the setup step for every CI job: unit tests, linting, type checks, all of them. Fix: drop the \`databricks-sql-kernel\` dep declaration and the \`[kernel]\` extra from pyproject.toml until the wheel is on PyPI. The lazy import in \`backend/kernel/client.py\` still raises a clear ImportError pointing at \`pip install databricks-sql-kernel\` (or local maturin) when use_sea=True is invoked without the kernel present. When the kernel is published, a small follow-up will add back: databricks-sql-kernel = {version = "^0.1.0", optional = true} [tool.poetry.extras] kernel = ["databricks-sql-kernel"] A pointed comment in pyproject.toml documents the deferred change. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Three CI failures after the poetry-lock fix uncovered three real issues: 1. pyarrow is optional in the connector. The default-deps CI test job installs without it; the +PyArrow job installs with. The kernel backend's result_set.py + type_mapping.py import pyarrow eagerly (the kernel always returns pyarrow), and the unit tests import the backend at collection time — which crashes the default-deps job at ModuleNotFoundError. Fix: gate the three kernel unit tests on `pytest.importorskip( "pyarrow")` so they skip on default-deps and run on +PyArrow. Verified locally: 39 pass with pyarrow, 3 skipped without. No change to the backend module itself — nothing imports it until use_sea=True is invoked, and pyarrow is on the kernel wheel's runtime dep list so use_sea=True can't hit this either. 2. mypy: KernelDatabricksClient.open_session returns self._session_id, which mypy types as Optional[SessionId] because the field starts as None. Fix: bind the new id to a local non-Optional variable, assign to the field, return the local. CI's check-types runs cleanly on backend/kernel/ now; pre-existing mypy noise elsewhere isn't mine. 3. black --check: black 22.12.0 (the version CI pins) wants reformatting on result_set.py / type_mapping.py / client.py. Applied. Verified locally with the same black version. All 39 kernel unit tests + 619 pre-existing unit tests pass. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

The +PyArrow CI matrix installs pyarrow but not the databricks-sql-kernel wheel (the wheel isn't on PyPI yet, and the [kernel] extra is deferred — see commit 31ca581). The previous fix gated unit tests on `pytest.importorskip("pyarrow")` but test_kernel_auth_bridge.py was still pulled into a kernel-wheel ImportError because: src/databricks/sql/backend/kernel/__init__.py -> from databricks.sql.backend.kernel.client import KernelDatabricksClient -> import databricks_sql_kernel # ImportError on +PyArrow CI The eager re-export from `__init__.py` was a convenience that broke every consumer that only needed a submodule (type_mapping, result_set, auth_bridge) — they all triggered the kernel wheel import for no reason. Fix: - Drop the eager re-export from `kernel/__init__.py`. Comment documents why and points callers (= session.py::_create_backend, already this shape) at the direct `from .client import ...`. - Drop the no-longer-needed `pytest.importorskip("pyarrow")` / `importorskip("databricks_sql_kernel")` from test_kernel_auth_bridge.py — auth_bridge.py itself has neither dep, so the test now runs on every CI matrix variant. - test_kernel_result_set.py and test_kernel_type_mapping.py keep the pyarrow importorskip because they themselves use pyarrow. Verified locally across the three matrix shapes: - both pyarrow + kernel installed: 39 pass. - pyarrow only (no kernel wheel — the +PyArrow CI shape): 39 pass. - neither: 9 pass (auth_bridge only), 2 modules skip (the others use pyarrow). Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

…sing The connector's coverage CI job runs the full e2e suite, several of whose test classes parametrize ``extra_params`` over ``{}`` and ``{"use_sea": True}``. With ``use_sea=True`` now routing through the Rust kernel via PyO3, those cases die at ``connect()`` with our pointed ImportError because the ``databricks-sql-kernel`` wheel isn't yet on PyPI — and that CI job (sensibly) doesn't try to build it from a sibling repo. Fix: ``pytest_collection_modifyitems`` hook in the top-level ``conftest.py`` that adds a ``skip`` marker to any parametrize case with ``extra_params={"use_sea": True, ...}`` when ``importlib.util.find_spec("databricks_sql_kernel")`` returns ``None``. Behavior change is CI-only — local dev with the kernel wheel installed (via ``maturin develop`` from the kernel repo) runs those cases as before. Once the kernel wheel is published, the [kernel] extra in pyproject.toml gets enabled (see comment block there) and the default-deps CI matrix will install it; the skip then becomes a no-op. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Lifts the NotSupportedError that execute_command currently raises for parametrized queries. The kernel-side PyO3 binding for Statement.bind_param landed in databricks-sql-kernel#18; this commit wires the connector's TSparkParameter shape through to it. Implementation: - New `bind_tspark_params(kernel_stmt, parameters)` in type_mapping.py forwards each TSparkParameter to the kernel as `(ordinal, value.stringValue, type)`. ordinal is the 1-based position in the parameters list; the connector's `ordinal: bool` flag is checked only to reject named bindings (kernel v0 doesn't accept them on the wire). - execute_command no longer raises on `parameters=[...]`. The query_tags branch stays — that's a separate gap. Tests: - 6 new unit tests in tests/unit/test_kernel_type_mapping.py for the mapper: - positional forwarding preserves ordering and (ordinal, value, type) - None value forwards as SQL NULL - VOID passes through verbatim (kernel parser ignores value for VOID) - named bindings raise NotSupportedError with a pointed message - missing TSparkParameter.type defaults to STRING (defensive) - empty parameters list is a no-op - 3 new e2e tests in tests/e2e/test_kernel_backend.py against dogfood: - mixed-type round-trip (INT, STRING, BOOLEAN) via the connector's native IntegerParameter/StringParameter/BooleanParameter - None parameter (VoidParameter → SQL NULL) - DECIMAL parameter with precision/scale carried in the SQL type string (auto-inferred — explicit-arg path has a pre-existing bug in native.py where format-args are swapped) Unit pass: 45/45 (was 39). Full unit suite: 625 pass. Live e2e: 14/14 (was 11). Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

One-shot benchmark script under scripts/ that runs each (backend × SQL-shape) combination N+1 times against a live warehouse, drops the first run (cache warm-up), and reports min/median/max for session-open, time-to-first-row, drain, and RSS delta. Not a CI gate — single-machine, single-warehouse, high-variance script meant to be re-run by hand when we want a baseline. Shapes: - SELECT 1 (pure round-trip latency, no data) - range(10k) (inline result, ~10K rows) - range(1M) (crosses CloudFetch threshold; currently panics on the kernel backend — see kernel issue #19, nested block_on bug) - wide-uuid(100k) (wider rows, Arrow serialization) - metadata.catalogs (metadata round-trip) Output is a Markdown table you can paste into a PR. Run with: set -a && source ~/.databricks/pecotesting-creds && set +a export DATABRICKS_SERVER_HOSTNAME=${DATABRICKS_HOST#https://} .venv/bin/python scripts/bench_kernel_vs_thrift.py Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

vikrantpuppala · 2026-05-15T10:14:56Z

Update: kernel issue #19 fixed in databricks-sql-kernel#20.

Re-ran the benchmark with the fix applied:

Shape	Thrift drain	Kernel drain (was)	Kernel drain (after fix)	Ratio
`SELECT 1`	396 ms	1085 ms	1079 ms	2.7×
`range(10k)`	1106 ms	1347 ms	1395 ms	1.26×
`range(1M)`	15009 ms	panic	7317 ms	0.49× (2× faster!)
`wide-uuid(100k)`	9776 ms	9925 ms	9781 ms	1.00×
`metadata.catalogs`	420 ms	550 ms	546 ms	1.30×

The big result: on range(1M) — the shape that previously panicked — the kernel-backed path is now 2× faster than Thrift at drain (7.3s vs 15.0s; 137K rows/s vs 67K rows/s). CloudFetch's parallel chunk download is paying off.

Other observations unchanged:

Fixed ~700ms kernel TTFR overhead on small queries (SELECT 1 is the cleanest signal). Probably per-Session tokio runtime construction; separate follow-up.
wide-uuid(100k) lines up at ~9.8s on both backends — server-side dominates here.
RSS delta is larger on kernel for the 1M shape (+21MB vs +4KB) because pyarrow holds the whole drained table in scope at end of drain. Both backends should converge if we drain batch-by-batch instead.

vikrantpuppala added 9 commits May 14, 2026 15:10

vikrantpuppala had a problem deploying to azure-prod May 15, 2026 09:52 — with GitHub Actions Failure

vikrantpuppala force-pushed the feat/kernel-bind-param branch from aa8bd24 to 19c0bb6Compare May 18, 2026 10:24

Navigation Menu

Uh oh!

chore(bench): kernel-vs-Thrift performance baseline harness + results - #790

chore(bench): kernel-vs-Thrift performance baseline harness + results#790
vikrantpuppala wants to merge 9 commits into
feat/kernel-bind-paramfrom
bench/kernel-vs-thrift

vikrantpuppala commented May 15, 2026

Uh oh!

vikrantpuppala commented May 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

vikrantpuppala commented May 15, 2026

Summary

Results (median of 5 samples, warm-up dropped, dogfood)

Three findings worth flagging

1. Fixed kernel TTFR overhead (~700ms)

2. CloudFetch panic on large results (kernel issue #19)

3. RSS overhead ~+1MB per kernel session

What this doesn't measure

Uh oh!

vikrantpuppala commented May 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant