arrowbricks 3.1.2__tar.gz → 3.1.3__tar.gz

{arrowbricks-3.1.2 → arrowbricks-3.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 3.1.2
+Version: 3.1.3
 Requires-Dist: arro3-core>=0.8 ; extra == 'arro3'
 Provides-Extra: arro3
 License-File: LICENSE
@@ -271,12 +271,17 @@ The [official driver](https://github.com/databricks/databricks-sql-python) is th
 
 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:
 
-| | avg | stdev | range | peak RSS during the query |
-|---|---|---|---|---|
-| `databricks-sql-connector` | 0.88s | 0.06s | 0.83s - 0.95s | 16 MB |
-| arrowbricks | 0.50s | 0.02s | 0.48s - 0.52s | 16 MB |
+| | avg | stdev | range |
+|---|---|---|---|
+| `databricks-sql-connector` | 0.88s | 0.06s | 0.83s - 0.95s |
+| arrowbricks | 0.50s | 0.02s | 0.48s - 0.52s |
 
-**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:
+These historical latency measurements compare the connector's `fetchall()`
+(Python rows) with arrowbricks' `fetchall_arrow()` (Arrow), which do different
+amounts of materialization. The previously published 16 MB memory figures
+were invalid: the script sampled peak RSS before importing the libraries or
+running queries. It now measures after the workload; rerun it for actual
+peak memory on your machine. The installed-footprint measurements were:
 
 | | installed size (package + all required deps) |
 |---|---|
@@ -292,6 +297,25 @@ DATABRICKS_TOKEN=dapiXXXXXXXXXXXXXXXXXXXXXXXXXXXX \
 python examples/benchmark_vs_connector.py
 ```
 
+For comparing arrowbricks versions with identical APIs, use
+[`examples/benchmark_versions.py`](examples/benchmark_versions.py). It keeps
+one connection per version, discards warm-ups, alternates execution order,
+passes concurrency explicitly, and checks row counts. See
+[`benchmarks/2026-09-06.md`](benchmarks/2026-09-06.md) for replay/cache measurements
+and [`benchmarks/2026-09-06-downloads.md`](benchmarks/2026-09-06-downloads.md)
+for the subsequent cloud-fetch scheduling measurements and limits.
+
+Cached IPC replay can be measured without a warehouse:
+
+```bash
+uv run python examples/benchmark_replay.py --rows 200000 --columns 16 --repeats 8
+```
+
+Replays now share immutable input bytes across decoded tables, reducing
+repeated copies and memory use. Arrow may still copy misaligned fixed-width
+buffers or decompress IPC-compressed bodies. Keeping a small slice of a
+decoded array can retain the full source allocation until that slice is released.
+
 `BENCHMARK_SQL` overrides the query, `BENCHMARK_RUNS` (default 3) controls how many timed runs to average.
 
 ## A note on Arrow IPC compression

{arrowbricks-3.1.2 → arrowbricks-3.1.3}/README.md

@@ -259,12 +259,17 @@ The [official driver](https://github.com/databricks/databricks-sql-python) is th
 
 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:
 
-| | avg | stdev | range | peak RSS during the query |
-|---|---|---|---|---|
-| `databricks-sql-connector` | 0.88s | 0.06s | 0.83s - 0.95s | 16 MB |
-| arrowbricks | 0.50s | 0.02s | 0.48s - 0.52s | 16 MB |
-
-**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:
+| | avg | stdev | range |
+|---|---|---|---|
+| `databricks-sql-connector` | 0.88s | 0.06s | 0.83s - 0.95s |
+| arrowbricks | 0.50s | 0.02s | 0.48s - 0.52s |
+
+These historical latency measurements compare the connector's `fetchall()`
+(Python rows) with arrowbricks' `fetchall_arrow()` (Arrow), which do different
+amounts of materialization. The previously published 16 MB memory figures
+were invalid: the script sampled peak RSS before importing the libraries or
+running queries. It now measures after the workload; rerun it for actual
+peak memory on your machine. The installed-footprint measurements were:
 
 | | installed size (package + all required deps) |
 |---|---|
@@ -280,6 +285,25 @@ DATABRICKS_TOKEN=dapiXXXXXXXXXXXXXXXXXXXXXXXXXXXX \
 python examples/benchmark_vs_connector.py
 ```
 
+For comparing arrowbricks versions with identical APIs, use
+[`examples/benchmark_versions.py`](examples/benchmark_versions.py). It keeps
+one connection per version, discards warm-ups, alternates execution order,
+passes concurrency explicitly, and checks row counts. See
+[`benchmarks/2026-09-06.md`](benchmarks/2026-09-06.md) for replay/cache measurements
+and [`benchmarks/2026-09-06-downloads.md`](benchmarks/2026-09-06-downloads.md)
+for the subsequent cloud-fetch scheduling measurements and limits.
+
+Cached IPC replay can be measured without a warehouse:
+
+```bash
+uv run python examples/benchmark_replay.py --rows 200000 --columns 16 --repeats 8
+```
+
+Replays now share immutable input bytes across decoded tables, reducing
+repeated copies and memory use. Arrow may still copy misaligned fixed-width
+buffers or decompress IPC-compressed bodies. Keeping a small slice of a
+decoded array can retain the full source allocation until that slice is released.
+
 `BENCHMARK_SQL` overrides the query, `BENCHMARK_RUNS` (default 3) controls how many timed runs to average.
 
 ## A note on Arrow IPC compression

{arrowbricks-3.1.2 → arrowbricks-3.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "3.1.2"
+version = "3.1.3"
 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-3.1.2 → arrowbricks-3.1.3}/rust/arrowbricks_core/Cargo.lock

@@ -34,39 +34,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "arrow"
-version = "59.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "61d285d16bce7d0be61912f7928342b673067b6b7d7ef6cc179258ba7de1fecf"
-dependencies = [
- "arrow-arith",
- "arrow-array",
- "arrow-buffer",
- "arrow-cast",
- "arrow-data",
- "arrow-ipc",
- "arrow-ord",
- "arrow-row",
- "arrow-schema",
- "arrow-select",
- "arrow-string",
-]
-
-[[package]]
-name = "arrow-arith"
-version = "59.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "757ef1836251e88222542a7da2623bc1c9cb9e20afefa6db2c41e79991cd91d4"
-dependencies = [
- "arrow-array",
- "arrow-buffer",
- "arrow-data",
- "arrow-schema",
- "chrono",
- "num-traits",
-]
-
 [[package]]
 name = "arrow-array"
 version = "59.2.0"
@@ -186,19 +153,6 @@ dependencies = [
  "arrow-select",
 ]
 
-[[package]]
-name = "arrow-row"
-version = "59.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bbec439386df71ad570e6758a946111322b9e9dc8db83b5527321f0b4c9119c2"
-dependencies = [
- "arrow-array",
- "arrow-buffer",
- "arrow-data",
- "arrow-schema",
- "half",
-]
-
 [[package]]
 name = "arrow-schema"
 version = "59.2.0"
@@ -223,28 +177,14 @@ dependencies = [
 ]
 
 [[package]]
-name = "arrow-string"
-version = "59.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c838a25bb3691e919e0f617616ac51a4ff8517a952e29ca133cf0c22b2ce65b1"
+name = "arrowbricks_core"
+version = "3.1.3"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
- "arrow-data",
- "arrow-schema",
- "arrow-select",
- "memchr",
- "num-traits",
- "regex",
- "regex-syntax",
-]
-
-[[package]]
-name = "arrowbricks_core"
-version = "3.1.2"
-dependencies = [
- "arrow",
+ "arrow-ipc",
  "arrow-json",
+ "arrow-schema",
  "base64 0.23.1",
  "bytes",
  "chrono",

{arrowbricks-3.1.2 → arrowbricks-3.1.3}/rust/arrowbricks_core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "arrowbricks_core"
-version = "3.1.2"
+version = "3.1.3"
 edition = "2024"
 readme = "README.md"
 
@@ -9,10 +9,12 @@ name = "arrowbricks_core"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-# default-features off: arrow's defaults are csv+ipc+json, and we only ever
-# decode Arrow-IPC chunks -- csv/json parsing (plus their transitive deps)
-# would just be dead weight in the shipped .so.
-arrow = { version = "59.1.0", default-features = false, features = ["ipc"] }
+# Depend on the components we use. The umbrella `arrow` crate also builds
+# arithmetic, string, and row kernels even with default-features disabled.
+arrow-array = "59.1.0"
+arrow-buffer = "59.1.0"
+arrow-ipc = "59.1.0"
+arrow-schema = "59.1.0"
 arrow-json = "59.1.0"
 bytes = "1.12.1"
 # Already an unconditional transitive dependency of arrow-cast (pulled in by
@@ -103,6 +105,11 @@ redundant_clone = "warn"
 large_enum_variant = "warn"
 needless_collect = "warn"
 
+[profile.dev]
+# Keep source locations in backtraces without caching full debug type data
+# for Arrow and the HTTP stack. Incremental compilation remains enabled.
+debug = "line-tables-only"
+
 [profile.release]
 # opt-level="s" over the implicit default (3) or "z" -- a size-vs-speed
 # tradeoff previously rejected here on the reasoning that "3" was what the
@@ -126,3 +133,9 @@ opt-level = "s"
 strip = true
 lto = true
 codegen-units = 1
+
+[profile.release.package.pyo3-arrow]
+# This bridge is dominated by interoperability helpers. Optimizing it for
+# size saves ~32 KiB on macOS arm64; the IPC decoder keeps opt-level="s".
+# Replay and live-query benchmarks are recorded in benchmarks/2026-09-06.md.
+opt-level = "z"

{arrowbricks-3.1.2 → arrowbricks-3.1.3}/rust/arrowbricks_core/README.md

@@ -32,6 +32,12 @@ This builds a mixed Python/Rust project: the compiled extension lands at
 `mypy`/`pyright`/`ty` type-checking -- a compiled PyO3 extension has no type
 info of its own without them.
 
+Development/test builds retain source line information for backtraces but
+omit full debug type data to reduce the Cargo cache. For full debugger
+variable inspection, run with `CARGO_PROFILE_DEV_DEBUG=2`. Incremental
+compilation stays enabled. Release builds optimize the Python/Arrow bridge
+for size while keeping the IPC decoder at its existing optimization level.
+
 ## Quickstart
 
 Importable directly (bypassing the `arrowbricks` Python wrapper) as

{arrowbricks-3.1.2 → arrowbricks-3.1.3}/rust/arrowbricks_core/src/client/download.rs

@@ -121,28 +121,24 @@ impl DbClient {
 
     /// Downloads one cloud-fetch link, splitting it across parallel HTTP
     /// Range requests when the shared `download_slots` budget has room to
-    /// spare -- see that field's own doc comment for the measured win and
-    /// why the large-result case is safe.
+    /// spare. `split_limit` reserves a fair share for the other links in
+    /// the same discovered batch; available permits alone do not tell us
+    /// how many workers have yet to start requesting their own slots.
     pub(crate) async fn fetch_link_bytes_budgeted(
         self: &Arc<Self>,
         url: &str,
         compressed: bool,
+        split_limit: usize,
         stats: &Arc<QueryStatsAccumulator>,
     ) -> 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. `acquire()`'s `Err` case (the semaphore closed) can't
-        // happen -- nothing in this crate ever calls `.close()` on
-        // `download_slots` -- but propagating it as a real error instead of
-        // an `.expect()` costs nothing and avoids a panic if that ever
-        // changes.
+        // One permit per file is mandatory. This can wait because other
+        // files' extra Range requests consume the same shared budget.
         let _base = self
             .download_slots
             .acquire()
             .await
             .map_err(|e| ApiError::permanent(format!("download_slots semaphore closed unexpectedly: {e}")))?;
-        let want = (MAX_SPLIT_PARTS - 1).min(self.download_slots.available_permits());
+        let want = (MAX_SPLIT_PARTS.min(split_limit.max(1)) - 1).min(self.download_slots.available_permits());
         let extra = if want > 0 {
             self.download_slots.try_acquire_many(want as u32).ok()
         } else {
@@ -156,8 +152,8 @@ impl DbClient {
             .await
     }
 
-    /// Downloads one cloud-fetch link as concurrent HTTP Range requests of
-    /// `part_size` bytes each, concatenated in order. The real object size
+    /// Downloads a `part_size` probe and spreads the remainder across
+    /// bounded parallel Range requests, 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.
@@ -181,10 +177,11 @@ impl DbClient {
         max_parts: u64,
         stats: &Arc<QueryStatsAccumulator>,
     ) -> 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) = self
+        // Retry the probe as one request (headers and body). Start tail
+        // downloads as soon as its headers reveal the file size, so reading
+        // the first MiB overlaps the rest. JoinSet aborts those requests if
+        // the probe fails and retries, or the surrounding fetch is cancelled.
+        let (total, head, mut downloads) = self
             .retry_call_tracked(Some(stats), || async {
                 let resp = self
                     .http
@@ -205,96 +202,84 @@ impl DbClient {
                     .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?;
-        stats.bytes_downloaded.fetch_add(head.len() as u64, Ordering::Relaxed);
-
-        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(),
-            ));
-        }
+                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",
+                    ));
+                }
 
