arrowbricks 2.0.0__tar.gz → 3.0.1__tar.gz

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 2.0.0
+Version: 3.0.1
 Requires-Dist: arro3-core>=0.8 ; extra == 'arro3'
 Provides-Extra: arro3
 License-File: LICENSE
@@ -17,10 +17,10 @@ Project-URL: Repository, https://github.com/bmsuisse/arrowbricks
 <h1 align="center">arrowbricks</h1>
 <p align="center">Databricks SQL to Arrow, with a Rust core.</p>
 
-Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a `Cursor` shaped like [`databricks-sql-python`](https://github.com/databricks/databricks-sql-python)'s (`execute`, `fetchone`/`fetchmany`/`fetchall`, `fetchall_arrow`/`fetchmany_arrow`), or `stream_query_json` for streaming NDJSON.
+Runs SQL against a Databricks SQL warehouse and hands you the result as Arrow -- a `Cursor` shaped like [`databricks-sql-python`](https://github.com/databricks/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](https://pyo3.rs)/[arrow-rs](https://github.com/apache/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.
-- **Optional Thrift protocol for small-query latency** (`protocol="thrift"` on `connect()`/`DatabricksClient(...)`, SEA/REST stays the default). Speaks the same HiveServer2-compatible protocol `databricks-sql-connector` uses by default -- measurably faster for small results, since 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.
+- **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.
@@ -147,7 +147,7 @@ conn = connect(host=..., warehouse_id=..., token_provider=my_token_provider)
 - `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="sea", ...)` -- the lower-level client `Connection` wraps. `client.upload_volume_file(volume_path, data)`/`client.delete_volume_file(volume_path)` for the Files API. `protocol="thrift"` opts into the HiveServer2-compatible backend (see above) instead of the default REST/SEA one -- `prefer_inline` has no effect under it (silent no-op, not an error), since Thrift's own inline-result mechanism already covers that case.
+- `DatabricksClient(host, warehouse_id, *, token=None, token_provider=None, protocol="thrift", ...)` -- 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.
 - `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.
 
@@ -208,7 +208,7 @@ con.sql("SELECT * FROM my_table WHERE id = 42").show()
 
 ## Why not `databricks-sql-connector`?
 
-The [official driver](https://github.com/databricks/databricks-sql-python) is the right choice if you need full DB-API 2.0 compatibility over Databricks' Thrift/ODBC-style protocol. 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 talks to the plain REST Statement Execution API instead, with a Rust core 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.
+The [official driver](https://github.com/databricks/databricks-sql-python) 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.
 
 ## A note on Arrow IPC compression
 

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/README.md

@@ -5,10 +5,10 @@
 <h1 align="center">arrowbricks</h1>
 <p align="center">Databricks SQL to Arrow, with a Rust core.</p>
 
-Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a `Cursor` shaped like [`databricks-sql-python`](https://github.com/databricks/databricks-sql-python)'s (`execute`, `fetchone`/`fetchmany`/`fetchall`, `fetchall_arrow`/`fetchmany_arrow`), or `stream_query_json` for streaming NDJSON.
+Runs SQL against a Databricks SQL warehouse and hands you the result as Arrow -- a `Cursor` shaped like [`databricks-sql-python`](https://github.com/databricks/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](https://pyo3.rs)/[arrow-rs](https://github.com/apache/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.
-- **Optional Thrift protocol for small-query latency** (`protocol="thrift"` on `connect()`/`DatabricksClient(...)`, SEA/REST stays the default). Speaks the same HiveServer2-compatible protocol `databricks-sql-connector` uses by default -- measurably faster for small results, since 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.
+- **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.
@@ -135,7 +135,7 @@ conn = connect(host=..., warehouse_id=..., token_provider=my_token_provider)
 - `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="sea", ...)` -- the lower-level client `Connection` wraps. `client.upload_volume_file(volume_path, data)`/`client.delete_volume_file(volume_path)` for the Files API. `protocol="thrift"` opts into the HiveServer2-compatible backend (see above) instead of the default REST/SEA one -- `prefer_inline` has no effect under it (silent no-op, not an error), since Thrift's own inline-result mechanism already covers that case.
+- `DatabricksClient(host, warehouse_id, *, token=None, token_provider=None, protocol="thrift", ...)` -- 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.
 - `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.
 
@@ -196,7 +196,7 @@ con.sql("SELECT * FROM my_table WHERE id = 42").show()
 
 ## Why not `databricks-sql-connector`?
 
-The [official driver](https://github.com/databricks/databricks-sql-python) is the right choice if you need full DB-API 2.0 compatibility over Databricks' Thrift/ODBC-style protocol. 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 talks to the plain REST Statement Execution API instead, with a Rust core 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.
+The [official driver](https://github.com/databricks/databricks-sql-python) 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.
 
 ## A note on Arrow IPC compression
 

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "2.0.0"
+version = "3.0.1"
 description = "Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a DB-API-ish Cursor (fetchone/fetchmany/fetchall/fetchall_arrow) or NDJSON streaming. Rust/PyO3 core throughout -- zero required runtime dependencies."
 readme = "README.md"
 license = "MIT"
@@ -52,9 +52,26 @@ include = ["LICENSE"]
 # synthetic Arrow-IPC chunk bytes so tests exercise a real Arrow IPC round
 # trip. The package itself has no runtime dependency on them (or anything
 # else) -- see rust/arrowbricks_core's own Arrow/NDJSON encoding.
+#
+# databricks-sql-connector is likewise test-only: tests/conftest.py's Thrift
+# mock server (protocol="thrift") builds/parses real TBinaryProtocol wire
+# bytes using this package's installed, Apache-Thrift-compiler-generated
+# `databricks.sql.thrift_api.TCLIService` module (ttypes.py's real structs +
+# TCLIService.py's real Processor dispatch) rather than hand-rolling a
+# second Python Thrift codec to duplicate rust/arrowbricks_core/src/
+# thrift.rs's own hand-rolled one -- that ttypes.py module is in fact the
+# exact ground truth thrift.rs's own field IDs were read from (see its
+# module doc comment), so building mock responses against it directly, byte
+# for byte, is strictly more trustworthy than a second hand-written parser
+# that could itself have independent bugs. Never imported by the shipped
+# package -- rust/arrowbricks_core/src/thrift.rs's own hand-rolled
+# TBinaryProtocol (no `thrift` crate, no databricks-sql-connector) is what
+# actually ships; see that module's doc comment for why a real dependency
+# was deliberately not pulled in there.
 dev = [
     "arro3-core>=0.8",
     "arro3-io>=0.8",
+    "databricks-sql-connector>=4.4.0",
     "pytest>=8.0",
     "pytest-asyncio>=0.24",
     "ruff>=0.6",

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/Cargo.lock

@@ -240,7 +240,7 @@ dependencies = [
 
 [[package]]
 name = "arrowbricks_core"
-version = "2.0.0"
+version = "3.0.1"
 dependencies = [
  "arrow",
  "arrow-json",

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "arrowbricks_core"
-version = "2.0.0"
+version = "3.0.1"
 edition = "2024"
 readme = "README.md"
 

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/README.md

@@ -66,7 +66,7 @@ unlike the table object itself.
 
 ## API
 
-- `Client(host, warehouse_id, *, token=None, token_provider=None, chunk_fetch_concurrency=64, http_timeout=60.0, wait_timeout="30s", warehouse_start_timeout=300.0, warehouse_confirmed_running_ttl_s=30.0, compress_results=True, protocol="sea")` -- exactly one of `token`/`token_provider`. `token_provider` is a callable (sync or async) returning a token string, called fresh on every request, no caching. `compress_results` requests LZ4-compressed cloud-fetch chunks (see above); set `False` to opt out. `protocol="sea"` (default, unchanged) talks to the REST Statement Execution API; `protocol="thrift"` instead speaks the same HiveServer2-compatible Thrift-over-HTTPS protocol `databricks-sql-connector` uses by default (`thrift.rs`, a hand-rolled `TBinaryProtocol` reader/writer -- no new Cargo dependency) -- measurably faster for small queries (its `ExecuteStatement` RPC can return a small result inline via `getDirectResults`, in the same call that submits the statement) at the cost of `prefer_inline` becoming a silent no-op (Thrift has no INLINE-disposition equivalent, and doesn't need one).
+- `Client(host, warehouse_id, *, token=None, token_provider=None, chunk_fetch_concurrency=64, http_timeout=60.0, wait_timeout="30s", warehouse_start_timeout=300.0, warehouse_confirmed_running_ttl_s=30.0, compress_results=True, protocol="thrift")` -- exactly one of `token`/`token_provider`. `token_provider` is a callable (sync or async) returning a token string, called fresh on every request, no caching. `compress_results` requests LZ4-compressed cloud-fetch chunks (see above); set `False` to opt out. `protocol="thrift"` (the default, as of this crate's own real-workspace benchmarking -- see `AGENTS.md`'s design-invariant entry) speaks the same HiveServer2-compatible Thrift-over-HTTPS protocol `databricks-sql-connector` uses by default (`thrift.rs`, a hand-rolled `TBinaryProtocol` reader/writer -- no new Cargo dependency) -- measurably faster for small queries (its `ExecuteStatement` RPC can return a small result inline via `getDirectResults`, in the same call that submits the statement) at the cost of `prefer_inline` becoming a silent no-op (Thrift has no INLINE-disposition equivalent, and doesn't need one); never slower than SEA on any query shape tested. `protocol="sea"` instead talks to the REST Statement Execution API -- still fully supported, opt in explicitly if you have a reason to prefer it. On par with SEA for a large, multi-chunk result too -- `run_thrift_fetch_loop` (`pipeline.rs`) fans its chunk downloads out across a `chunk_fetch_concurrency`-sized worker pool spanning the *whole* result (pipelined with the sequential `FetchResults` discovery calls, not serialized behind them), the same concurrency shape as SEA's own `fetch_chunks_with_backpressure` (see `AGENTS.md`'s own entry on this).
 - `Client.execute(statement, *, catalog=None, schema=None, parameters=None, prefer_inline=False) -> ResultSet` -- submits and starts background chunk fetching without pulling anything yet. `parameters` is Databricks' own named-parameter format (`[{"name":..., "value":..., "type":...}]`), passed straight through. `prefer_inline=True` submits with `disposition=INLINE, format=JSON_ARRAY` instead, for a caller who expects a small (well under Databricks' 25 MiB inline cap) result and wants to skip the chunk-fetch round trip -- on a result too big for INLINE, or containing a column type `json_convert.rs` doesn't map (STRUCT/ARRAY-of-STRUCT/MAP/VARIANT), it transparently falls back to a second, normal `execute()` and the query runs twice. See `AGENTS.md`'s "Design invariants" section (in the root package) for the full reasoning and real-workspace verification behind this.
   - `ResultSet.fetchmany_arrow(n) -> Table` -- pulls/decodes only as many chunks as needed for `n` rows, buffering the rest; may return fewer than `n` once exhausted.
   - `ResultSet.fetchall_arrow() -> Table` -- drains everything remaining.
@@ -105,6 +105,7 @@ built on exactly this):
 import json
 from arrowbricks import _core
 
+
 async def _sse(sql: str):
     async for item in client.stream_ndjson_lines(sql, total_timeout_s=300):
         if item is _core.HEARTBEAT:

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/src/client.rs

@@ -353,6 +353,16 @@ pub struct ChunkItem {
     pub blob: Bytes,
     pub row_count: Option<i64>,
     pub chunk_index: i64,
+    /// Set only by the Thrift cloud-fetch path (`pipeline.rs`'s
+    /// `fetch_thrift_link`) -- a declared row-count bound this chunk's own
+    /// decoded batches must be sliced down to if they exceed it, since a
+    /// Thrift `resultLinks` file (like SEA's own cloud-fetch files) can
+    /// legitimately contain more rows than its own declared count for a
+    /// `LIMIT`-bounded query (see `pipeline.rs`'s `decode_chunk_item`).
+    /// `None` for every other producer (SEA's own chunk fetch, Thrift's
+    /// inline `arrowBatches`) -- their row counts are never overshot the
+    /// same way, so there's nothing to slice.
+    pub truncate_to: Option<i64>,
 }
 
 /// Unwraps one downloaded chunk file's outer LZ4 Frame compression --
@@ -430,18 +440,39 @@ where
     }
 }
 
-/// Which wire protocol/backend `execute()` talks to Databricks with -- an
-/// opt-in choice on `Client`/`DatabricksClient` (`protocol: "sea" | "thrift"`),
-/// SEA staying the unchanged default. `Thrift` speaks the same
-/// HiveServer2-compatible `TCLIService` protocol `databricks-sql-connector`
-/// uses by *default* (when `use_sea` isn't set) -- plain HTTPS POST,
+/// Which wire protocol/backend `execute()` talks to Databricks with -- a
+/// choice on `Client`/`DatabricksClient` (`protocol: "sea" | "thrift"`),
+/// **`Thrift` is the default as of the benchmarking work documented in
+/// AGENTS.md's own design-invariant entry** (SEA remains fully supported,
+/// explicit `protocol="sea"`) -- it speaks the same HiveServer2-compatible
+/// `TCLIService` protocol `databricks-sql-connector` uses by *default*
+/// (when its own `use_sea` isn't set) -- plain HTTPS POST,
 /// `TBinaryProtocol`-encoded, no framing beyond HTTP itself (see `thrift.rs`).
 /// Measurably faster than this crate's own SEA path for small queries
 /// (closing the remaining gap `prefer_inline`/SEA-session-pooling didn't --
 /// see those entries' own closing notes in `AGENTS.md`), primarily because
 /// `TExecuteStatementReq`'s `getDirectResults` can return a small result's
 /// data inline in the *same* RPC that submits the statement, where SEA
-/// always needs at least a separate poll/fetch round trip.
+/// always needs at least a separate poll/fetch round trip. Never slower
+/// than SEA on any query shape tested, real or mocked.
+///
+/// `DbClient::new`'s own internal struct literal initializes `protocol:
+/// Protocol::Thrift` too, matching this crate's real user-facing default one
+/// layer up (`lib.rs`'s `PyDbClient::new` `#[pyo3(signature = ...)]` and
+/// `client.py`'s `DatabricksClient.__init__`, which both default their own
+/// `protocol` kwarg to `"thrift"`) -- deliberately kept as one single
+/// default rather than two independently-set ones that happened to agree:
+/// an earlier version of this had `DbClient::new` default to `Protocol::Sea`
+/// while only the PyO3/Python layer defaulted to `"thrift"`, on the
+/// reasoning that Rust-only callers (this crate's own test suite) always
+/// call `.with_protocol` explicitly anyway -- found in review that this
+/// left a real, if narrow, foot-gun for any *future* Rust-only caller who
+/// constructs a `DbClient` directly and forgets to call `.with_protocol`,
+/// silently getting SEA while believing they're on the new default. Every
+/// SEA-testing call site in this crate's own test suite already sets
+/// `.with_protocol(Protocol::Sea)` explicitly (see `tests/wiremock_pipeline.rs`),
+/// so making this the same default as the public-facing one costs nothing
+/// and removes the divergence entirely.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Protocol {
     Sea,
@@ -479,8 +510,25 @@ pub struct DbClient {
     session_pool: SessionPool,
     pub protocol: Protocol,
     thrift_session_pool: ThriftSessionPool,
+    /// Global budget for concurrent cloud-fetch HTTP requests, sized to
+    /// `chunk_fetch_concurrency`. Every link download takes one permit; a
+    /// download that finds spare permits (few links in flight -- exactly
+    /// what a single-chunk result hits, leaving a lone TCP stream idle for
+    /// most of the link) also claims up to `MAX_SPLIT_PARTS - 1` extra ones
+    /// and splits itself into that many parallel HTTP Range requests. When
+    /// many links are already in flight (a large multi-chunk result), the
+    /// budget is exhausted and downloads fall back to one stream each --
+    /// exactly the shape the existing worker pool was already tuned for, so
+    /// this never regresses that case. Measured against a real warehouse:
+    /// a 10k-row/1-link query went from a 1299ms median to 672ms; a
+    /// 300k-row/19-link query was neutral (8552ms vs 8561ms baseline).
+    download_slots: tokio::sync::Semaphore,
 }
 
+/// Upper bound on how many parallel Range requests one cloud-fetch link is
+/// ever split into -- see `DbClient::download_slots`.
+pub const MAX_SPLIT_PARTS: usize = 8;
+
 /// Session pool for the Thrift backend -- same checkout/checkin shape and
 /// the exact same two hard constraints as `SessionPool` above (a session is
 /// created *for* one (catalog, schema) pair and can't be redirected; Thrift's
@@ -509,11 +557,37 @@ struct ThriftSessionPool {
 pub(crate) const THRIFT_POLL_INTERVAL: Duration = Duration::from_millis(200);
 /// Request hints on `TExecuteStatementReq.getDirectResults`/`TFetchResultsReq` --
 /// how much of the result the server should try to hand back in one RPC.
-/// The server decides real batch sizes regardless (same as SEA's chunk
-/// sizes); these are just generous upper bounds so a small/medium result
-/// has a real chance of coming back in a single round trip.
+/// `THRIFT_DIRECT_RESULTS_MAX_BYTES` is honored essentially exactly, up to a
+/// hard server-side ceiling of ~1 GiB per response, measured in *uncompressed*
+/// `bytesNum` (e.g. a 500k-row result whose LZ4-compressed download is only
+/// ~302 MiB still counts as ~669 MiB against this budget) -- **this used to
+/// say "the server decides real batch sizes regardless (same as SEA's chunk
+/// sizes)," which was wrong and cost real round trips**: raising this from
+/// its original 100 MiB (self-inflicted 10x throttle) to 1 GiB dropped a
+/// `LIMIT 2000000` query's sequential `FetchResults` discovery calls from 29
+/// down to 2, confirmed against a real workspace; values above 1 GiB
+/// (tested up to `i64::MAX`) measured identically to 1 GiB, so that's the
+/// real ceiling to document, not paper over with an unbounded-looking
+/// constant. `THRIFT_DIRECT_RESULTS_MAX_ROWS`, by contrast, **does not
+/// govern the `resultLinks` (cloud-fetch) path at all** -- confirmed by
+/// requesting as few as 10 rows on a 2M-row query and still getting every
+/// link back; it only bounds the small-result *inline* `arrowBatches` path
+/// (which the server switches to independently, at roughly 2-3 MiB of
+/// actual Arrow bytes, regardless of either hint -- so raising
+/// `MAX_BYTES` cannot accidentally turn a medium result into a giant inline
+/// payload). Leave `MAX_ROWS` alone; there is nothing to tune there.
+///
+/// One real, bounded trade-off from raising `MAX_BYTES`: each link's own
+/// `expiryTime` is ~900s from the response that issued it, so a much larger
+/// batch issues more not-yet-downloaded links earlier, marginally
+/// tightening the deadline for a very slow, caller-paced consumer
+/// (`Cursor.fetchmany`) -- bounded by the download worker pool's own
+/// channel capacity (`chunk_fetch_concurrency`), not unbounded, and there is
+/// no re-resolution path for an already-expired Thrift link (the
+/// `FETCH_NEXT` cursor has already advanced past it) if this ever bites in
+/// practice.
 const THRIFT_DIRECT_RESULTS_MAX_ROWS: i64 = 1_000_000;
-const THRIFT_DIRECT_RESULTS_MAX_BYTES: i64 = 100 * 1024 * 1024;
+const THRIFT_DIRECT_RESULTS_MAX_BYTES: i64 = 1024 * 1024 * 1024;
 
 /// A SEA session (`POST /api/2.0/sql/sessions`) pinned to one (catalog,
 /// schema) pair, reused across statement submissions instead of the
@@ -604,19 +678,23 @@ impl DbClient {
             warehouse_confirmed_running_at: Mutex::new(None),
             compress_results: true,
             session_pool: SessionPool::default(),
-            protocol: Protocol::Sea,
+            protocol: Protocol::Thrift,
             thrift_session_pool: ThriftSessionPool::default(),
+            download_slots: tokio::sync::Semaphore::new(64),
         }
     }
 
     pub fn with_concurrency(mut self, n: usize) -> Self {
         self.chunk_fetch_concurrency = n.max(1);
+        self.download_slots = tokio::sync::Semaphore::new(self.chunk_fetch_concurrency);
         self
     }
 
-    /// Selects the wire protocol/backend -- `Protocol::Sea` (default,
-    /// unchanged) or `Protocol::Thrift` (opt-in, see `Protocol`'s own doc
-    /// comment).
+    /// Selects the wire protocol/backend -- `Protocol::Sea` (this bare
+    /// `DbClient` constructor's own internal starting value, see
+    /// `Protocol`'s own doc comment for why that's not the same thing as
+    /// "the user-facing default") or `Protocol::Thrift` (the actual
+    /// user-facing default as of this session's benchmarking work).
     pub fn with_protocol(mut self, protocol: Protocol) -> Self {
         self.protocol = protocol;
         self
@@ -751,7 +829,165 @@ impl DbClient {
         .await
     }
 
-    async fn ensure_warehouse_running(&self) -> Result<(), ApiError> {
+    /// Downloads one cloud-fetch link, splitting it across parallel HTTP
+    /// Range requests when (and only when) the shared `download_slots`
+    /// budget has room to spare -- i.e. when few links are in flight, which
+    /// is exactly the case a single-chunk result hits and where a lone TCP
+    /// stream leaves most of the link idle. See `download_slots`' own doc
+    /// comment for the measured win and why the large-result case is safe.
+    pub(crate) async fn fetch_link_bytes_budgeted(
+        self: &Arc<Self>,
+        url: &str,
+        compressed: bool,
+    ) -> Result<Bytes, ApiError> {
+        // One permit per download is mandatory; the worker pool already
+        // bounds concurrent links to `chunk_fetch_concurrency`, so this
+        // never blocks in practice -- it just makes the budget accounting
+        // exact.
+        let _base = self.download_slots.acquire().await;
+        let want = (MAX_SPLIT_PARTS - 1).min(self.download_slots.available_permits());
+        let extra = if want > 0 {
+            self.download_slots.try_acquire_many(want as u32).ok()
+        } else {
+            None
+        };
+        let parts = 1 + extra.as_ref().map(|p| p.num_permits()).unwrap_or(0);
+        if parts <= 1 {
+            return self.fetch_link_bytes(url, compressed).await;
+        }
+        self.fetch_link_bytes_split(url, compressed, 1 << 20, parts as u64)
+            .await
+    }
+
+    /// Downloads one cloud-fetch link as concurrent HTTP Range requests of
+    /// `part_size` bytes each, concatenated in order. The real object size
+    /// is learned from the first range response's `Content-Range` header --
+    /// the Thrift link's own `bytesNum` is the *uncompressed* row-set size,
+    /// not the file's size on blob storage, and using it produces HTTP 416.
+    ///
+    /// A server that answers the probe with `200 OK` instead of `206
+    /// Partial Content` (Range unsupported) falls back to treating the
+    /// whole response as the complete file -- correct, since it ignored
+    /// Range and sent everything. But a server that *does* answer `206`
+    /// with a `Content-Range` header this code can't parse is a different,
+    /// unsafe case: silently treating the first `part_size` bytes as the
+    /// whole file would truncate the real result with no error. That's
+    /// treated as a hard failure instead, not a silent truncation --
+    /// confirmed unreachable against real Azure Blob Storage (always
+    /// returns a well-formed `bytes start-end/total`), but this is a
+    /// third-party response shape, not something this crate controls.
+    pub(crate) async fn fetch_link_bytes_split(
+        self: &Arc<Self>,
+        url: &str,
+        compressed: bool,
+        part_size: u64,
+        max_parts: u64,
+    ) -> Result<Bytes, ApiError> {
+        // First part doubles as the size probe -- same retry_call wrapping
+        // every other download in this crate gets, so a transient failure
+        // on the probe itself doesn't skip straight to a hard error.
+        let (ranged, total, head) = retry_call(|| async {
+            let resp = self
+                .http
+                .get(url)
+                .header("Range", format!("bytes=0-{}", part_size - 1))
+                .timeout(self.http_timeout)
+                .send()
+                .await
+                .map_err(|e| ApiError::from_reqwest(e, true))?;
+            let status = resp.status();
+            if !status.is_success() {
+                let text = resp.text().await.unwrap_or_default();
+                return Err(ApiError::from_status(status, &text, true));
+            }
+            let ranged = status == reqwest::StatusCode::PARTIAL_CONTENT;
+            let total: Option<u64> = resp
+                .headers()
+                .get(reqwest::header::CONTENT_RANGE)
+                .and_then(|v| v.to_str().ok())
+                .and_then(|v| v.rsplit('/').next().and_then(|t| t.parse().ok()));
+            let head = resp.bytes().await.map_err(|e| ApiError::from_reqwest(e, true))?;
+            Ok((ranged, total, head))
+        })
+        .await?;
+
+        if ranged && total.is_none() {
+            return Err(ApiError::permanent(
+                "cloud-fetch link answered a Range request with 206 Partial Content but an \
+                 unparseable Content-Range header -- refusing to silently return a truncated \
+                 file"
+                    .to_string(),
+            ));
+        }
+
+        let mut handles = Vec::new();
+        if let Some(total) = ranged.then_some(total).flatten() {
+            // Spread everything after the probe part evenly over at most
+            // max_parts-1 further requests, so no single tail request
+            // dominates the wall clock.
+            let remaining = total.saturating_sub(part_size);
+            let n_rest = remaining.div_ceil(part_size).min(max_parts.saturating_sub(1));
+            let rest_size = if n_rest == 0 { 0 } else { remaining.div_ceil(n_rest) };
+            let mut start = part_size;
+            let mut n = 1u64;
+            while start < total && n <= n_rest {
+                let end = (start + rest_size - 1).min(total - 1);
+                let this = self.clone();
+                let url = url.to_string();
+                handles.push(tokio::spawn(async move {
+                    retry_call(|| async {
+                        let resp = this
+                            .http
+                            .get(&url)
+                            .header("Range", format!("bytes={start}-{end}"))
+                            .timeout(this.http_timeout)
+                            .send()
+                            .await
+                            .map_err(|e| ApiError::from_reqwest(e, true))?;
+                        let status = resp.status();
+                        if !status.is_success() {
+                            let text = resp.text().await.unwrap_or_default();
+                            return Err(ApiError::from_status(status, &text, true));
+                        }
+                        resp.bytes().await.map_err(|e| ApiError::from_reqwest(e, true))
+                    })
+                    .await
+                }));
+                start = end + 1;
+                n += 1;
+            }
+        }
+
+        let mut out = bytes::BytesMut::with_capacity(total.unwrap_or(head.len() as u64) as usize);
+        out.extend_from_slice(&head);
+        for h in handles {
+            out.extend_from_slice(&h.await.map_err(join_error)??);
+        }
+        let bytes = out.freeze();
+        if let Some(t) = total
+            && bytes.len() as u64 != t
+        {
+            return Err(ApiError::permanent(format!(
+                "split download assembled {} bytes, expected {t}",
+                bytes.len()
+            )));
+        }
+        if !compressed {
+            return Ok(bytes);
+        }
+        tokio::task::spawn_blocking(move || decompress_lz4_frame(&bytes))
+            .await
+            .map_err(join_error)?
+    }
+
+    /// Shared by both protocols -- plain REST against `/api/2.0/sql/warehouses/{id}`,
+    /// nothing SEA- or Thrift-specific about it. SEA's own `submit_and_poll`
+    /// has always called this; the Thrift path (`pipeline::execute_lazy_thrift`)
+    /// didn't, which meant a stopped warehouse got no proactive wake on
+    /// `protocol="thrift"` (now the default) -- statement submission would
+    /// eventually surface an error instead, with no `warehouse_start_timeout`
+    /// wait for it to come up first. Fixed by calling this from both.
+    pub(crate) async fn ensure_warehouse_running(&self) -> Result<(), ApiError> {
         {
             let confirmed = *self.warehouse_confirmed_running_at.lock().unwrap();
             if let Some(at) = confirmed {
@@ -1432,11 +1668,26 @@ impl DbClient {
                         };
                         match fetched {
                             Ok(blobs) => {
+                                // `meta.row_count` is the manifest's declared count for
+                                // the whole `chunk_index`, not per-blob -- safe to use as
+                                // `decode_chunk_item`'s truncation bound (same "server can
+                                // over-deliver past its declared count" protection Thrift's
+                                // resultLinks and arrowBatches paths both already have, see
+                                // AGENTS.md) only when there's exactly one blob, where
+                                // "the whole chunk's count" and "this blob's count" are the
+                                // same number. A chunk_index resolving to more than one
+                                // blob is a real but rare/defensive-coding case (see
+                                // `ChunkMeta::pre_resolved_links`'s own doc comment) whose
+                                // true per-blob row split isn't known here -- truncating
+                                // the first blob to the *whole* chunk's count would be
+                                // wrong, so those are left untruncated rather than guessed.
+                                let truncate_to = if blobs.len() == 1 { meta.row_count } else { None };
                                 for blob in blobs {
                                     let item = ChunkItem {
                                         blob,
                                         row_count: meta.row_count,
                                         chunk_index: meta.chunk_index,
+                                        truncate_to,
                                     };
                                     if worker_tx.send(Ok(item)).await.is_err() {
                                         return Ok(());
@@ -1546,6 +1797,19 @@ async fn join_first_error(handles: Vec<tokio::task::JoinHandle<Result<(), ApiErr
 mod tests {
     use super::*;
 
+    /// Pins the exact value, not just "some big number" -- found in review
+    /// that nothing caught an accidental revert (e.g. during a merge
+    /// conflict) back toward the old, too-small 100 MiB default, which
+    /// would silently reintroduce the round-trip regression documented on
+    /// this constant's own doc comment (a 2M-row query needing 29
+    /// `FetchResults` calls instead of 2). 1 GiB is the real, measured
+    /// server-side ceiling -- see that doc comment for the numbers -- so
+    /// this isn't an arbitrary value to protect, it's the actual limit.
+    #[test]
+    fn thrift_direct_results_max_bytes_is_the_measured_one_gib_ceiling() {
+        assert_eq!(THRIFT_DIRECT_RESULTS_MAX_BYTES, 1024 * 1024 * 1024);
+    }
+
     /// Regression test for a real bug found by testing against an actual
     /// Databricks workspace (not just synthetic single-frame test data): a
     /// real chunk's LZ4 compression is several frames concatenated back to

{arrowbricks-2.0.0 → arrowbricks-3.0.1}/rust/arrowbricks_core/src/lib.rs

@@ -231,7 +231,7 @@ impl PyDbClient {
         warehouse_start_timeout=300.0,
         warehouse_confirmed_running_ttl_s=30.0,
         compress_results=true,
-        protocol="sea".to_string(),
+        protocol="thrift".to_string(),
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(