arrowbricks 3.1.3__tar.gz → 3.2.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 3.1.3
+Version: 3.2.0
 Requires-Dist: arro3-core>=0.8 ; extra == 'arro3'
 Provides-Extra: arro3
 License-File: LICENSE
@@ -160,9 +160,9 @@ conn = connect(host=..., warehouse_id=..., token=..., on_event=log_query)
 
 ## Cancellation
 
-When a chunk download times out (`total_timeout_s`) or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while it's in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
+When `total_timeout_s` elapses or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while a query is in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
 
-This covers `Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, `client.stream_query_json(...)`, and the lower-level `._core.Client`'s own streamed APIs -- everything that wraps the chunk-download phase in the Rust-level heartbeat. Neither `total_timeout_s` nor a bare `task.cancel()` on `Cursor.execute()`/`execute_streamed()` itself (the initial submit/poll wait) triggers a cancel -- by the time a query's chunks are being fetched at all, the statement has typically already finished running server-side, so there's usually nothing left to cancel there in practice.
+This covers both phases of a query: the submit/poll wait while the statement is still running (`Cursor.execute()`/`execute_streamed()`, and the start of `client.stream_query_json(...)`), and the chunk download afterwards (`Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, the rest of `stream_query_json`). `stream_query_json`'s `total_timeout_s` is one budget across both phases. On `protocol="sea"` there is one short blind spot: the submit request itself can wait server-side for up to `wait_timeout` (default 30s) before Databricks returns a statement id, and a query abandoned inside that window can't be cancelled.
 
 ## Errors
 
@@ -300,11 +300,19 @@ 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
+passes concurrency explicitly, and checks row counts. Add `--verify-ipc`
+(requires `arro3-core`) to compare serialized Arrow results in memory after
+timing. Results must have stable values, schema metadata, row order, and batch boundaries;
+verification contributes to peak process memory, and checksums are not printed. 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.
 
+Follow-up experiments cover [spare request slots](benchmarks/2026-09-06-spare-slots.md),
+[concurrent replay and NDJSON encoding](benchmarks/2026-09-06-lowlevel.md),
+and the rejected [LZ4 capacity](benchmarks/2026-09-06-lz4-capacity.md) and
+[direct-fill buffer](benchmarks/2026-09-06-direct-fill.md) changes.
+
 Cached IPC replay can be measured without a warehouse:
 
 ```bash
@@ -315,6 +323,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.
+Decoding releases the Python interpreter lock, so independent replay calls
+can run concurrently on separate Python threads.
 
 `BENCHMARK_SQL` overrides the query, `BENCHMARK_RUNS` (default 3) controls how many timed runs to average.
 

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

@@ -148,9 +148,9 @@ conn = connect(host=..., warehouse_id=..., token=..., on_event=log_query)
 
 ## Cancellation
 
-When a chunk download times out (`total_timeout_s`) or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while it's in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
+When `total_timeout_s` elapses or your code cancels the surrounding coroutine (`task.cancel()`/`asyncio.wait_for`) while a query is in flight, arrowbricks fires a best-effort server-side cancel in the background (Thrift's `CancelOperation`, or SEA's `POST .../cancel`) so Databricks stops running the query instead of finishing it for nobody. This is fire-and-forget: the `QueryTimeout`/cancellation still reaches you immediately, and the cancel call's own result (success or failure) is never surfaced or awaited.
 
-This covers `Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, `client.stream_query_json(...)`, and the lower-level `._core.Client`'s own streamed APIs -- everything that wraps the chunk-download phase in the Rust-level heartbeat. Neither `total_timeout_s` nor a bare `task.cancel()` on `Cursor.execute()`/`execute_streamed()` itself (the initial submit/poll wait) triggers a cancel -- by the time a query's chunks are being fetched at all, the statement has typically already finished running server-side, so there's usually nothing left to cancel there in practice.
+This covers both phases of a query: the submit/poll wait while the statement is still running (`Cursor.execute()`/`execute_streamed()`, and the start of `client.stream_query_json(...)`), and the chunk download afterwards (`Cursor.fetchall_streamed()`/`fetchall_arrow_streamed()`, the rest of `stream_query_json`). `stream_query_json`'s `total_timeout_s` is one budget across both phases. On `protocol="sea"` there is one short blind spot: the submit request itself can wait server-side for up to `wait_timeout` (default 30s) before Databricks returns a statement id, and a query abandoned inside that window can't be cancelled.
 
 ## Errors
 
