arrowbricks 3.1.4__tar.gz → 4.0.0__tar.gz

{arrowbricks-3.1.4 → arrowbricks-4.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arrowbricks
-Version: 3.1.4
+Version: 4.0.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
 
@@ -212,6 +212,15 @@ except ArrowbricksError:
 
 ## Arrow vs. row-tuple fetches
 
+From 4.0, Arrow fetches and `ReplayableArrowChunk.to_table()` return
+`arrowbricks._core.Table`. It exposes `num_rows`, `num_columns`,
+`column_names`, `len(table)` and the Arrow capsule interfaces. For methods
+such as `.column()`, `.schema` or `.to_batches()`, first convert with
+`arro3.core.Table.from_arrow(table)`, `pyarrow.table(table)` or
+`polars.from_arrow(table)`. DuckDB can consume it directly.
+`write_ipc_stream` requires `__arrow_c_stream__`; wrap an array-only producer
+in a table before passing it in.
+
 Everything above works with zero dependencies installed *except* row-tuple fetches. `fetchall_arrow`/`fetchmany_arrow` return an Arrow table straight from the Rust core -- the faster path if your code can consume Arrow directly (DuckDB, pyarrow, polars, a Parquet writer, ...):
 
 ```python

{arrowbricks-3.1.4 → arrowbricks-4.0.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
 
@@ -200,6 +200,15 @@ except ArrowbricksError:
 
 ## Arrow vs. row-tuple fetches
 
+From 4.0, Arrow fetches and `ReplayableArrowChunk.to_table()` return
+`arrowbricks._core.Table`. It exposes `num_rows`, `num_columns`,
+`column_names`, `len(table)` and the Arrow capsule interfaces. For methods
+such as `.column()`, `.schema` or `.to_batches()`, first convert with
+`arro3.core.Table.from_arrow(table)`, `pyarrow.table(table)` or
+`polars.from_arrow(table)`. DuckDB can consume it directly.
+`write_ipc_stream` requires `__arrow_c_stream__`; wrap an array-only producer
+in a table before passing it in.
+
 Everything above works with zero dependencies installed *except* row-tuple fetches. `fetchall_arrow`/`fetchmany_arrow` return an Arrow table straight from the Rust core -- the faster path if your code can consume Arrow directly (DuckDB, pyarrow, polars, a Parquet writer, ...):
 
 ```python

{arrowbricks-3.1.4 → arrowbricks-4.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "3.1.4"
+version = "4.0.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"
@@ -15,12 +15,8 @@ requires-python = ">=3.11"
 dependencies = []
 
 [project.optional-dependencies]
-# pyo3-arrow's own Table.column()/schema methods are designed to hand back
-# the caller's *real* arro3.core objects for further arro3-ecosystem
-# interop (not a value pyo3-arrow can materialize itself), so
-# fetchone/fetchmany/fetchall need this installed. `pip install
-# arrowbricks[arro3]` -- or handle rows as Arrow via fetchall_arrow/
-# fetchmany_arrow instead, which don't need it at all.
+# Row materialization converts our capsule table through arro3-core.
+# Arrow fetches and streaming do not need this optional dependency.
 arro3 = ["arro3-core>=0.8"]
 
 [project.urls]

{arrowbricks-3.1.4 → arrowbricks-4.0.0}/rust/arrowbricks_core/Cargo.lock

@@ -81,7 +81,6 @@ dependencies = [
  "atoi",
  "base64 0.23.1",
  "chrono",
- "comfy-table",
  "half",
  "lexical-core",
  "num-traits",
@@ -178,7 +177,7 @@ dependencies = [
 
 [[package]]
 name = "arrowbricks_core"
-version = "3.1.4"
+version = "4.0.0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -193,7 +192,6 @@ dependencies = [
  "memchr",
  "proptest",
  "pyo3",
- "pyo3-arrow",
  "pyo3-async-runtimes",
  "pythonize",
  "reqwest",
@@ -303,9 +301,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327"
 dependencies = [
  "iana-time-zone",
- "js-sys",
  "num-traits",
- "wasm-bindgen",
  "windows-link",
 ]
 
@@ -329,16 +325,6 @@ dependencies = [
  "memchr",
 ]
 
-[[package]]
-name = "comfy-table"
-version = "7.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "958c5d6ecf1f214b4c2bbbbf6ab9523a864bd136dcf71a7e8904799acfe1ad47"
-dependencies = [
- "unicode-segmentation",
- "unicode-width",
-]
-
 [[package]]
 name = "const-random"
 version = "0.1.18"
@@ -889,7 +875,7 @@ dependencies = [
  "jni-sys",
  "log",
  "simd_cesu8",
- "thiserror 2.0.20",
+ "thiserror",
  "walkdir",
  "windows-link",
 ]
@@ -1039,16 +1025,6 @@ dependencies = [
  "twox-hash",
 ]
 
-[[package]]
-name = "matrixmultiply"
-version = "0.3.11"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7"
-dependencies = [
- "autocfg",
- "rawpointer",
-]
-
 [[package]]
 name = "memchr"
 version = "2.8.3"
@@ -1066,21 +1042,6 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
-[[package]]
-name = "ndarray"
-version = "0.17.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
-dependencies = [
- "matrixmultiply",
- "num-complex",
- "num-integer",
- "num-traits",
- "portable-atomic",
- "portable-atomic-util",
- "rawpointer",
-]
-
 [[package]]
 name = "num-bigint"
 version = "0.5.1"
@@ -1129,23 +1090,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "numpy"
-version = "0.29.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4"
-dependencies = [
- "half",
- "libc",
- "ndarray",
- "num-complex",
- "num-integer",
- "num-traits",
- "pyo3",
- "pyo3-build-config",
- "rustc-hash",
-]
-
 [[package]]
 name = "once_cell"
 version = "1.21.4"
@@ -1194,15 +1138,6 @@ version = "1.14.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3"
 
-[[package]]
-name = "portable-atomic-util"
-version = "0.2.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618"
-dependencies = [
- "portable-atomic",
-]
-
 [[package]]
 name = "potential_utf"
 version = "0.1.5"
@@ -1255,9 +1190,6 @@ version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
 dependencies = [
- "chrono",
- "chrono-tz",
- "indexmap",
  "libc",
  "once_cell",
  "portable-atomic",
@@ -1266,27 +1198,6 @@ dependencies = [
  "pyo3-macros",
 ]
 
-[[package]]
-name = "pyo3-arrow"
-version = "0.19.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3d5ddf226a2dbf7607570d0657c2bf6fbe299208368b2914f3dd7e7ba0b57688"
-dependencies = [
- "arrow-array",
- "arrow-buffer",
- "arrow-cast",
- "arrow-data",
- "arrow-schema",
- "arrow-select",
- "chrono",
- "chrono-tz",
- "half",
- "indexmap",
- "numpy",
- "pyo3",
- "thiserror 1.0.69",
-]
-
 [[package]]
 name = "pyo3-async-runtimes"
 version = "0.29.0"
@@ -1413,12 +1324,6 @@ dependencies = [
  "rand_core",
 ]
 
-[[package]]
-name = "rawpointer"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
-
 [[package]]
 name = "regex"
 version = "1.13.1"
@@ -1499,12 +1404,6 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
-[[package]]
-name = "rustc-hash"
-version = "2.1.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d"
-
 [[package]]
 name = "rustc_version"
 version = "0.4.1"
@@ -1850,33 +1749,13 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
-[[package]]
-name = "thiserror"
-version = "1.0.69"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
-dependencies = [
- "thiserror-impl 1.0.69",
-]
-
 [[package]]
 name = "thiserror"
 version = "2.0.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f"
 dependencies = [
- "thiserror-impl 2.0.20",
-]
-
-[[package]]
-name = "thiserror-impl"
-version = "1.0.69"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.119",
+ "thiserror-impl",
 ]
 
 [[package]]
@@ -2047,18 +1926,6 @@ version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
 
-[[package]]
-name = "unicode-segmentation"
-version = "1.13.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
-
-[[package]]
-name = "unicode-width"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
-
 [[package]]
 name = "untrusted"
 version = "0.9.0"

{arrowbricks-3.1.4 → arrowbricks-4.0.0}/rust/arrowbricks_core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "arrowbricks_core"
-version = "3.1.4"
+version = "4.0.0"
 edition = "2024"
 readme = "README.md"
 
@@ -11,7 +11,10 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 # 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"
+# `ffi` backs src/arrow_ffi.rs, this crate's own Arrow PyCapsule bridge.
+# `chrono-tz` lets arrow-json format Databricks' named-zone timestamps
+# ("Etc/UTC"); pyo3-arrow used to enable it implicitly.
+arrow-array = { version = "59.1.0", features = ["ffi", "chrono-tz"] }
 arrow-buffer = "59.1.0"
 arrow-ipc = "59.1.0"
 arrow-schema = "59.1.0"
@@ -20,15 +23,14 @@ 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
+# arrow-json) at this exact version -- declared directly here too so
 # json_convert.rs can decode BINARY columns without hand-rolling base64. No
 # new code ships that wasn't already linked in (feature unification, same
 # reasoning as chrono/rustls below).
 base64 = "0.23"
-# Already an unconditional transitive dependency of pyo3-arrow (see its own
-# comment below) -- declared directly here too so json_convert.rs can parse
-# DATE/TIMESTAMP/TIMESTAMP_NTZ strings without hand-rolling calendar
-# arithmetic. No new code ships that wasn't already linked in.
+# Already an unconditional transitive dependency of arrow-array -- declared
+# directly here too so json_convert.rs can parse DATE/TIMESTAMP/TIMESTAMP_NTZ
+# strings without hand-rolling calendar arithmetic. No new code ships that wasn't already linked in.
 chrono = { version = "0.4", default-features = false, features = ["std"] }
 # Decompresses cloud-fetch chunk bytes when the server honors our
 # `result_compression: LZ4_FRAME` request (see client.rs's execute_statement)
@@ -36,12 +38,6 @@ chrono = { version = "0.4", default-features = false, features = ["std"] }
 # Pure Rust, no C toolchain dependency -- unlike reqwest's TLS stack below.
 lz4_flex = { version = "0.14", default-features = false, features = ["frame", "std"] }
 pyo3 = { version = "0.29.1", features = ["abi3-py311"] }
-# pyo3-arrow's only optional feature is buffer_protocol (on by default),
-# which we don't use -- numpy/chrono/chrono-tz/arrow-cast(prettyprint) below
-# it are unconditional dependencies of the crate itself, not features, so
-# they ship regardless; not worth dropping pyo3-arrow over that (reimplementing
-# the Arrow C Data Interface ourselves is a much bigger risk than the MBs saved).
-pyo3-arrow = { version = "0.19.0", default-features = false }
 pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
 pythonize = "0.29.0"
 # reqwest's own `rustls` feature is hardcoded to `__rustls-aws-lc-rs`, no way
@@ -135,9 +131,3 @@ 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.4 → arrowbricks-4.0.0}/rust/arrowbricks_core/README.md

@@ -63,12 +63,11 @@ asyncio.run(main())
 ```
 
 `table` implements the [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html)
-(`__arrow_c_stream__`/`__arrow_c_array__`) -- any consumer that speaks that
-protocol (DuckDB, pyarrow, arro3) can import it with zero copies, no extra
-dependency needed. Note: `table.column(i)`/`table.schema` (materializing
-native Python values) are `pyo3-arrow` methods designed to hand back the
-caller's own *real* `arro3.core` objects -- they need `arro3-core` installed,
-unlike the table object itself.
+(`__arrow_c_stream__`/`__arrow_c_schema__`) -- any consumer that speaks that
+protocol (DuckDB, pyarrow, polars, arro3) can import it with zero copies, no
+extra dependency needed. The table itself only exposes `num_rows`,
+`num_columns`, `column_names` and `len()`; to read values, import it into an
+Arrow library first (e.g. `arro3.core.Table.from_arrow(table)`).
 
 ## API
 

arrowbricks-4.0.0/rust/arrowbricks_core/src/arrow_ffi.rs

@@ -0,0 +1,106 @@
+//! Arrow PyCapsule interface, both directions, on arrow-rs's own
+//! `ffi_stream` -- replaces `pyo3-arrow`, whose `Table.__arrow_c_stream__`
+//! honors `requested_schema` via `arrow_cast::cast` and so kept the whole cast
+//! kernel (~1.8 MiB, a third of `.text`) linked in for a feature no caller of
+//! this crate used.
+
+use std::ffi::CStr;
+
+use arrow_array::ffi::FFI_ArrowSchema;
+use arrow_array::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
+use arrow_array::{RecordBatch, RecordBatchIterator};
+use arrow_schema::SchemaRef;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3::types::{PyCapsule, PyCapsuleMethods};
+
+const STREAM_CAPSULE: &CStr = c"arrow_array_stream";
+const SCHEMA_CAPSULE: &CStr = c"arrow_schema";
+
+/// An in-memory Arrow table. Consume it through `__arrow_c_stream__`
+/// (arro3, pyarrow, polars, DuckDB, ...); the attributes are just enough to
+/// inspect a result without importing another Arrow library.
+#[pyclass(name = "Table", module = "arrowbricks._core", frozen)]
+pub struct PyTable {
+    batches: Vec<RecordBatch>,
+    schema: SchemaRef,
+}
+
+impl PyTable {
+    pub fn try_new(batches: Vec<RecordBatch>, schema: SchemaRef) -> PyResult<Self> {
+        // A stream advertises one schema for every batch. Exporting different
+        // buffer layouts under it can make consumers misinterpret memory.
+        // Like the previous bridge, ignore metadata and nullability differences.
+        if batches.iter().any(|batch| {
+            let fields = batch.schema_ref().fields();
+            fields.len() != schema.fields().len()
+                || fields.iter().zip(schema.fields()).any(|(actual, expected)| {
+                    actual.name() != expected.name() || !actual.data_type().equals_datatype(expected.data_type())
+                })
+        }) {
+            return Err(PyRuntimeError::new_err("All batches must have same schema"));
+        }
+        Ok(Self { batches, schema })
+    }
+}
+
+#[pymethods]
+impl PyTable {
+    #[getter]
+    fn num_rows(&self) -> usize {
+        self.batches.iter().map(RecordBatch::num_rows).sum()
+    }
+
+    #[getter]
+    fn num_columns(&self) -> usize {
+        self.schema.fields().len()
+    }
+
+    #[getter]
+    fn column_names(&self) -> Vec<String> {
+        self.schema.fields().iter().map(|f| f.name().clone()).collect()
+    }
+
+    fn __len__(&self) -> usize {
+        self.num_rows()
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "arrowbricks.Table(num_rows={}, num_columns={})",
+            self.num_rows(),
+            self.num_columns()
+        )
+    }
+
+    /// `requested_schema` is ignored, as the PyCapsule interface allows: the
+    /// consumer checks the schema it actually gets.
+    #[pyo3(signature = (requested_schema=None))]
+    fn __arrow_c_stream__<'py>(
+        &self,
+        py: Python<'py>,
+        requested_schema: Option<Bound<'py, PyAny>>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let _ = requested_schema;
+        let reader = RecordBatchIterator::new(self.batches.clone().into_iter().map(Ok), self.schema.clone());
+        PyCapsule::new_with_value(py, FFI_ArrowArrayStream::new(Box::new(reader)), STREAM_CAPSULE)
+    }
+
+    fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
+        let schema =
+            FFI_ArrowSchema::try_from(self.schema.as_ref()).map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        PyCapsule::new_with_value(py, schema, SCHEMA_CAPSULE)
+    }
+}
+
+/// Imports any object implementing `__arrow_c_stream__`.
+pub fn import_stream(obj: &Bound<'_, PyAny>) -> PyResult<ArrowArrayStreamReader> {
+    let capsule = obj.call_method0("__arrow_c_stream__")?.cast_into::<PyCapsule>()?;
+    let ptr = capsule
+        .pointer_checked(Some(STREAM_CAPSULE))?
+        .cast::<FFI_ArrowArrayStream>();
+    // SAFETY: the capsule name check guarantees an `ArrowArrayStream`;
+    // `from_raw` moves it out and marks the capsule's copy released, so the
+    // capsule destructor does not release it a second time.
+    unsafe { ArrowArrayStreamReader::from_raw(ptr.as_ptr()) }.map_err(|e| PyRuntimeError::new_err(e.to_string()))
+}

{arrowbricks-3.1.4 → arrowbricks-4.0.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.4 → arrowbricks-4.0.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,17 @@ 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());
+        if result.is_ok() {
+            self.note_warehouse_running();
         }
+        checkin.finish(result.is_ok());
         result
     }
 
@@ -390,6 +398,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 +411,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 +611,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 {