arrowbricks

Databricks SQL to Arrow, with a Rust core.

Runs SQL against a Databricks SQL warehouse and hands you the result as Arrow -- a Cursor shaped like databricks-sql-python's (execute, fetchone/fetchmany/fetchall, fetchall_arrow/fetchmany_arrow), or stream_query_json for streaming NDJSON. Talks to Databricks over Thrift by default (protocol="thrift"), or the REST Statement Execution API if you prefer (protocol="sea") -- see below.

• Rust core. Statement submit/poll, bounded-concurrency chunk fetch, the reorder buffer, and Arrow-IPC decode all run in a PyO3/arrow-rs extension bundled in this same package -- 1.6x-2.5x faster than a pure-Python/asyncio client on a multi-chunk result, scaling further with chunk count and concurrency where asyncio+GIL plateaus.
• Thrift by default, SEA/REST as a fully-supported alternative (protocol="sea" on connect()/DatabricksClient(...)). Thrift speaks the same HiveServer2-compatible protocol databricks-sql-connector uses by default -- measurably faster for small results (a statement's data can come back inline in the very same call that submits it, instead of SEA's separate poll-then-fetch round trip), and on par with SEA for a large, multi-chunk result too (its chunk downloads fan out concurrently across the whole result, same as SEA's own bounded-concurrency fetch, not serialized batch by batch) -- never slower than SEA on any query shape tested. Confirmed against a real production warehouse; see AGENTS.md's design-invariant entry for the benchmarking history behind the switch.
• Compressed cloud-fetch transport by default. Every statement requests LZ4-compressed chunk downloads (same default as the official databricks-sql-connector) and decompresses them in Rust before you ever see the bytes -- less data over the wire, which matters more than local decode speed for a large result. Measured ~2x faster chunk-fetch time against a real 120-column/100k-row table. Disable per-client with compress_results=False (connect()/DatabricksClient(...)) if your link to the warehouse is fast enough that decompression CPU time stops paying for itself.
• Zero required dependencies. pip install arrowbricks and go.
• Bring-your-own-auth -- a static token or your own token-refresh callable. No cloud-SDK dependency baked in.
• Result order preserved even though chunks can complete out of order over the network.
• Lazy fetching -- chunks are pulled only as fetchone/fetchmany/fetchall actually need them, not all upfront.
• Heartbeats between slow chunks (execute_streamed/stream_query_json), so a caller streaming this over e.g. SSE never goes silent during a cold warehouse start.

Install

pip install arrowbricks

Ships as precompiled platform wheels (Linux/macOS/Windows) -- no Rust toolchain needed, and nothing else to install for most of the API. Row-tuple fetches (fetchone/fetchmany/fetchall) need one optional extra: pip install arrowbricks[arro3] -- see Arrow vs. row-tuple fetches below.

Quickstart

import asyncio
from arrowbricks import connect


async def main():
    conn = connect(
        host="adb-1234567890.1.azuredatabricks.net",
        warehouse_id="abcd1234efgh5678",
        token="dapi...",  # or token_provider=... -- see Auth below
    )
    cursor = conn.cursor()

    await cursor.execute("SELECT * FROM my_catalog.my_schema.my_table LIMIT 100")
    async for row in cursor:
        print(row)

    await cursor.execute("SELECT * FROM my_catalog.my_schema.my_table LIMIT 100")
    table = await cursor.fetchall_arrow()  # an Arrow table (arro3/pyarrow/DuckDB-compatible)


asyncio.run(main())

For streaming NDJSON (e.g. a FastAPI SSE endpoint, first row out as soon as its chunk arrives):

from arrowbricks import HEARTBEAT, DatabricksClient

client = DatabricksClient(host=..., warehouse_id=..., token=...)

async for item in client.stream_query_json("SELECT * FROM my_catalog.my_schema.big_table"):
    if item is HEARTBEAT:
        continue  # forward as an SSE keep-alive comment, e.g.
    print(item)  # one ready-to-send JSON string per row

See examples/basic.py for a runnable version, examples/cursor_paging.py for paging a large result with fetchmany/fetchmany_arrow without buffering it all upfront, examples/azure_auth.py for a caching token_provider built on Azure AD (DefaultAzureCredential, needs azure-identity), or examples/oauth_m2m_auth.py for the same idea using Databricks' own OAuth machine-to-machine client-credentials flow instead -- like databricks-sql-connector's auth_type="databricks-oauth", but stdlib-only (urllib.request), no extra dependency.

FastAPI SSE example

stream_query_json is the full-speed way to serve a query over HTTP: the first row reaches the client after roughly one chunk's fetch/decode time, not the whole query's -- the Rust core is fetching, decoding, and reordering chunks concurrently the entire time, and the HEARTBEATs keep the connection alive through a slow cold warehouse start instead of the client just seeing dead air:

import os
from collections.abc import AsyncIterator

from fastapi import FastAPI
from fastapi.responses import StreamingResponse

from arrowbricks import HEARTBEAT, DatabricksClient

app = FastAPI()
client = DatabricksClient(
    host=os.environ["DATABRICKS_HOST"],
    warehouse_id=os.environ["DATABRICKS_WAREHOUSE_ID"],
    token=os.environ["DATABRICKS_TOKEN"],
)


async def _sse(sql: str) -> AsyncIterator[str]:
    async for item in client.stream_query_json(sql, total_timeout_s=300):
        if item is HEARTBEAT:
            yield ": keep-alive\n\n"  # SSE comment line -- clients ignore it, it just keeps the connection open
        else:
            yield f"data: {item}\n\n"


@app.get("/query")
async def query(sql: str) -> StreamingResponse:
    return StreamingResponse(_sse(sql), media_type="text/event-stream")

uvicorn app:app --reload
curl -N "http://localhost:8000/query?sql=SELECT+*+FROM+range(1000000)"

This example takes sql straight from the request for brevity -- arrowbricks does no SQL validation by design, so a real deployment must validate/allowlist it (or accept fixed query names + params) before exposing a route like this publicly. See examples/fastapi_sse.py for the runnable version, examples/fastapi_sse_pivot.py for the same over a buffered Cursor.fetchall_streamed result with one combined heartbeat/timeout budget across both the wait and the download, or examples/fastapi_sse_validated.py for one way to do that validation, using sqlglot to require a single read-only SELECT against an allowlist of fully-qualified tables.

Auth

connect/DatabricksClient take either:

• token: str -- a static personal access token or pre-issued OAuth token, or
• token_provider -- a callable (sync or async) returning a token string, called on every request.

arrowbricks has no opinion on how you get a token and no cloud-SDK dependency of its own. If your provider is expensive to call, cache/refresh inside it -- arrowbricks does no caching on your behalf.

conn = connect(host=..., warehouse_id=..., token_provider=my_token_provider)

Observability

Pass on_event to connect/DatabricksClient to get a QueryStats snapshot once per query, at completion (success, cancelled, timed out, or errored):

def log_query(stats):
    print(
        f"{stats.statement_id}: {stats.outcome} in {stats.fetch_s:.2f}s, {stats.bytes_downloaded} bytes, {stats.retry_count} retries"
    )


conn = connect(host=..., warehouse_id=..., token=..., on_event=log_query)

on_event is sync or async, same as token_provider, and applies to every query run through that client (not passed per-call). It's strictly fire-and-forget: any exception it raises is caught and swallowed on the Rust side, never surfaced to your query, and dispatching it never blocks or measurably slows the actual fetch -- a slow or broken logging/metrics callback can't make your queries slower or fail them. QueryStats fields: statement_id: str, protocol: "thrift" | "sea", warehouse_wait_s: float (0.0 if the client already knew the warehouse was running), submit_to_ready_s: float (statement submit -> ready to fetch), fetch_s: float (download + decode), num_chunks: int, bytes_downloaded: int, retry_count: int, concurrency_used: int, and outcome: "success" | "cancelled" | "timeout" | "error".

on_event fires once the fetch of a given query actually completes (or is abandoned) -- a Cursor whose result is never fully drained (e.g. only a partial fetchmany(), then abandoned) never fires one, same category as this package's existing documented gap around GC-based cleanup.

Cancellation

When a chunk download times out (total_timeout_s) or your code cancels the surrounding coroutine (task.cancel()/asyncio.wait_for) while it's in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's CancelOperation, or SEA's POST .../cancel) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the QueryTimeout/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.

This covers Cursor.fetchall_streamed()/fetchall_arrow_streamed(), client.stream_query_json(...), and the lower-level ._core.Client's own streamed APIs -- everything that wraps the chunk-download phase in the Rust-level heartbeat. Neither total_timeout_s nor a bare task.cancel() on Cursor.execute()/execute_streamed() itself (the initial submit/poll wait) triggers a cancel -- by the time a query's chunks are being fetched at all, the statement has typically already finished running server-side, so there's usually nothing left to cancel there in practice.

Errors

Every exception arrowbricks raises itself is an ArrowbricksError -- a RuntimeError subclass, so an except RuntimeError written before this hierarchy existed keeps catching everything unchanged. Three subclasses split by what actually changes what your code should do next:

• TransientError -- a network blip, connection reset, or 5xx that survived every internal retry (retry_attempts, exponential backoff up to retry_max_wait_s -- see below) before ever reaching you. Back off further and try again later; retrying immediately just repeats what already failed.
• AuthError -- HTTP 401/403, even after every internal retry re-fetched a token from your token_provider, or your token_provider itself raised (e.g. its own OAuth refresh call came back unauthorized). Treat the credential itself as bad/expired, not a transient blip.
• StatementError -- the SQL statement itself failed or was canceled server-side (bad SQL, a permissions error, a warehouse-side query failure). Retrying the identical statement will fail identically; this isn't a transport problem.

Anything else (a generic non-auth/non-statement HTTP error, an internal parse/decode failure) raises the plain ArrowbricksError base.

from arrowbricks import ArrowbricksError, AuthError, StatementError, TransientError

try:
    await cursor.execute(sql)
except AuthError:
    refresh_credentials()
except StatementError:
    raise  # bad SQL -- don't retry
except TransientError:
    await asyncio.sleep(5)
    await cursor.execute(sql)  # safe to retry
except ArrowbricksError:
    raise

QueryTimeout (total_timeout_s exceeded -- see Cancellation above) also subclasses ArrowbricksError, so except ArrowbricksError is a genuine catch-all for every exception this package raises itself.

API

• connect(host, warehouse_id, *, token=None, token_provider=None, ...) -> Connection
• Connection.cursor() -> Cursor
• Connection.client -> DatabricksClient -- the same client cursor() uses, for lower-level access (e.g. stream_query_json, upload_volume_file).
• Cursor.execute(sql, parameters=None, *, row_limit=None, offset=None, catalog=None, schema=None, total_timeout_s=None, prefer_inline=False) -> Cursor -- submits and waits for the statement, like a real DB-API cursor. parameters, if given, is Databricks' own named-parameter format -- [{"name": ..., "value": ..., "type": ...}] bound against :name markers in sql. prefer_inline=True tries fetching a small result (well under Databricks' 25 MiB inline cap) in the same round trip as the submission itself, skipping the chunk-fetch entirely. If the result turns out too big, the statement fails server-side before ever really executing, and arrowbricks transparently re-runs it the normal way -- safe to double-submit, since nothing committed the first time; a caller who sets prefer_inline without actually expecting a small result just pays for the query twice in this case. If instead the result comes back but has a column type this can't convert (nested ARRAY/MAP/STRUCT, VARIANT), the statement already succeeded -- arrowbricks does not silently re-run it (that could duplicate a write for non-idempotent SQL); it raises ArrowbricksError naming the statement instead, so you know it ran and can decide yourself whether re-running is safe. Leave prefer_inline off unless you know the result is small, and unless you know the query is a SELECT (or otherwise idempotent) if you want the byte-limit fallback's automatic retry to be safe too.
• Cursor.execute_streamed(...) -- same args, but an async generator yielding HEARTBEAT while waiting on a slow cold start, then the ready Cursor -- for bridging e.g. an SSE connection. Its timeout/heartbeats stop the moment the statement is ready, before any chunk has been downloaded -- see fetchall_streamed below for the download phase itself.
• Cursor.fetchone() -> tuple | None, Cursor.fetchmany(size) -> list[tuple], Cursor.fetchall() -> list[tuple], and iterating a Cursor directly -- row tuples; needs the arro3 extra.
• Cursor.fetchmany_arrow(size) -> Table, Cursor.fetchall_arrow() -> Table -- an Arrow table (implements __arrow_c_stream__, so arro3/pyarrow/DuckDB can all consume it directly, zero-copy).
• Cursor.fetchall_streamed(*, total_timeout_s=None) / Cursor.fetchall_arrow_streamed(*, total_timeout_s=None) -- like fetchall()/fetchall_arrow(), but yield HEARTBEAT while pulling chunks instead of blocking silently, then the final rows/Table -- for a caller downloading a large result over SSE who needs heartbeats (and a timeout) through the download, not just the initial wait. Compose with execute_streamed and a shared deadline if you want one combined budget across both phases (see examples/fastapi_sse_pivot.py).
• Cursor.description -- DB-API-style [(name, type_name, None, None, None, None, None), ...] after execute().
• client.stream_query_json(sql, **kwargs) (or the equivalent free function stream_query_json(client, sql, **kwargs)) -- yields HEARTBEAT, then each row as a JSON string, as soon as its chunk arrives. Timestamps come out as full ISO-8601, every column key is always present ("col":null for a null value, never an omitted key). JSON has no literal for NaN/Infinity/-Infinity, so those come back as "col":null by default -- pass non_finite_floats="string" to get "col":"NaN"/"col":"Infinity"/"col":"-Infinity" instead if you need to tell them apart from a real NULL.
• DatabricksClient(host, warehouse_id, *, token=None, token_provider=None, protocol="thrift", on_event=None, retry_attempts=6, retry_max_wait_s=20.0, ...) -- the lower-level client Connection wraps. client.upload_volume_file(volume_path, data)/client.delete_volume_file(volume_path) for the Files API. Pass protocol="sea" to opt into the REST Statement Execution API backend instead of the default Thrift one (see above) -- prefer_inline (SEA-only) has no effect under protocol="thrift" (silent no-op, not an error), since Thrift's own inline-result mechanism already covers that case. on_event -- see Observability below. retry_attempts/retry_max_wait_s tune the retry policy behind TransientError/AuthError (see Errors above) -- total attempts (the first try plus retry_attempts - 1 retries) and the exponential-backoff ceiling in seconds, for every retryable request this client makes (statement submit/poll, chunk-index resolution, chunk download, volume file ops). retry_attempts must be at least 1 (raises ValueError otherwise).
• write_ipc_stream(table, buf) -- writes any Arrow-C-Data-Interface-compatible object as an uncompressed Arrow-IPC stream (see below).
• ReplayableArrowChunk(data: bytes, chunk_index, declared_row_count=None) -- wraps raw Arrow-IPC stream bytes (e.g. previously downloaded and stored) so they can be read more than once via __arrow_c_stream__ (a schema peek, then the actual scan -- DuckDB's registration path does this), and .to_table() for a one-shot parse. No extra dependency needed.

Cursor.execute/execute_streamed/stream_query_json all accept catalog, schema, row_limit, offset, and total_timeout_s.

Arrow vs. row-tuple fetches

Everything above works with zero dependencies installed except row-tuple fetches. fetchall_arrow/fetchmany_arrow return an Arrow table straight from the Rust core -- the faster path if your code can consume Arrow directly (DuckDB, pyarrow, polars, a Parquet writer, ...):

import duckdb

table = await cursor.fetchall_arrow()
duckdb.sql("SELECT count(*) FROM table").show()  # DuckDB reads it zero-copy

fetchone/fetchmany/fetchall (and iterating a Cursor directly) materialize actual Python tuples instead -- ("id", "label")-style rows you can index into, print, or pass to code that doesn't know about Arrow at all. That conversion needs arro3-core (pip install arrowbricks[arro3]):

await cursor.execute("SELECT id, label FROM my_catalog.my_schema.my_table")
async for row in cursor:  # or: rows = await cursor.fetchall()
    print(row[0], row[1])

Calling a row-tuple method without arro3-core installed raises a ModuleNotFoundError naming the exact install command, rather than failing silently or with a confusing traceback.

Using with DuckDB

Anything arrowbricks hands back as Arrow (fetchall_arrow/fetchmany_arrow, ReplayableArrowChunk) implements __arrow_c_stream__, so DuckDB can register and query it directly -- zero-copy, no intermediate materialization, and no arro3-core/pyarrow install needed on top:

import duckdb
from arrowbricks import connect

conn = connect(host=..., warehouse_id=..., token=...)
cursor = conn.cursor()
await cursor.execute("SELECT * FROM my_catalog.my_schema.my_table LIMIT 100")
table = await cursor.fetchall_arrow()

con = duckdb.connect()
con.register("my_table", table)
con.sql("SELECT count(*) FROM my_table").show()

ReplayableArrowChunk works the same way for Arrow-IPC bytes you fetched and stored earlier (e.g. a raw chunk's bytes, cached in Redis/a file/wherever) -- DuckDB's registration path calls __arrow_c_stream__ twice (a schema peek, then the actual scan), which is exactly what ReplayableArrowChunk exists to support:

from arrowbricks import ReplayableArrowChunk

chunk = ReplayableArrowChunk(stored_bytes, chunk_index=0)
con.register("my_table", chunk)
con.sql("SELECT * FROM my_table WHERE id = 42").show()

Rust core

rust/arrowbricks_core is the crate implementing the hot path above, built into this same arrowbricks wheel as a compiled submodule -- not a separate PyPI package. See its own README for the crate-level design, plus standalone DuckDB and FastAPI SSE examples against the compiled extension directly.

Why not databricks-sql-connector?

The official driver is the right choice if you need full DB-API 2.0 compatibility. If you just want a query result as Arrow/JSON in your own async app, it drags in a lot for that: pandas, thrift, openpyxl, pybreaker, pyjwt, oauthlib, lz4, requests, urllib3 as hard dependencies. arrowbricks speaks the same wire protocols (Thrift by default, or the REST Statement Execution API via protocol="sea") with a hand-rolled Rust implementation instead, and zero required dependencies of its own. The Cursor API is deliberately shaped like the official driver's so switching between them is mostly a constructor change, but arrowbricks is async throughout (execute, fetchone, etc. are all coroutines) -- there's no sync escape hatch.

Benchmarks

Measured against a real Databricks SQL warehouse (Azure Databricks, 2X-Small Pro serverless warehouse, Photon on, 1-4 auto-scaling clusters -- the smallest/cheapest warehouse tier, deliberately: a bigger warehouse would narrow the gap by making the query itself slower and the client-side overhead this compares proportionally smaller). Query: SELECT id, id * 2 AS doubled, CAST(id AS STRING) AS label FROM range(200000) (200k rows, 3 columns), 3 timed runs after 1 discarded warm-up run, one connection reused per library:

arrowbricks: ~1.8x faster, same order-of-magnitude peak memory for this query size -- at 200k rows the Python interpreter's own baseline footprint dominates over the actual result data for both libraries, so this particular number doesn't show a difference. The real memory/footprint difference is in what gets installed, not what a single small query allocates:

Run it yourself: examples/benchmark_vs_connector.py (needs both packages installed: pip install arrowbricks databricks-sql-connector). Reads DATABRICKS_HOST/DATABRICKS_WAREHOUSE_ID/DATABRICKS_TOKEN from the environment or a .env file (never commit one with a real token in it) and runs each library in its own subprocess so peak memory reflects that library alone:

DATABRICKS_HOST=adb-1234567890.1.azuredatabricks.net \
DATABRICKS_WAREHOUSE_ID=abcd1234efgh5678 \
DATABRICKS_TOKEN=dapiXXXXXXXXXXXXXXXXXXXXXXXXXXXX \
python examples/benchmark_vs_connector.py

BENCHMARK_SQL overrides the query, BENCHMARK_RUNS (default 3) controls how many timed runs to average.

A note on Arrow IPC compression

write_ipc_stream (and everything in this package that serializes Arrow-IPC bytes) always writes uncompressed bodies. A compressed body (arro3's own default is compression="LZ4") is transparently decompressed by some Arrow readers (e.g. DuckDB's) but not necessarily by every other Arrow IPC reader -- notably, duckdb-wasm's browser-side decoder silently fails to parse LZ4-compressed bodies. Since arrowbricks' bytes might end up read by anything, plain uncompressed is the safe default.

License

MIT