@@ -288,11 +288,19 @@ 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
+passes concurrency explicitly, and checks row counts. Add `--verify-ipc`
+(requires `arro3-core`) to compare serialized Arrow results in memory after
+timing. Results must have stable values, schema metadata, row order, and batch boundaries;
+verification contributes to peak process memory, and checksums are not printed. 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.
 
+Follow-up experiments cover [spare request slots](benchmarks/2026-09-06-spare-slots.md),
+[concurrent replay and NDJSON encoding](benchmarks/2026-09-06-lowlevel.md),
+and the rejected [LZ4 capacity](benchmarks/2026-09-06-lz4-capacity.md) and
+[direct-fill buffer](benchmarks/2026-09-06-direct-fill.md) changes.
+
 Cached IPC replay can be measured without a warehouse:
 
 ```bash
@@ -303,6 +311,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.
+Decoding releases the Python interpreter lock, so independent replay calls
+can run concurrently on separate Python threads.
 
 `BENCHMARK_SQL` overrides the query, `BENCHMARK_RUNS` (default 3) controls how many timed runs to average.
 

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

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

@@ -178,7 +178,7 @@ dependencies = [
 
 [[package]]
 name = "arrowbricks_core"
-version = "3.1.3"
+version = "3.2.0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -190,6 +190,7 @@ dependencies = [
  "chrono",
  "hyper-rustls",
  "lz4_flex",
+ "memchr",
  "proptest",
  "pyo3",
  "pyo3-arrow",

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

@@ -1,6 +1,6 @@
 [package]
 name = "arrowbricks_core"
-version = "3.1.3"
+version = "3.2.0"
 edition = "2024"
 readme = "README.md"
 
@@ -17,6 +17,8 @@ arrow-ipc = "59.1.0"
 arrow-schema = "59.1.0"
 arrow-json = "59.1.0"
 bytes = "1.12.1"
+# Already in the dependency graph; vectorized newline search in NDJSON output.
+memchr = "2"
 # Already an unconditional transitive dependency of arrow-cast (pulled in by
 # pyo3-arrow) at this exact version -- declared directly here too so
 # json_convert.rs can decode BINARY columns without hand-rolling base64. No

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

@@ -73,7 +73,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="thrift", retry_attempts=6, retry_max_wait_s=20.0)` -- exactly one of `token`/`token_provider`. `retry_attempts`/`retry_max_wait_s` tune the retry policy (total attempts, exponential-backoff ceiling in seconds) behind every retryable request this client makes; `retry_attempts` must be at least 1. `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.
+- `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 -- the recognized INLINE byte-limit failure falls back to a second, normal `execute()`. If an already-succeeded result cannot be converted (including unsupported empty STRUCT arrays), it raises `ArrowbricksError` without resubmitting the statement. Use the byte-limit retry only with SQL that is safe to execute again. 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.
   - `ResultSet.fetchall_arrow_streamed(*, total_timeout_s=None)` -- same as `fetchall_arrow()`, but an async iterator yielding the `HEARTBEAT` singleton while pulling chunks instead of blocking silently (bridge e.g. an SSE connection through the download), then a `Table` exactly once. Raises if `total_timeout_s` elapses first.