-        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();
-                let part_stats = stats.clone();
-                handles.push(tokio::spawn(async move {
-                    let bytes = this
-                        .retry_call_tracked(Some(&part_stats), || 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));
+                let mut downloads = tokio::task::JoinSet::new();
+                if let Some(total) = ranged.then_some(total).flatten() {
+                    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;
+                    for index in 0..n_rest {
+                        if start >= total {
+                            break;
+                        }
+                        let end = (start + rest_size - 1).min(total - 1);
+                        let this = self.clone();
+                        let url = url.to_string();
+                        let part_stats = stats.clone();
+                        downloads.spawn(async move {
+                            let bytes = this
+                                .retry_call_tracked(Some(&part_stats), || 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;
+                            if let Ok(b) = &bytes {
+                                part_stats.bytes_downloaded.fetch_add(b.len() as u64, Ordering::Relaxed);
                             }
-                            resp.bytes().await.map_err(|e| ApiError::from_reqwest(e, true))
-                        })
-                        .await;
-                    if let Ok(b) = &bytes {
-                        part_stats.bytes_downloaded.fetch_add(b.len() as u64, Ordering::Relaxed);
-                    }
-                    bytes
-                }));
-                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);
-        // Parts must concatenate in order, so this can't use `JoinSet`
-        // (which yields in completion order) without tracking indices --
-        // simpler to keep the `Vec` and explicitly `.abort()` every
-        // not-yet-awaited sibling the moment one part fails, rather than
-        // silently leaving them running (a bare `?` here would return
-        // early and just drop the rest, which does NOT cancel them --
-        // `JoinHandle::drop` detaches, it doesn't abort -- leaving up to
-        // `MAX_SPLIT_PARTS - 1` sibling Range downloads, each with their
-        // own `retry_call` backoff, still in flight for a link the caller
-        // has already given up on).
-        let mut iter = handles.into_iter();
-        while let Some(h) = iter.next() {
-            match h.await {
-                Ok(Ok(part)) => out.extend_from_slice(&part),
-                Ok(Err(e)) => {
-                    for remaining in iter {
-                        remaining.abort();
+                            (index, bytes)
+                        });
+                        start = end + 1;
                     }
-                    return Err(e);
                 }
-                Err(join_err) => {
-                    for remaining in iter {
-                        remaining.abort();
+                let head = match resp.bytes().await {
+                    Ok(bytes) => bytes,
+                    Err(error) => {
+                        // Finish cancelling the old parts before retrying
+                        // the probe under the same download-slot budget.
+                        downloads.shutdown().await;
+                        return Err(ApiError::from_reqwest(error, true));
                     }
-                    return Err(join_error(join_err));
-                }
+                };
+                Ok((total, head, downloads))
+            })
+            .await?;
+        stats.bytes_downloaded.fetch_add(head.len() as u64, Ordering::Relaxed);
+
+        let bytes = if downloads.is_empty() {
+            head
+        } else {
+            let mut parts = Vec::with_capacity(downloads.len());
+            while let Some(result) = downloads.join_next().await {
+                let (index, bytes) = result.map_err(join_error)?;
+                parts.push((index, bytes?));
             }
-        }
-        let bytes = out.freeze();
+            parts.sort_unstable_by_key(|(index, _)| *index);
+            let mut out = bytes::BytesMut::with_capacity(total.unwrap_or(head.len() as u64) as usize);
+            out.extend_from_slice(&head);
+            for (_, part) in parts {
+                out.extend_from_slice(&part);
+            }
+            out.freeze()
+        };
         if let Some(t) = total
             && bytes.len() as u64 != t
         {
@@ -316,6 +301,71 @@ impl DbClient {
 mod tests {
     use super::*;
 
+    #[tokio::test]
+    async fn split_download_starts_tail_requests_before_the_probe_body_finishes() {
+        use tokio::io::{AsyncReadExt, AsyncWriteExt};
+        use tokio::net::TcpListener;
+
+        for fail_first_probe in [false, true] {
+            let probes = Arc::new(std::sync::atomic::AtomicUsize::new(0));
+            let server_probes = probes.clone();
+            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+            let address = listener.local_addr().unwrap();
+            let tail_started = Arc::new(tokio::sync::Notify::new());
+            tokio::spawn(async move {
+                for _ in 0..if fail_first_probe { 6 } else { 3 } {
+                    let (mut socket, _) = listener.accept().await.unwrap();
+                    let tail_started = tail_started.clone();
+                    let probes = server_probes.clone();
+                    tokio::spawn(async move {
+                        let mut request = Vec::new();
+                        while !request.ends_with(b"\r\n\r\n") {
+                            request.push(socket.read_u8().await.unwrap());
+                        }
+                        let request = String::from_utf8(request).unwrap().to_ascii_lowercase();
+                        let range = request.lines().find(|s| s.starts_with("range:")).unwrap();
+                        let (start, end) = range.split_once("bytes=").unwrap().1.split_once('-').unwrap();
+                        let start: usize = start.parse().unwrap();
+                        let end: usize = end.parse().unwrap();
+                        let headers = format!(
+                            "HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {start}-{end}/12\r\nConnection: close\r\n\r\n",
+                            end - start + 1,
+                        );
+                        socket.write_all(headers.as_bytes()).await.unwrap();
+                        if start == 0 {
+                            let first = probes.fetch_add(1, Ordering::Relaxed) == 0;
+                            tail_started.notified().await;
+                            if fail_first_probe && first {
+                                socket.write_all(b"a").await.unwrap();
+                                return; // truncated body: the whole probe must retry
+                            }
+                        } else {
+                            tail_started.notify_one();
+                        }
+                        let _ = socket.write_all(&b"abcdefghijkl"[start..=end]).await;
+                    });
+                }
+            });
+            let client = Arc::new(
+                DbClient::new(&format!("http://{address}"), "wh", "token")
+                    .with_retry_attempts(2)
+                    .with_retry_max_wait_s(0.0),
+            );
+            let stats = Arc::new(QueryStatsAccumulator::default());
+            let bytes = tokio::time::timeout(
+                std::time::Duration::from_secs(2),
+                client.fetch_link_bytes_split(&format!("http://{address}/blob"), false, 4, 3, &stats),
+            )
+            .await
+            .expect("tail requests must start from the probe headers, without waiting for its body")
+            .unwrap();
+            assert_eq!(&bytes[..], b"abcdefghijkl");
+            assert_eq!(probes.load(Ordering::Relaxed), if fail_first_probe { 2 } else { 1 });
+            assert_eq!(stats.retry_count.load(Ordering::Relaxed), u32::from(fail_first_probe));
+            assert!(stats.bytes_downloaded.load(Ordering::Relaxed) >= 12);
+        }
+    }
+
     /// 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-3.1.2 → arrowbricks-3.1.3}/rust/arrowbricks_core/src/client.rs

@@ -416,7 +416,7 @@ impl DbClient {
             // aws-lc-rs to ring, on the theory that either transport change
             // could have moved the optimum: a first pass reported 96 as
             // ~13-16% faster than 64 on both this table and a second,
-            // larger one (fact_sales_order_invoiced, 20M rows/295 cols) --
+            // larger one (large_benchmark_table, 20M rows/295 cols) --
             // **found on review to be a false positive**. The benchmark
             // script called `connect()`/`arrowbricks.connect()` without
             // ever passing `chunk_fetch_concurrency=` explicitly, so every
@@ -431,7 +431,7 @@ impl DbClient {
             // run-to-run noise, not a code effect. Caught by re-running a
             // controlled, interleaved A/B (`chunk_fetch_concurrency=`
             // passed explicitly each time, no rebuild needed) on
-            // `dim_article`: 64/96/64/96 measured 125.60s/129.92s/137.50s/
+            // `benchmark_table`: 64/96/64/96 measured 125.60s/129.92s/137.50s/
             // 134.30s -- no consistent winner, well within run-to-run
             // noise. Reverted to 64. If re-attempting this again, always
             // pass `chunk_fetch_concurrency=` explicitly in the benchmark

{arrowbricks-3.1.2 → arrowbricks-3.1.3}/rust/arrowbricks_core/src/json_convert.rs

@@ -17,13 +17,14 @@
 
 use std::sync::Arc;
 
-use arrow::array::{
+use arrow_array::RecordBatch;
+use arrow_array::types::Date32Type;
+use arrow_array::{
     ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Float32Array, Float64Array, Int8Array,
     Int16Array, Int32Array, Int64Array, StringArray, StructArray, TimestampMicrosecondArray,
 };
-use arrow::buffer::NullBuffer;
-use arrow::datatypes::{DataType, Date32Type, Field, Fields, Schema, TimeUnit};
-use arrow::record_batch::RecordBatch;
+use arrow_buffer::NullBuffer;
+use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit};
 use chrono::{DateTime, NaiveDate, NaiveDateTime};
 use serde_json::Value as JsonValue;
 
@@ -524,7 +525,7 @@ mod tests {
         assert_eq!(batch.num_rows(), 1);
         assert_eq!(batch.num_columns(), 13);
 
-        use arrow::array::Array;
+        use arrow_array::Array;
         assert_eq!(
             batch.column(0).as_any().downcast_ref::<Int8Array>().unwrap().value(0),
             1
@@ -635,7 +636,7 @@ mod tests {
         )]];
 
         let batch = json_array_to_record_batch(&rows, &columns).unwrap();
-        use arrow::array::Array;
+        use arrow_array::Array;
         let s = batch.column(0).as_any().downcast_ref::<StructArray>().unwrap();
         assert_eq!(
             s.column_by_name("a")
@@ -691,7 +692,7 @@ mod tests {
             vec![None], // the whole struct is NULL for this row
         ];
         let batch = json_array_to_record_batch(&rows, &columns).unwrap();
-        use arrow::array::Array;
+        use arrow_array::Array;
         let s = batch.column(0).as_any().downcast_ref::<StructArray>().unwrap();
         assert!(
             s.column_by_name("a").unwrap().is_null(0),