feat(rust-pyo3): add POC PyO3 Python bindings for the kernel#423
Draft
vikrantpuppala wants to merge 2 commits intomainfrom
Draft
feat(rust-pyo3): add POC PyO3 Python bindings for the kernel#423vikrantpuppala wants to merge 2 commits intomainfrom
vikrantpuppala wants to merge 2 commits intomainfrom
Conversation
Four small kernel changes that together close the per-query gap vs the existing Thrift backend on small/medium results, with no regressions on large CloudFetch queries. 1. Skip redundant DELETE for inline-Closed statements. The SEA server returns status=Closed alongside inline result data — the statement is already cleaned up server-side, so issuing a DELETE is a wasted round-trip (~250ms). Plumb a `server_side_closed: bool` through ExecuteResult; Statement::execute_single skips registering the statement_id for cleanup when set. 2. Make Drop for Statement non-blocking. Drop previously block_on(close_statement(...)), forcing every caller to pay a synchronous cleanup round-trip even when nothing was waiting for the result. Spawn the close on the runtime instead — best-effort fire-and-forget. Saves ~250ms on every CloudFetch query. 3. Coalesce small batches on the inline path. InlineArrowProvider was emitting 200+ tiny RecordBatches per 100K-row result (one per IPC message). Add a batch_merge_target_rows knob and apply the same coalescing logic the CloudFetch download path uses. Reduces per-batch overhead at language bindings (e.g. PyO3, ODBC). 4. Enable batch_merge_target_rows by default (128k rows). Was 0 (disabled). All consumers now get coalesced batches by default; no API change. Measured on dogfood warehouse, randomized interleaved benchmark vs Thrift backend (median wall time, fetchall_arrow path): size Rust (before) -> Rust (after) Thrift ratio SELECT 1 500ms -> 394ms 387ms 1.02x 10K 950ms -> 893ms 1014ms 0.88x 500K 2600ms -> 2178ms 3305ms 0.66x 1M 3700ms -> 3579ms 3814ms 0.94x 10M 8700ms -> 8677ms 8802ms 0.99x
This is a proof-of-concept satellite crate that exposes the Rust kernel
to Python via PyO3 + abi3, validating the multi-language strategy in
docs/kernel-strategy-final-recommendation.md (Decision 2: PyO3 over
adbc_driver_manager for Python).
Status: DRAFT / RFC. Not for merge as-is.
Surface:
Connection(host, http_path, access_token, *, catalog=None, schema=None)
.execute(sql) -> ResultSet
ResultSet
.num_columns(), .column_names(), .arrow_schema()
.fetch_next_batch() -> pyarrow.RecordBatch | None (streaming)
.fetch_all_arrow() -> pyarrow.Table (drains rest)
Working:
- End-to-end smoke against a live SEA warehouse (PAT auth).
- Inline + CloudFetch result paths.
- pyarrow.Table return via Arrow C Data Interface (zero-copy where
schema permits).
- abi3-py39 wheel — one wheel covers Python 3.9+.
- Performance at parity with or better than the existing Thrift-based
path on most query sizes.
Deferred (POC scope):
- PAT only — no OAuth M2M / U2M / Azure SP / external credential providers.
- No metadata methods (get_objects, get_table_schema, get_table_types).
- No async execute, cancel, Ctrl-C signal handling, logging bridge.
- No prepared statements / parameter binding.
- No tests, no CI integration, not packaged for PyPI.
Design:
- Wraps the kernel's ADBC Optionable layer (string-keyed config).
Should switch to typed config structs once Phase 0a (DatabricksConfig,
AuthConfig, ...) lands per the kernel-strategy doc.
- Uses a new `Statement::execute_owned` inherent method (in this PR's
parent perf branch) that returns `Box<dyn RecordBatchReader + Send +
'static>`, decoupling the reader's lifetime from the Statement so the
binding can hold the reader past the Statement's drop.
- Releases the GIL during all kernel-side work.
Depends on the perf-small-query-optimizations branch for the
`execute_owned` API and `server_side_closed` cleanup short-circuit.
Draft
8 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This is a proof-of-concept / RFC. Not for merge as-is.
Adds a new satellite crate
rust-pyo3/that exposes the Databricks ADBC Rust kernel to Python via PyO3 + abi3. Validates the multi-language strategy indocs/kernel-strategy-final-recommendation.md— Decision 2: first-party PyO3 binding over routing throughadbc_driver_manager.What works
pyarrow.Tablereturn via Arrow C Data Interface (zero-copy where the schema permits).Public surface
Explicitly out of scope (POC)
get_objects,get_table_schema,get_table_types).cancel(), no Ctrl-C signal handling, no logging bridge into Pythonlogging.Design notes for reviewers
Wraps the kernel's ADBC
Optionablelayer (string-keyed config). The kernel-strategy doc proposes typed Rust config structs (DatabricksConfig,AuthConfig, …) as Phase 0a. Once that lands, this binding should switch to typed construction. Currently this layering choice is a deliberate POC compromise — Phase 0a is a separate workstream and didn't seem worth blocking the binding spike on.Uses a new
Statement::execute_ownedinherent method (in perf(rust): reduce per-query overhead and coalesce small batches #422) that returnsBox<dyn RecordBatchReader + Send + 'static>, decoupling the reader's lifetime from the Statement. This lets the PyO3 wrapper hold the reader past the Statement's drop, which is needed for streaming.Releases the GIL during all kernel-side work (
execute_owned, batch drain) and reacquires once for the pyarrow conversion phase.Dependency
Depends on #422 for the
execute_ownedAPI and theserver_side_closedcleanup short-circuit. This branch is based on that one; if #422 lands first, this rebases trivially onto main.Open questions
databricks-adbc/rust-pyo3/long-term, or move it to its own repo (databricks-adbc-pythonor similar) once it grows beyond a POC?python-driver-rust-adbc-sea-design.md§4.2-4.3) still describes theadbc_driver_managerintegration. The newer doc supersedes that, but the older sections need a rewrite.Test plan
maturin develop --releaseproduces an abi3-py39 wheel.import databricks_adbc_pyo3works in a fresh venv.examples/e2e_smoke.pyround-trips small (SELECT 1) and large (1M rows viarange()) queries against a dogfood warehouse.This pull request and its description were written by Isaac.