@@ -82,7 +82,7 @@ unlike the table object itself.
 - `Client.stream_ndjson_lines(statement, *, catalog=None, schema=None, parameters=None, total_timeout_s=None)` -- chunk-at-a-time: an async iterator yielding the `HEARTBEAT` singleton while waiting on the statement or any individual chunk, then a `list[str]` of NDJSON lines (one per row, explicit nulls, ISO-8601 timestamps) per chunk in logical order. Decode and JSON encoding both happen in Rust -- backs `stream_query_json` end to end.
 - `Client.upload_volume_file(volume_path, data: bytes)` / `Client.delete_volume_file(volume_path)` -- Unity Catalog volume files via the Files API. Delete treats a 404 as success (idempotent). Both raise a plain `RuntimeError` (message only) on failure.
 - `write_ipc_stream(stream, buf)` -- free function; writes any object implementing `__arrow_c_stream__` (a `Table` from this crate, arro3, pyarrow, ...) as uncompressed Arrow-IPC stream bytes to a Python file-like object. No dependency needed regardless of the input's origin.
-- `read_ipc_stream(data: bytes) -> Table` -- free function; the exact inverse of `write_ipc_stream`, parsing raw Arrow-IPC stream bytes back into a `Table`. No dependency needed regardless of where the bytes came from -- backs `arrowbricks.ReplayableArrowChunk`, which needs to re-parse the same cached bytes on every `__arrow_c_stream__` call.
+- `read_ipc_stream(data: bytes) -> Table` -- free function; the exact inverse of `write_ipc_stream`, parsing raw Arrow-IPC stream bytes back into a `Table`. No dependency needed regardless of where the bytes came from -- backs `arrowbricks.ReplayableArrowChunk`, which needs to re-parse the same cached bytes on every `__arrow_c_stream__` call. Decoding releases the Python interpreter lock so independent replay calls can run concurrently; decoded arrays retain ownership of the immutable input bytes.
 - `HEARTBEAT` -- module-level singleton; compare with `is`, e.g. `if item is _core.HEARTBEAT: ...`.
 
 ## With DuckDB

arrowbricks-3.2.0/rust/arrowbricks_core/proptest-regressions/json_convert.txt

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc c23ce2dbe22693eff6644d643de77c6c113dc918c876434391fb93a05438d706 # shrinks to type_name = "STRUCT", precision = None, scale = None, type_text = Some(""), values = []

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

@@ -36,18 +36,15 @@ use super::model::QueryStatsAccumulator;
 /// after each `EndMark` and picks up the next concatenated frame on a
 /// subsequent `read_to_end` call against the *same* instance (verified: the
 /// decoder's position in the underlying byte slice carries over across
