arrowbricks 1.2.0__tar.gz → 1.3.0__tar.gz

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 1.2.0
+Version: 1.3.0
 Requires-Dist: arro3-core>=0.8 ; extra == 'arro3'
 Provides-Extra: arro3
 License-File: LICENSE

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "1.2.0"
+version = "1.3.0"
 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"

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/rust/arrowbricks_core/Cargo.lock

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

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/rust/arrowbricks_core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "arrowbricks_core"
-version = "1.2.0"
+version = "1.3.0"
 edition = "2024"
 readme = "README.md"
 
@@ -45,7 +45,7 @@ default = ["extension-module"]
 extension-module = ["pyo3/extension-module"]
 
 [dev-dependencies]
-tokio = { version = "1.53.1", features = ["rt-multi-thread", "macros", "test-util"] }
+tokio = { version = "1.53.1", features = ["rt-multi-thread", "macros", "test-util", "net", "io-util"] }
 wiremock = "0.6.5"
 
 [profile.release]

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/rust/arrowbricks_core/README.md

@@ -65,7 +65,7 @@ unlike the table object itself.
 
 ## API
 
-- `Client(host, warehouse_id, *, token=None, token_provider=None, chunk_fetch_concurrency=32, http_timeout=60.0, wait_timeout="30s", warehouse_start_timeout=300.0, warehouse_confirmed_running_ttl_s=30.0, compress_results=True)` -- 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.
+- `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)` -- 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.
 - `Client.execute_arrow(statement, *, catalog=None, schema=None, parameters=None) -> Table` -- eager: fetches and assembles the whole result before returning. `parameters` is Databricks' own named-parameter format (`[{"name":..., "value":..., "type":...}]`), passed straight through.
 - `Client.execute(statement, *, catalog=None, schema=None, parameters=None) -> ResultSet` -- submits and starts background chunk fetching without pulling anything yet.
   - `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.

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/rust/arrowbricks_core/src/client.rs

@@ -84,6 +84,41 @@ struct StatementResponseBody {
     status: StatementStatusBody,
     #[serde(default)]
     manifest: Option<ManifestBody>,
+    /// SEA embeds whichever chunks are already ready straight in the
+    /// submit/poll response body -- confirmed against a real workspace: a
+    /// SUCCEEDED response's `result.external_links` already contained chunk
+    /// 0's presigned URL, disposition EXTERNAL_LINKS same as always. Used to
+    /// skip that chunk's own `GET .../result/chunks/{i}` resolution request
+    /// entirely (see `ChunkMeta::pre_resolved_links`) -- a full round trip
+    /// saved for whichever chunks land here, which for a fast/small query is
+    /// a meaningful fraction of total latency.
+    #[serde(default)]
+    result: Option<ResultBody>,
+}
+
+#[derive(Deserialize)]
+struct ResultLinkBody {
+    // `#[serde(default)]` (not `Option<i64>`) deliberately: an omitempty-style
+    // server serializer would drop a zero-valued `chunk_index` field entirely
+    // rather than emit `0` -- exactly the chunk this optimization targets
+    // most (chunk 0), and exactly the case that would otherwise turn one
+    // optional fast-path field into a hard parse failure for the *whole*
+    // statement response (`authed_json` fails the entire `StatementResponseBody`
+    // deserialize on any missing required field, no retry). Defaulting to 0
+    // reconstructs the omitted value correctly either way.
+    #[serde(default)]
+    chunk_index: i64,
+    // Same reasoning, plus an empty link is filtered out where this is
+    // consumed (`execute_statement`) rather than trusted -- a link that
+    // somehow came through empty is worse than just re-resolving normally.
+    #[serde(default)]
+    external_link: String,
+}
+
+#[derive(Deserialize, Default)]
+struct ResultBody {
+    #[serde(default)]
+    external_links: Vec<ResultLinkBody>,
 }
 
 #[derive(Deserialize)]