-/// calls) -- so looping `read_to_end` on one decoder until it stops growing
-/// `out` reads every frame without reconstructing a decoder per frame.
+/// calls) -- so looping `read_to_end` on one decoder until its underlying
+/// reader is exhausted reads every frame without reconstructing a decoder.
 pub(crate) fn decompress_lz4_frame(compressed: &Bytes) -> Result<Bytes, ApiError> {
     use std::io::Read;
-    // `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.
+    // LZ4 on Arrow-IPC data often compresses several-fold. Starting at the
+    // compressed size can therefore pay for repeated buffer growth. Neither
+    // size is a bound on the other: incompressible input can expand slightly.
+    // `* 4` remains a heuristic; Vec grows normally when it underestimates.
+    // See benchmark_lz4_capacity_hypotheses for the allocation tradeoffs.
     let mut out = Vec::with_capacity(compressed.len() * 4);
     let mut decoder = lz4_flex::frame::FrameDecoder::new(&compressed[..]);
     // Terminate on the *reader* being exhausted, not on "output stopped
@@ -301,6 +298,65 @@ impl DbClient {
 mod tests {
     use super::*;
 
+    #[test]
+    #[ignore = "manual release-mode decompression allocation experiment"]
+    fn benchmark_lz4_capacity_hypotheses() {
+        use std::hint::black_box;
+        use std::io::{Read, Write};
+        use std::time::Instant;
+
+        // Deterministic random bytes mixed with repeated bytes cover different
+        // compression ratios without depending on warehouse data.
+        for random_fraction in [0, 25, 50, 100] {
+            let mut state = 0x12345678_u64;
+            let input: Vec<u8> = (0..8 * 1024 * 1024)
+                .map(|i| {
+                    state ^= state << 13;
+                    state ^= state >> 7;
+                    state ^= state << 17;
+                    if i % 100 < random_fraction { state as u8 } else { 0 }
+                })
+                .collect();
+            let mut compressed = Vec::new();
+            for part in input.chunks(512 * 1024) {
+                let mut encoder = lz4_flex::frame::FrameEncoder::new(Vec::new());
+                encoder.write_all(part).unwrap();
+                compressed.extend(encoder.finish().unwrap());
+            }
+            let variants = [
+                ("compressed_x1", compressed.len()),
+                ("compressed_x2", compressed.len() * 2),
+                ("compressed_x4", compressed.len() * 4),
+                ("exact", input.len()),
+                ("short_hint", input.len() - 64),
+            ];
+            for round in 0..10 {
+                for index in 0..variants.len() {
+                    let (variant, capacity) = variants[(index + round) % variants.len()];
+                    let start = Instant::now();
+                    let mut output = Vec::with_capacity(capacity);
+                    let mut decoder = lz4_flex::frame::FrameDecoder::new(black_box(compressed.as_slice()));
+                    while !decoder.get_ref().is_empty() {
+                        decoder.read_to_end(&mut output).unwrap();
+                    }
+                    let seconds = start.elapsed().as_secs_f64();
+                    assert_eq!(output, input);
+                    if round > 1 {
+                        println!(
+                            "{}",
+                            serde_json::json!({
+                                "random_percent": random_fraction, "variant": variant,
+                                "round": round, "seconds": seconds,
+                                "compressed_bytes": compressed.len(), "decoded_bytes": output.len(),
+                                "retained_capacity": output.capacity()
+                            })
+                        );
+                    }
+                }
+            }
+        }
+    }
+
     #[tokio::test]
     async fn split_download_starts_tail_requests_before_the_probe_body_finishes() {
         use tokio::io::{AsyncReadExt, AsyncWriteExt};

{arrowbricks-3.1.3 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/client/model.rs

@@ -6,6 +6,7 @@
 //! statement gets you before any chunk is fetched (`ChunkMeta`,
 //! `StatementSubmitResult`, `InlineOrExternal`, `ChunkItem`).
 
+use std::sync::Mutex;
 use std::sync::atomic::{AtomicU8, AtomicU32, AtomicU64, AtomicUsize, Ordering};
 
 use bytes::Bytes;
@@ -55,7 +56,7 @@ pub struct ColumnDescription {
 /// `Drop for HeartbeatWait`/`Drop for HeartbeatStream`) can fire
 /// `DbClient::cancel_statement` without either of those generic structs
 /// needing to know anything protocol-specific themselves.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum CancelHandle {
     Sea { statement_id: String },
     Thrift { operation: thrift::OperationHandle },
@@ -113,6 +114,17 @@ pub struct QueryStatsAccumulator {
     /// anything itself.
     warehouse_wait_bits: AtomicU64,
     outcome: AtomicU8,
+    /// The statement/operation handle while it's submitted but not yet
+    /// terminal -- set by the submit/poll loops (`client/sea.rs`'s
+    /// `submit_and_poll_inner`, `pipeline/thrift_exec.rs`'s
+    /// `submit_and_await_thrift_statement`) as soon as Databricks hands one
+    /// back, cleared once the statement reaches a terminal state. Whatever
+    /// is still here when `pipeline/stats.rs`'s `CancelInFlightOnDrop`
+    /// drops (the submit future was abandoned -- a `total_timeout_s`, a
+    /// Python-side `task.cancel()`/`asyncio.wait_for` -- or a poll failed
+    /// mid-wait) gets a best-effort server-side cancel, so a query nobody
+    /// is waiting for anymore stops running on the warehouse.
+    in_flight: Mutex<Option<CancelHandle>>,
 }
 
 const OUTCOME_UNSET: u8 = 0;
@@ -170,6 +182,18 @@ impl QueryStatsAccumulator {
     pub fn warehouse_wait_s(&self) -> f64 {
         f64::from_bits(self.warehouse_wait_bits.load(Ordering::Relaxed))
     }
+
+    pub fn set_in_flight(&self, handle: CancelHandle) {
+        *self.in_flight.lock().unwrap() = Some(handle);
+    }
+
+    pub fn clear_in_flight(&self) {
+        *self.in_flight.lock().unwrap() = None;
+    }
+
+    pub fn take_in_flight(&self) -> Option<CancelHandle> {
+        self.in_flight.lock().unwrap().take()
+    }
 }
 
 /// One query's worth of timing/counters, handed to `EventSink::on_event`

{arrowbricks-3.1.3 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/client/sea.rs

@@ -21,7 +21,8 @@ use super::DbClient;
 use super::POLL_INTERVAL;
 use super::error::{ApiError, join_error};
 use super::model::{
-    ChunkItem, ChunkMeta, ColumnDescription, InlineOrExternal, QueryStatsAccumulator, StatementSubmitResult,
+    CancelHandle, ChunkItem, ChunkMeta, ColumnDescription, InlineOrExternal, QueryStatsAccumulator,
+    StatementSubmitResult,
 };
 
 /// Typed response shapes -- replaces navigating a dynamic `serde_json::Value`
@@ -373,10 +374,14 @@ impl DbClient {
             }
         }
 
+        let mut checkin = SessionCheckin {
+            client: self,
+            catalog,
+            schema,
+            session_id,
+        };
         let result = self.submit_and_poll_inner(body, stats).await;
-        if let Some(id) = session_id {
-            self.checkin_session(catalog, schema, id, result.is_ok());
-        }
+        checkin.finish(result.is_ok());
         result
     }
 
@@ -390,6 +395,9 @@ impl DbClient {
             .authed_json(reqwest::Method::POST, &url, Some(&body), Some(stats))
             .await?;
 
+        stats.set_in_flight(CancelHandle::Sea {
+            statement_id: data.statement_id.clone(),
+        });
         while !matches!(
             data.status.state.as_str(),
             "SUCCEEDED" | "FAILED" | "CANCELED" | "CLOSED"
@@ -400,6 +408,7 @@ impl DbClient {
                 .authed_json(reqwest::Method::GET, &poll_url, None, Some(stats))
                 .await?;
         }
+        stats.clear_in_flight();
 
         match data.status.state.as_str() {
             "FAILED" => {
@@ -599,6 +608,31 @@ impl DbClient {
 /// panic case would let that worker's unfetched work vanish with no error at
 /// all: the channel closing normally looks to the consumer exactly like a
 /// complete, successful result instead of a truncated one.
+/// Returns `submit_and_poll`'s session to the pool even when its future is
+/// dropped mid-poll (a timeout or cancellation) -- without it the pool's
+/// reservation for that key leaks, and after `MAX_SESSIONS_PER_KEY` such
+/// drops every later query for the key runs session-less.
+struct SessionCheckin<'a> {
+    client: &'a DbClient,
+    catalog: Option<&'a str>,
+    schema: Option<&'a str>,
+    session_id: Option<String>,
+}
+
+impl SessionCheckin<'_> {
+    fn finish(&mut self, keep: bool) {
+        if let Some(id) = self.session_id.take() {
+            self.client.checkin_session(self.catalog, self.schema, id, keep);
+        }
+    }
+}
+
+impl Drop for SessionCheckin<'_> {
+    fn drop(&mut self) {
+        self.finish(false);
+    }
+}
+
 async fn join_first_error(handles: Vec<tokio::task::JoinHandle<Result<(), ApiError>>>) -> Option<ApiError> {
     let mut first_err = None;
     for h in handles {

{arrowbricks-3.1.3 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/heartbeat.rs

@@ -175,6 +175,10 @@ impl<T: Send + 'static> HeartbeatWait<T> {
         self
     }
 
+    pub fn deadline(&self) -> Option<Instant> {
+        self.deadline
+    }
+
     /// One step: `Ok(Some(Tick::Heartbeat))` if still waiting,
     /// `Ok(Some(Tick::Ready(value)))` exactly once when the wrapped future
     /// completes, `Ok(None)` if already exhausted (caller should raise
@@ -254,6 +258,14 @@ impl<T: Send + 'static> HeartbeatStream<T> {
         }
     }
 
+    /// Shares an existing deadline (e.g. `HeartbeatWait::deadline` of the
+    /// submit/poll wait that preceded this stream) instead of starting a
+    /// fresh `total_timeout_s` budget at construction.
+    pub fn with_deadline(mut self, deadline: Option<Instant>) -> Self {
+        self.deadline = deadline;
+        self
+    }
+
     /// See `HeartbeatWait::with_cancel`'s own doc comment -- identical
     /// contract.
     pub fn with_cancel(mut self, on_cancel: impl FnOnce(bool) + Send + 'static) -> Self {

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

@@ -334,7 +334,8 @@ fn build_column(
 
             let validity: Vec<bool> = rows_parsed.iter().map(Option::is_some).collect();
             let fields = Fields::from(child_fields);
-            let struct_array = StructArray::new(fields.clone(), child_arrays, Some(NullBuffer::from(validity)));
+            let struct_array = StructArray::try_new(fields.clone(), child_arrays, Some(NullBuffer::from(validity)))
+                .map_err(|e| ApiError::permanent(format!("column `{name}`: invalid STRUCT array: {e}")))?;
             Ok((DataType::Struct(fields), Arc::new(struct_array)))
         }
         other => Err(ApiError {
@@ -702,6 +703,15 @@ mod tests {
         assert!(s.is_null(1), "row 1's whole struct must be null");
     }
 
+    #[test]
+    fn empty_struct_returns_a_conversion_error_instead_of_panicking() {
+        let columns = vec![struct_col("s", "STRUCT<>")];
+        for rows in [vec![], vec![vec![Some("{}".to_string())]], vec![vec![None]]] {
+            let err = json_array_to_record_batch(&rows, &columns).unwrap_err();
+            assert!(err.message.contains("invalid STRUCT array"));
+        }
+    }
+
     #[test]
     fn struct_with_a_nested_unsupported_composite_field_errors_instead_of_guessing() {
         let columns = vec![struct_col("s", "STRUCT<a: INT, nested: STRUCT<x: INT>>")];

{arrowbricks-3.1.3 → arrowbricks-3.2.0}/rust/arrowbricks_core/src/lib.rs

@@ -62,11 +62,14 @@ fn write_ipc_stream(py: Python<'_>, stream: Bound<'_, PyAny>, buf: Bound<'_, PyA
 #[pyfunction]
 #[pyo3(signature = (data))]
 fn read_ipc_stream(data: Bound<'_, PyBytes>) -> PyResult<PyTable> {
+    let py = data.py();
     // The immutable Python bytes own the memory for as long as any decoded
     // array needs it, including arrays exported through the C Data Interface.
     // PyBackedBytes provides that ownership without copying or custom unsafe code.
     let blob = bytes::Bytes::from_owner(pyo3::pybacked::PyBackedBytes::from(data));
-    let (batches, schema) = pipeline::decode_ipc_stream(&blob).map_err(|e| PyRuntimeError::new_err(e.message))?;
+    let (batches, schema) = py
+        .detach(|| pipeline::decode_ipc_stream(&blob))
+        .map_err(|e| PyRuntimeError::new_err(e.message))?;
     PyTable::try_new(batches, schema).map_err(|e| PyRuntimeError::new_err(e.to_string()))
 }
 
@@ -908,13 +911,10 @@ impl PyFetchallArrowStreamedIter {
     }
 }
 
-/// Not-yet-started vs. running state for `PyNdjsonStreamIter`. The
-/// submit/poll/spawn-workers step (`Pending` -> `Running`) happens on the
-/// iterator's first `__anext__` call, un-heartbeated -- matching
-/// `stream_query_json`'s pre-cutover behavior (its submit/poll wait was never
-/// heartbeat-wrapped, only its chunk loop was). The `total_timeout_s` budget
-/// starts counting from `Running`, not from construction, for the same
-/// reason.
+/// State for `PyNdjsonStreamIter`. The iterator's first `__anext__` starts
+/// the submit/poll/spawn-workers step (`Pending` -> `Submitting`), which is
+/// heartbeated and timed like the chunk pulls in `Running`; one
+/// `total_timeout_s` budget, counted from that first call, covers both.
 enum PyNdjsonStreamState {
     Pending {
         client: Arc<DbClient>,
@@ -925,6 +925,15 @@ enum PyNdjsonStreamState {
         total_timeout_s: Option<f64>,
         non_finite_as_string: bool,
     },
+    /// Submit/poll in flight under the same `total_timeout_s` budget (and
+    /// heartbeats) as the chunk pulls that follow -- a timeout here drops
+    /// the submit future, which cancels the statement server-side (see
+    /// `pipeline/stats.rs`'s `CancelInFlightOnDrop`).
+    Submitting {
+        client: Arc<DbClient>,
+        wait: HeartbeatWait<NdjsonStream>,
+        total_timeout_s: Option<f64>,
+    },
     Running {
         stream: Arc<AsyncMutex<NdjsonStream>>,
         heartbeat: HeartbeatStream<Vec<String>>,
@@ -967,22 +976,54 @@ impl PyNdjsonStreamIter {
                         else {
                             unreachable!()
                         };
-                        let client_for_cancel = client.clone();
-                        let stream = pipeline::execute_ndjson_stream(
+                        let submit_client = client.clone();
+                        let wait = HeartbeatWait::new(
+                            async move {
+                                pipeline::execute_ndjson_stream(
+                                    submit_client,
+                                    &statement,
+                                    catalog.as_deref(),
+                                    schema.as_deref(),
+                                    parameters,
+                                    non_finite_as_string,
+                                )
+                                .await
+                            },
+                            total_timeout_s,
+                        );
+                        *guard = PyNdjsonStreamState::Submitting {
                             client,
-                            &statement,
-                            catalog.as_deref(),
-                            schema.as_deref(),
-                            parameters,
-                            non_finite_as_string,
-                        )
-                        .await
-                        .map_err(api_error_to_pyerr)?;
-                        let heartbeat = HeartbeatStream::new(total_timeout_s).with_cancel(pipeline::cancel_hook(
-                            client_for_cancel,
-                            stream.cancel_handle.clone(),
-                            stream.stats.clone(),
-                        ));
+                            wait,
+                            total_timeout_s,
+                        };
+                    }
+                    PyNdjsonStreamState::Submitting { wait, .. } => {
+                        let stream = match wait.tick().await {
+                            Ok(Some(Tick::Heartbeat)) => {
+                                return Python::attach(|py| heartbeat_singleton(py).map(|h| h.into_any()));
+                            }
+                            Ok(Some(Tick::Ready(stream))) => stream,
+                            Ok(None) => unreachable!("Submitting is replaced as soon as its wait yields Ready"),
+                            Err(e) => {
+                                *guard = PyNdjsonStreamState::Done;
+                                return Err(api_error_to_pyerr(e));
+                            }
+                        };
+                        let PyNdjsonStreamState::Submitting {
+                            client,
+                            wait,
+                            total_timeout_s,
+                        } = std::mem::replace(&mut *guard, PyNdjsonStreamState::Done)
+                        else {
+                            unreachable!()
+                        };
+                        let heartbeat = HeartbeatStream::new(total_timeout_s)
+                            .with_deadline(wait.deadline())
+                            .with_cancel(pipeline::cancel_hook(
+                                client,
+                                stream.cancel_handle.clone(),
+                                stream.stats.clone(),
+                            ));
                         *guard = PyNdjsonStreamState::Running {
                             stream: Arc::new(AsyncMutex::new(stream)),
                             heartbeat,