@@ -149,10 +184,38 @@ impl ApiError {
         }
     }
 
-    fn from_reqwest(e: reqwest::Error) -> Self {
+    /// `idempotent` must be `true` only when replaying the *whole* request is
+    /// actually safe -- a GET, or a PUT/DELETE whose retried body is byte-for-
+    /// byte the same operation either way (`upload_volume_file`'s
+    /// overwrite=true PUT, `delete_volume_file`'s DELETE). It must be `false`
+    /// for the statement-submit POST: a decode/mid-flight-send failure there
+    /// means the server may have already accepted and started executing the
+    /// statement before the response broke -- for arbitrary caller SQL
+    /// (INSERT/MERGE/COPY INTO), blindly replaying that POST risks a second,
+    /// duplicate execution, not just a duplicate read. `is_decode()` alone
+    /// isn't the whole idempotent-failure surface either -- `is_request()`
+    /// (excluding `is_connect()`/`is_timeout()`, already handled) also covers
+    /// hyper's own "connection closed before message completed" pooled-
+    /// connection-reuse race, which fires before any body is even read and
+    /// is equally safe to retry on an idempotent request.
+    fn from_reqwest(e: reqwest::Error, idempotent: bool) -> Self {
+        // `is_decode()` (reqwest's `Kind::Decode`) isn't only content-decoding
+        // -- `Response::bytes()`/`.text()`/`.json()` also wrap a body read
+        // that fails mid-stream (connection reset, truncated transfer) in
+        // the same `Kind::Decode`, message "error decoding response body"
+        // (see reqwest's `async_impl::response::Response::do_bytes` and
+        // `error::decode`). Confirmed against a real workspace: a 400-chunk/
+        // 5.6M-row fetch hit exactly this on one of many large concurrent
+        // blob downloads, and it reproduced identically before this change
+        // existed -- a genuine transient network blip, not a permanent
+        // problem with the request. Connect/timeout errors stay non-transient,
+        // unchanged from before -- those mean the endpoint genuinely isn't
+        // responding, where failing fast on the caller's own timeout is still
+        // right.
+        let transient = idempotent && (e.is_decode() || (e.is_request() && !e.is_connect() && !e.is_timeout()));
         Self {
             message: e.to_string(),
-            transient: false,
+            transient,
         }
     }
 
@@ -181,6 +244,16 @@ pub(crate) fn join_error(e: tokio::task::JoinError) -> ApiError {
 pub struct ChunkMeta {
     pub chunk_index: i64,
     pub row_count: Option<i64>,
+    /// Set when the statement submit/poll response already embedded one or
+    /// more of this chunk's presigned URLs (see `StatementResponseBody::result`)
+    /// -- the fetch worker downloads them directly instead of first resolving
+    /// via `GET .../result/chunks/{i}`. `Vec`, not `Option<String>` -- a
+    /// `chunk_index` can carry more than one blob (same reason
+    /// `fetch_chunk_index` returns `Vec<Bytes>` and `ReorderBuffer` keys on
+    /// `VecDeque`, not a single item: collapsing to one would silently drop
+    /// every link but the last for the same index). Empty means "not
+    /// pre-resolved, fetch it the normal way."
+    pub pre_resolved_links: Vec<String>,
 }
 
 /// What submitting a statement gets you before any chunk is fetched:
@@ -233,10 +306,15 @@ pub struct ChunkItem {
 /// `out` reads every frame without reconstructing a decoder per frame.
 fn decompress_lz4_frame(compressed: &Bytes) -> Result<Bytes, ApiError> {
     use std::io::Read;
-    // Decompressed size is always >= compressed size for real chunk data, so
-    // this is a real (if approximate) lower bound, not a guess -- avoids
-    // repeated doubling-and-copy growth across frames.
-    let mut out = Vec::with_capacity(compressed.len());
+    // `compressed.len()` is a real lower bound, but LZ4 on Arrow-IPC data
+    // (long dictionary/offset-buffer runs, mostly-repeated bytes) typically
+    // compresses several-fold -- estimating just the lower bound means the
+    // real decompressed size almost always blows past initial capacity,
+    // paying for repeated doubling-and-copy growth on every chunk. `* 4` is
+    // a heuristic, not a guarantee (`Vec` still grows normally if it's wrong
+    // either way) -- just a better starting point than the guaranteed-too-
+    // small lower bound.
+    let mut out = Vec::with_capacity(compressed.len() * 4);
     let mut decoder = lz4_flex::frame::FrameDecoder::new(&compressed[..]);
     loop {
         let before = out.len();
@@ -315,18 +393,22 @@ impl DbClient {
             wait_timeout: "30s".to_string(),
             // Python's DatabricksClient defaults to 6, tuned for asyncio+GIL
             // where higher concurrency stops paying off past single digits
-            // (see its own comment). Ad hoc benchmarking against a mocked
-            // warehouse (80 chunks, 8ms simulated per-chunk latency) showed
-            // this Rust core's real OS-thread parallelism keeps improving up
-            // to ~64 concurrent fetches (119ms@8, 70ms@16, 48ms@32, 42ms@64),
-            // regressing slightly at 128 once concurrency exceeds the chunk
-            // count -- not reproducible from a script in this repo, so take
-            // the specific numbers as directional, not a checked-in
-            // benchmark. 32 is a reasonable default headroom below that peak
-            // without chasing a number that's workload-shaped; callers with
-            // very large chunk counts
+            // (see its own comment). This Rust core's real OS-thread
+            // parallelism keeps paying off well past that -- measured against
+            // a real 400-chunk/5.6M-row/120-column table, repeated runs (full
+            // download time, same query, same warehouse): 16=140s,
+            // 32=~113s (avg of 3), 64=114s, 96=~102s (avg of 2), 128=122s.
+            // 16 is clearly worse and 128 clearly regresses (concurrency
+            // outrunning what the warehouse/network can actually keep fed);
+            // 32-96 all land in roughly the same band with 96 nominally
+            // fastest on this table/warehouse, but the gap over 32 is modest
+            // (~10%) and the exact peak is workload/warehouse-shaped, not a
+            // fixed constant -- 64 is picked as a safe middle-ground default
+            // headroom above the old value with no observed downside, not
+            // a claim that 64 is the true optimum. A caller with a very
+            // large chunk count or a fast/low-latency link to the warehouse
             // may still want to raise it further.
-            chunk_fetch_concurrency: 32,
+            chunk_fetch_concurrency: 64,
             warehouse_start_timeout: Duration::from_secs(300),
             warehouse_confirmed_running_ttl: Duration::from_secs(30),
             warehouse_confirmed_running_at: Mutex::new(None),
@@ -394,6 +476,11 @@ impl DbClient {
         url: &str,
         body: Option<&Value>,
     ) -> Result<T, ApiError> {
+        // Only a GET is safe to blindly replay on a decode/mid-flight-send
+        // failure -- see `ApiError::from_reqwest`'s doc. The only POSTs this
+        // crate makes through here are statement submission and
+        // warehouse-start, neither of which is safe to risk double-executing.
+        let idempotent = method == reqwest::Method::GET;
         retry_call(|| async {
             // Fetched fresh on every attempt, not just once before the retry
             // loop -- matches Python's _bearer_token being called on every
@@ -408,9 +495,9 @@ impl DbClient {
             if let Some(b) = body {
                 req = req.json(b);
             }
-            let resp = req.send().await.map_err(ApiError::from_reqwest)?;
+            let resp = req.send().await.map_err(|e| ApiError::from_reqwest(e, idempotent))?;
             let status = resp.status();
-            let text = resp.text().await.map_err(ApiError::from_reqwest)?;
+            let text = resp.text().await.map_err(|e| ApiError::from_reqwest(e, idempotent))?;
             if !status.is_success() {
                 return Err(ApiError::from_status(status, &text));
             }
@@ -439,13 +526,13 @@ impl DbClient {
                 .timeout(self.http_timeout)
                 .send()
                 .await
-                .map_err(ApiError::from_reqwest)?;
+                .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));
             }
-            let bytes = resp.bytes().await.map_err(ApiError::from_reqwest)?;
+            let bytes = resp.bytes().await.map_err(|e| ApiError::from_reqwest(e, true))?;
             if !compressed {
                 return Ok(bytes);
             }
@@ -593,10 +680,25 @@ impl DbClient {
 
         let manifest = data.manifest.unwrap_or_default();
         let compressed = manifest.result_compression.as_deref() == Some("LZ4_FRAME");
+        // `Vec` per index, not a plain map entry -- see `ChunkMeta::pre_resolved_links`'s
+        // doc for why collapsing to one would silently lose rows.
+        let mut pre_resolved: std::collections::HashMap<i64, Vec<String>> = std::collections::HashMap::new();
+        if let Some(r) = data.result {
+            for link in r.external_links {
+                if link.external_link.is_empty() {
+                    continue;
+                }
+                pre_resolved
+                    .entry(link.chunk_index)
+                    .or_default()
+                    .push(link.external_link);
+            }
+        }
         let chunk_metas = manifest
             .chunks
             .into_iter()
             .map(|c| ChunkMeta {
+                pre_resolved_links: pre_resolved.remove(&c.chunk_index).unwrap_or_default(),
                 chunk_index: c.chunk_index,
                 row_count: c.row_count,
             })
@@ -629,6 +731,16 @@ impl DbClient {
         Ok(blobs)
     }
 
+    /// Same shape as `fetch_chunk_index`, but for links the statement submit/
+    /// poll response already embedded -- no resolution GET needed first.
+    async fn fetch_pre_resolved_links(&self, links: &[String], compressed: bool) -> Result<Vec<Bytes>, ApiError> {
+        let mut blobs = Vec::with_capacity(links.len());
+        for link in links {
+            blobs.push(self.fetch_link_bytes(link, compressed).await?);
+        }
+        Ok(blobs)
+    }
+
     /// Bounded-concurrency worker pool, mirroring
     /// `_fetch_chunks_with_backpressure`: a fixed pool of workers pulls from
     /// a shared queue of chunk metas, each pushing its blobs into a bounded
@@ -660,10 +772,16 @@ impl DbClient {
                     loop {
                         let meta = { queue.lock().unwrap().pop_front() };
                         let Some(meta) = meta else { return Ok(()) };
-                        match client
-                            .fetch_chunk_index(&statement_id, meta.chunk_index, compressed)
-                            .await
-                        {
+                        let fetched = if meta.pre_resolved_links.is_empty() {
+                            client
+                                .fetch_chunk_index(&statement_id, meta.chunk_index, compressed)
+                                .await
+                        } else {
+                            client
+                                .fetch_pre_resolved_links(&meta.pre_resolved_links, compressed)
+                                .await
+                        };
+                        match fetched {
                             Ok(blobs) => {
                                 for blob in blobs {
                                     let item = ChunkItem {
@@ -713,7 +831,9 @@ impl DbClient {
                 .timeout(self.http_timeout)
                 .send()
                 .await
-                .map_err(ApiError::from_reqwest)?;
+                // `overwrite=true` means retrying this exact PUT is safe --
+                // same idempotency reasoning as a GET.
+                .map_err(|e| ApiError::from_reqwest(e, true))?;
             let status = resp.status();
             if !status.is_success() {
                 let text = resp.text().await.unwrap_or_default();
@@ -738,7 +858,9 @@ impl DbClient {
                 .timeout(self.http_timeout)
                 .send()
                 .await
-                .map_err(ApiError::from_reqwest)?;
+                // DELETE is naturally idempotent here too (404 is already
+                // treated as success below).
+                .map_err(|e| ApiError::from_reqwest(e, true))?;
             let status = resp.status();
             if status == StatusCode::NOT_FOUND || status.is_success() {
                 return Ok(());
@@ -857,4 +979,57 @@ mod tests {
             err.message
         );
     }
+
+    /// Regression test for a real bug found against a real workspace: a
+    /// 400-chunk/5.6M-row fetch failed permanently on one of many large
+    /// concurrent blob downloads with reqwest's "error decoding response
+    /// body" -- a connection that closes early mid-body (`Kind::Decode`,
+    /// see `ApiError::from_reqwest`'s comment), not a genuinely dead
+    /// endpoint. Before the fix, `transient` was unconditionally `false` for
+    /// every reqwest error, so `retry_call` never got a second attempt and
+    /// the whole query failed outright. A raw truncated-response server
+    /// (rather than wiremock, which doesn't expose a way to violate its own
+    /// Content-Length) reproduces the same client-side error reqwest raised
+    /// against the real blob storage endpoint.
+    #[tokio::test]
+    async fn fetch_link_bytes_retries_after_a_connection_closed_mid_body() {
+        use tokio::io::{AsyncReadExt, AsyncWriteExt};
+        use tokio::net::TcpListener;
+
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+
+        tokio::spawn(async move {
+            let mut attempt = 0u32;
+            loop {
+                let Ok((mut socket, _)) = listener.accept().await else {
+                    return;
+                };
+                attempt += 1;
+                let mut buf = [0u8; 1024];
+                let _ = socket.read(&mut buf).await;
+                if attempt == 1 {
+                    // Claims 100 bytes, sends 10, then closes -- reqwest's
+                    // `.bytes()` surfaces exactly this as `Kind::Decode`.
+                    let _ = socket
+                        .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n0123456789")
+                        .await;
+                    let _ = socket.shutdown().await;
+                } else {
+                    let _ = socket
+                        .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nhello")
+                        .await;
+                    let _ = socket.shutdown().await;
+                    return;
+                }
+            }
+        });
+
+        let client = DbClient::new(&format!("http://{addr}"), "wh-test", "fake-token");
+        let bytes = client
+            .fetch_link_bytes(&format!("http://{addr}/data"), false)
+            .await
+            .expect("must retry past the truncated first attempt and succeed on the second");
+        assert_eq!(&bytes[..], b"hello");
+    }
 }

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/rust/arrowbricks_core/src/lib.rs

@@ -213,7 +213,7 @@ impl PyDbClient {
         warehouse_id,
         token=None,
         token_provider=None,
-        chunk_fetch_concurrency=32,
+        chunk_fetch_concurrency=64,
         http_timeout=60.0,
         wait_timeout="30s".to_string(),
         warehouse_start_timeout=300.0,

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/rust/arrowbricks_core/src/pipeline.rs

@@ -361,30 +361,64 @@ fn non_finite_token(v: f64) -> Option<&'static str> {
     }
 }
 
+/// One top-level Float32/Float64 column, downcast once per batch rather than
+/// once per (row, field) -- see `patch_non_finite_floats`.
+enum FloatColumn<'a> {
+    F64(usize, &'a arrow::array::Float64Array),
+    F32(usize, &'a arrow::array::Float32Array),
+}
+
 /// Rewrites each affected line's `null` (arrow-json's fixed encoding for a
 /// non-finite float, see `encode_ndjson_lines`) to a `"NaN"`/`"Infinity"`/
 /// `"-Infinity"` JSON string, in place. Only scans top-level Float32/Float64
 /// columns -- one row of `lines` per row of the batches, in the same order.
+/// Precomputes each batch's float columns once (schema scan + downcast) up
+/// front instead of redoing both per row -- for a wide, non-float-heavy
+/// schema (this feature's own motivating case: a 120-column real table) that
+/// was a dynamic downcast attempt on every column for every row, the large
+/// majority immediately discarded, and a schema with no float columns at all
+/// still paid for the full row x column scan for nothing.
 fn patch_non_finite_floats(batches: &[RecordBatch], lines: &mut [String]) {
     let mut global_row = 0usize;
     for batch in batches {
         let schema = batch.schema();
+        let float_columns: Vec<FloatColumn> = schema
+            .fields()
+            .iter()
+            .enumerate()
+            .filter_map(|(field_index, field)| match field.data_type() {
+                DataType::Float64 => Some(FloatColumn::F64(
+                    field_index,
+                    batch.column(field_index).as_primitive::<Float64Type>(),
+                )),
+                DataType::Float32 => Some(FloatColumn::F32(
+                    field_index,
+                    batch.column(field_index).as_primitive::<Float32Type>(),
+                )),
+                _ => None,
+            })
+            .collect();
+
+        if float_columns.is_empty() {
+            global_row += batch.num_rows();
+            continue;
+        }
+
         for row_in_batch in 0..batch.num_rows() {
-            for (field_index, field) in schema.fields().iter().enumerate() {
-                let token = match field.data_type() {
-                    DataType::Float64 => {
-                        let arr = batch.column(field_index).as_primitive::<Float64Type>();
+            for col in &float_columns {
+                let (field_index, token) = match col {
+                    FloatColumn::F64(field_index, arr) => (
+                        *field_index,
                         arr.is_valid(row_in_batch)
                             .then(|| non_finite_token(arr.value(row_in_batch)))
-                            .flatten()
-                    }
-                    DataType::Float32 => {
-                        let arr = batch.column(field_index).as_primitive::<Float32Type>();
+                            .flatten(),
+                    ),
+                    FloatColumn::F32(field_index, arr) => (
+                        *field_index,
                         arr.is_valid(row_in_batch)
                             .then(|| non_finite_token(arr.value(row_in_batch) as f64))
-                            .flatten()
-                    }
-                    _ => None,
+                            .flatten(),
+                    ),
                 };
                 if let Some(token) = token {
                     lines[global_row] = replace_nth_top_level_null(&lines[global_row], field_index, token);

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/rust/arrowbricks_core/tests/wiremock_pipeline.rs

@@ -173,6 +173,217 @@ async fn execute_statement_omits_compression_when_disabled() {
     );
 }
 
+/// Regression test for the SEA "embedded first chunk" optimization: a real
+/// workspace's SUCCEEDED statement response already carries chunk 0's
+/// presigned URL directly in `result.external_links`, so the client must
+/// skip that chunk's own `GET .../result/chunks/0` resolution call entirely
+/// and go straight to downloading the blob. `.expect(0)` on that mock fails
+/// the test (on `MockServer` drop) if it's ever hit -- proves the round trip
+/// is actually skipped, not just that the pipeline still produces correct
+/// rows some other way. Chunk 1 has no embedded link, so it must still go
+/// through the normal resolution path -- proves both branches coexist
+/// correctly in the same statement.
+#[tokio::test]
+async fn pre_resolved_chunk0_link_skips_the_extra_resolution_get() {
+    let server = MockServer::start().await;
+
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/warehouses/{WAREHOUSE_ID}")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"state": "RUNNING"})))
+        .mount(&server)
+        .await;
+
+    let uri = server.uri();
+    Mock::given(method("POST"))
+        .and(path("/api/2.0/sql/statements"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "statement_id": STATEMENT_ID,
+            "status": {"state": "SUCCEEDED"},
+            "manifest": {
+                "chunks": [{"chunk_index": 0, "row_count": 5}, {"chunk_index": 1, "row_count": 5}],
+                "schema": {"columns": [
+                    {"name": "id", "type_name": "LONG"},
+                    {"name": "label", "type_name": "STRING"},
+                ]},
+            },
+            "result": {
+                "external_links": [{"chunk_index": 0, "external_link": format!("{uri}/_data/chunk-0")}]
+            },
+        })))
+        .mount(&server)
+        .await;
+
+    // Never mounted for chunk 0 -- `pre_resolved_link` must mean the fetch
+    // worker never even attempts this request.
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/statements/{STATEMENT_ID}/result/chunks/0")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"external_links": []})))
+        .expect(0)
+        .mount(&server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/statements/{STATEMENT_ID}/result/chunks/1")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "external_links": [{"external_link": format!("{uri}/_data/chunk-1")}]
+        })))
+        .expect(1)
+        .mount(&server)
+        .await;
+
+    Mock::given(method("GET"))
+        .and(path_regex(r"^/_data/chunk-0$"))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_raw(build_chunk_bytes(0, 5), "application/vnd.apache.arrow.stream"),
+        )
+        .mount(&server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path_regex(r"^/_data/chunk-1$"))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_raw(build_chunk_bytes(5, 10), "application/vnd.apache.arrow.stream"),
+        )
+        .mount(&server)
+        .await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token"));
+    let summary = run_pipeline(client, "SELECT * FROM t", None, None, None).await.unwrap();
+
+    assert_eq!(summary.num_chunks, 2);
+    assert_eq!(summary.num_rows(), 10);
+    assert_ids_in_order(&summary.batches, 10);
+}
+
+/// Regression test for a bug caught in code review before it shipped:
+/// `pre_resolved` was originally a plain `HashMap<i64, String>`, which keeps
+/// only the *last* entry when `result.external_links` has more than one for
+/// the same `chunk_index` (a real, supported shape -- `fetch_chunk_index`'s
+/// own resolution path already returns `Vec<Bytes>` per chunk for exactly
+/// this reason). Two links for chunk_index 0 here -- if the fix regresses
+/// back to losing all but the last, this comes back with 5 rows instead of
+/// 10.
+#[tokio::test]
+async fn pre_resolved_links_supports_multiple_links_for_the_same_chunk_index() {
+    let server = MockServer::start().await;
+
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/warehouses/{WAREHOUSE_ID}")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"state": "RUNNING"})))
+        .mount(&server)
+        .await;
+
+    let uri = server.uri();
+    Mock::given(method("POST"))
+        .and(path("/api/2.0/sql/statements"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "statement_id": STATEMENT_ID,
+            "status": {"state": "SUCCEEDED"},
+            "manifest": {
+                "chunks": [{"chunk_index": 0, "row_count": 10}],
+                "schema": {"columns": [
+                    {"name": "id", "type_name": "LONG"},
+                    {"name": "label", "type_name": "STRING"},
+                ]},
+            },
+            "result": {
+                "external_links": [
+                    {"chunk_index": 0, "external_link": format!("{uri}/_data/chunk-0a")},
+                    {"chunk_index": 0, "external_link": format!("{uri}/_data/chunk-0b")},
+                ]
+            },
+        })))
+        .mount(&server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/statements/{STATEMENT_ID}/result/chunks/0")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"external_links": []})))
+        .expect(0)
+        .mount(&server)
+        .await;
+
+    Mock::given(method("GET"))
+        .and(path_regex(r"^/_data/chunk-0a$"))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_raw(build_chunk_bytes(0, 5), "application/vnd.apache.arrow.stream"),
+        )
+        .mount(&server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path_regex(r"^/_data/chunk-0b$"))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_raw(build_chunk_bytes(5, 10), "application/vnd.apache.arrow.stream"),
+        )
+        .mount(&server)
+        .await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token"));
+    let summary = run_pipeline(client, "SELECT * FROM t", None, None, None).await.unwrap();
+
+    assert_eq!(
+        summary.num_rows(),
+        10,
+        "both pre-resolved links for chunk_index 0 must be fetched, not just the last one"
+    );
+}
+
+/// Regression test for a bug caught in code review before it shipped: an
+/// omitempty-style server serializer could drop a zero-valued `chunk_index`
+/// field entirely rather than emit `0` -- exactly chunk 0, the chunk this
+/// optimization targets most. Before `chunk_index` had `#[serde(default)]`,
+/// this would either fail the whole `StatementResponseBody` parse (turning
+/// an optional fast path into a hard failure for the entire statement) or
+/// -- depending on the exact fix -- misattribute the link to the wrong
+/// chunk. Omitting `chunk_index` here must still resolve to chunk 0.
+#[tokio::test]
+async fn pre_resolved_link_with_omitted_chunk_index_defaults_to_chunk_zero() {
+    let server = MockServer::start().await;
+
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/warehouses/{WAREHOUSE_ID}")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"state": "RUNNING"})))
+        .mount(&server)
+        .await;
+
+    let uri = server.uri();
+    Mock::given(method("POST"))
+        .and(path("/api/2.0/sql/statements"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "statement_id": STATEMENT_ID,
+            "status": {"state": "SUCCEEDED"},
+            "manifest": {
+                "chunks": [{"chunk_index": 0, "row_count": 5}],
+                "schema": {"columns": [
+                    {"name": "id", "type_name": "LONG"},
+                    {"name": "label", "type_name": "STRING"},
+                ]},
+            },
+            // "chunk_index" deliberately omitted, not set to 0.
+            "result": {
+                "external_links": [{"external_link": format!("{uri}/_data/chunk-0")}]
+            },
+        })))
+        .mount(&server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path(format!("/api/2.0/sql/statements/{STATEMENT_ID}/result/chunks/0")))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"external_links": []})))
+        .expect(0)
+        .mount(&server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path_regex(r"^/_data/chunk-0$"))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_raw(build_chunk_bytes(0, 5), "application/vnd.apache.arrow.stream"),
+        )
+        .mount(&server)
+        .await;
+
+    let client = Arc::new(DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token"));
+    let summary = run_pipeline(client, "SELECT * FROM t", None, None, None).await.unwrap();
+
+    assert_eq!(summary.num_rows(), 5);
+    assert_ids_in_order(&summary.batches, 5);
+}
+
 #[tokio::test]
 async fn full_pipeline_happy_path() {
     let server = MockServer::start().await;

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/src/arrowbricks/_core.pyi

@@ -41,7 +41,7 @@ class Client:
         warehouse_id: str,
         token: str | None = None,
         token_provider: Callable[[], str | Awaitable[str]] | None = None,
-        chunk_fetch_concurrency: int = 32,
+        chunk_fetch_concurrency: int = 64,
         http_timeout: float = 60.0,
         wait_timeout: str = "30s",
         warehouse_start_timeout: float = 300.0,

{arrowbricks-1.2.0 → arrowbricks-1.3.0}/src/arrowbricks/client.py

@@ -33,7 +33,7 @@ class DatabricksClient:
         token_provider: TokenProvider | None = None,
         http_timeout: float = 60.0,
         wait_timeout: str = "30s",
-        chunk_fetch_concurrency: int = 32,
+        chunk_fetch_concurrency: int = 64,
         warehouse_start_timeout: float = 300.0,
         warehouse_confirmed_running_ttl_s: float = 30.0,
         compress_results: bool = True,
@@ -44,8 +44,10 @@ class DatabricksClient:
         self.http_timeout = http_timeout
         self.wait_timeout = wait_timeout
         # Each concurrent fetch slot holds a whole chunk's raw bytes in
-        # memory. 32 -- tuned for the Rust core's real OS-thread parallelism,
-        # which scales well past what asyncio+GIL concurrency used to buy.
+        # memory. 64 -- tuned for the Rust core's real OS-thread parallelism,
+        # which scales well past what asyncio+GIL concurrency used to buy;
+        # measured against a real 400-chunk/5.6M-row table (see client.rs's
+        # own comment for the numbers).
         self.chunk_fetch_concurrency = chunk_fetch_concurrency
         self.warehouse_start_timeout = warehouse_start_timeout
         # Requests LZ4-compressed cloud-fetch chunks (matches