arrowbricks 5.0.2__tar.gz → 5.0.3__tar.gz

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/PKG-INFO

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

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/pyproject.toml

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

@@ -177,7 +177,7 @@ dependencies = [
 
 [[package]]
 name = "arrowbricks_core"
-version = "5.0.2"
+version = "5.0.3"
 dependencies = [
  "arrow-array",
  "arrow-buffer",

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/Cargo.toml

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

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/client/download.rs

@@ -39,14 +39,21 @@ use super::model::QueryStatsAccumulator;
 /// 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;
     // 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[..]);
+    lz4_frame_decode_into(compressed, &mut out)?;
+    Ok(Bytes::from(out))
+}
+
+/// Appends every concatenated LZ4 frame in `src` to `out` -- the loop
+/// `decompress_lz4_frame` and the Thrift inline-blob path both need.
+pub(crate) fn lz4_frame_decode_into(src: &[u8], out: &mut Vec<u8>) -> Result<(), ApiError> {
+    use std::io::Read;
+    let mut decoder = lz4_flex::frame::FrameDecoder::new(src);
     // Terminate on the *reader* being exhausted, not on "output stopped
     // growing" -- found in code review that a frame which happens to decode
     // to zero bytes (a real, valid LZ4 Frame shape: header + immediate
@@ -61,10 +68,10 @@ pub(crate) fn decompress_lz4_frame(compressed: &Bytes) -> Result<Bytes, ApiError
     // between two real ones, see `decompress_lz4_frame_survives_a_zero_content_frame_in_the_middle`.
     while !decoder.get_ref().is_empty() {
         decoder
-            .read_to_end(&mut out)
+            .read_to_end(out)
             .map_err(|e| ApiError::permanent(format!("LZ4 frame decompress failed: {e}")))?;
     }
-    Ok(Bytes::from(out))
+    Ok(())
 }
 
 /// Reads the body into one buffer sized from `Content-Length`. Frames are

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/client/sea.rs

@@ -654,20 +654,13 @@ async fn join_first_error(handles: Vec<tokio::task::JoinHandle<Result<(), ApiErr
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::client::ApiErrorKind;
 
     fn ok_task() -> tokio::task::JoinHandle<Result<(), ApiError>> {
         tokio::spawn(async { Ok(()) })
     }
 
     fn err_task(msg: &'static str) -> tokio::task::JoinHandle<Result<(), ApiError>> {
-        tokio::spawn(async move {
-            Err(ApiError {
-                message: msg.to_string(),
-                transient: false,
-                kind: ApiErrorKind::Other,
-            })
-        })
+        tokio::spawn(async move { Err(ApiError::permanent(msg)) })
     }
 
     fn panicking_task() -> tokio::task::JoinHandle<Result<(), ApiError>> {

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/client.rs

@@ -44,6 +44,7 @@ mod sea;
 mod thrift_rpc;
 mod volume;
 
+pub(crate) use download::lz4_frame_decode_into;
 pub(crate) use error::join_error;
 pub use error::{ApiError, ApiErrorKind};
 pub use model::{

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/heartbeat.rs

@@ -12,7 +12,7 @@ use std::time::{Duration, Instant};
 
 use tokio::task::JoinHandle;
 
-use crate::client::{ApiError, ApiErrorKind, join_error};
+use crate::client::{ApiError, join_error};
 
 /// Matches `_streaming.py`'s `_HEARTBEAT_INTERVAL_S` -- well under typical
 /// PaaS idle-connection ceilings for a caller forwarding these as SSE
@@ -85,11 +85,7 @@ async fn check_deadline_or_wait<T>(
     let _ = handle.await;
     *handle_slot = None;
     let secs = total_timeout_s.unwrap_or(0.0);
-    Err(ApiError {
-        message: format!("Query exceeded {secs}s timeout"),
-        transient: false,
-        kind: ApiErrorKind::Other,
-    })
+    Err(ApiError::permanent(format!("Query exceeded {secs}s timeout")))
 }
 
 /// Shared `Drop` logic for `HeartbeatWait`/`HeartbeatStream`: best-effort, if
@@ -475,17 +471,8 @@ mod tests {
 
     #[tokio::test]
     async fn wrapped_future_error_propagates() {
-        let mut wait: HeartbeatWait<()> = HeartbeatWait::with_interval(
-            async {
-                Err(ApiError {
-                    message: "boom".into(),
-                    transient: false,
-                    kind: ApiErrorKind::Other,
-                })
-            },
-            None,
-            TEST_INTERVAL,
-        );
+        let mut wait: HeartbeatWait<()> =
+            HeartbeatWait::with_interval(async { Err(ApiError::permanent("boom")) }, None, TEST_INTERVAL);
         let err = wait.tick().await.unwrap_err();
         assert_eq!(err.message, "boom");
     }

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/json_convert.rs

@@ -28,14 +28,12 @@ use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit};
 use chrono::{DateTime, NaiveDate, NaiveDateTime};
 use serde_json::Value as JsonValue;
 
-use crate::client::{ApiError, ApiErrorKind, ColumnDescription};
+use crate::client::{ApiError, ColumnDescription};
 
 fn conv_err(column: &str, value: &str, type_name: &str, detail: impl std::fmt::Display) -> ApiError {
-    ApiError {
-        message: format!("column `{column}` (type {type_name}): could not parse {value:?}: {detail}"),
-        transient: false,
-        kind: ApiErrorKind::Other,
-    }
+    ApiError::permanent(format!(
+        "column `{column}` (type {type_name}): could not parse {value:?}: {detail}"
+    ))
 }
 
 pub fn json_array_to_record_batch(
@@ -46,11 +44,10 @@ pub fn json_array_to_record_batch(
     let mut arrays: Vec<ArrayRef> = Vec::with_capacity(columns.len());
 
     for (col_idx, col) in columns.iter().enumerate() {
-        let type_name = col.type_name.as_deref().ok_or_else(|| ApiError {
-            message: format!("column `{}` has no type_name in the manifest", col.name),
-            transient: false,
-            kind: ApiErrorKind::Other,
-        })?;
+        let type_name = col
+            .type_name
+            .as_deref()
+            .ok_or_else(|| ApiError::permanent(format!("column `{}` has no type_name in the manifest", col.name)))?;
         let values = rows.iter().map(|row| row.get(col_idx).cloned().flatten());
         let (data_type, array) = build_column(
             &col.name,
@@ -65,11 +62,8 @@ pub fn json_array_to_record_batch(
     }
 
     let schema = Arc::new(Schema::new(fields));
-    RecordBatch::try_new(schema, arrays).map_err(|e| ApiError {
-        message: format!("failed to assemble RecordBatch from JSON_ARRAY data: {e}"),
-        transient: false,
-        kind: ApiErrorKind::Other,
-    })
+    RecordBatch::try_new(schema, arrays)
+        .map_err(|e| ApiError::permanent(format!("failed to assemble RecordBatch from JSON_ARRAY data: {e}")))
 }
 
 #[allow(clippy::too_many_lines)]
@@ -153,11 +147,9 @@ fn build_column(
         }
         "DECIMAL" => {
             let (Some(precision), Some(scale)) = (precision, scale) else {
-                return Err(ApiError {
-                    message: format!("column `{name}`: DECIMAL type missing type_precision/type_scale in manifest"),
-                    transient: false,
-                    kind: ApiErrorKind::Other,
-                });
+                return Err(ApiError::permanent(format!(
+                    "column `{name}`: DECIMAL type missing type_precision/type_scale in manifest"
+                )));
             };
             let mut out = Vec::new();
             for v in values {
@@ -168,10 +160,10 @@ fn build_column(
             }
             let array = Decimal128Array::from(out)
                 .with_precision_and_scale(precision, scale)
-                .map_err(|e| ApiError {
-                    message: format!("column `{name}`: invalid DECIMAL(precision={precision}, scale={scale}): {e}"),
-                    transient: false,
-                    kind: ApiErrorKind::Other,
+                .map_err(|e| {
+                    ApiError::permanent(format!(
+                        "column `{name}`: invalid DECIMAL(precision={precision}, scale={scale}): {e}"
+                    ))
                 })?;
             Ok((DataType::Decimal128(precision, scale), Arc::new(array)))
         }
@@ -259,15 +251,13 @@ fn build_column(
             Ok((DataType::Binary, Arc::new(BinaryArray::from(refs))))
         }
         "STRUCT" => {
-            let type_text = type_text.ok_or_else(|| ApiError {
-                message: format!("column `{name}`: STRUCT type missing type_text in manifest"),
-                transient: false,
-                kind: ApiErrorKind::Other,
+            let type_text = type_text.ok_or_else(|| {
+                ApiError::permanent(format!("column `{name}`: STRUCT type missing type_text in manifest"))
             })?;
-            let field_defs = parse_struct_fields(type_text).map_err(|e| ApiError {
-                message: format!("column `{name}`: could not parse STRUCT type_text {type_text:?}: {e}"),
-                transient: false,
-                kind: ApiErrorKind::Other,
+            let field_defs = parse_struct_fields(type_text).map_err(|e| {
+                ApiError::permanent(format!(
+                    "column `{name}`: could not parse STRUCT type_text {type_text:?}: {e}"
+                ))
             })?;
 
             // One parsed JSON object per row (`None` = the whole struct is
@@ -338,11 +328,9 @@ fn build_column(
                 .map_err(|e| ApiError::permanent(format!("column `{name}`: invalid STRUCT array: {e}")))?;
             Ok((DataType::Struct(fields), Arc::new(struct_array)))
         }
-        other => Err(ApiError {
-            message: format!("column `{name}`: type {other} isn't supported by the INLINE/JSON_ARRAY fast path"),
-            transient: false,
-            kind: ApiErrorKind::Other,
-        }),
+        other => Err(ApiError::permanent(format!(
+            "column `{name}`: type {other} isn't supported by the INLINE/JSON_ARRAY fast path"
+        ))),
     }
 }
 

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/pipeline/ndjson.rs

@@ -16,7 +16,7 @@ use arrow_array::types::{Float32Type, Float64Type};
 use arrow_schema::DataType;
 use serde_json::Value;
 
-use crate::client::{ApiError, ApiErrorKind, CancelHandle, DbClient, Protocol, QueryStatsAccumulator, join_error};
+use crate::client::{ApiError, CancelHandle, DbClient, Protocol, QueryStatsAccumulator, join_error};
 
 use super::reorder::{ReorderBuffer, decode_chunk_item};
 use super::sea::submit_sea_and_report;
@@ -83,16 +83,12 @@ fn encode_ndjson_lines(batches: &[RecordBatch], non_finite_as_string: bool) -> R
         let builder = arrow_json::WriterBuilder::new().with_explicit_nulls(true);
         let mut writer = builder.build::<_, arrow_json::writer::LineDelimited>(&mut output);
         let refs: Vec<&RecordBatch> = batches.iter().collect();
-        writer.write_batches(&refs).map_err(|e| ApiError {
-            message: format!("NDJSON encode error: {e}"),
-            transient: false,
-            kind: ApiErrorKind::Other,
-        })?;
-        writer.finish().map_err(|e| ApiError {
-            message: format!("NDJSON encode error: {e}"),
-            transient: false,
-            kind: ApiErrorKind::Other,
-        })?;
+        writer
+            .write_batches(&refs)
+            .map_err(|e| ApiError::permanent(format!("NDJSON encode error: {e}")))?;
+        writer
+            .finish()
+            .map_err(|e| ApiError::permanent(format!("NDJSON encode error: {e}")))?;
     }
     if !output.pending.is_empty() {
         return Err(ApiError::permanent("NDJSON encode produced an unterminated row"));
@@ -407,23 +403,15 @@ mod tests {
             let builder = arrow_json::WriterBuilder::new().with_explicit_nulls(true);
             let mut writer = builder.build::<_, arrow_json::writer::LineDelimited>(&mut buf);
             let refs: Vec<&RecordBatch> = batches.iter().collect();
-            writer.write_batches(&refs).map_err(|e| ApiError {
-                message: format!("NDJSON encode error: {e}"),
-                transient: false,
-                kind: ApiErrorKind::Other,
-            })?;
-            writer.finish().map_err(|e| ApiError {
-                message: format!("NDJSON encode error: {e}"),
-                transient: false,
-                kind: ApiErrorKind::Other,
-            })?;
+            writer
+                .write_batches(&refs)
+                .map_err(|e| ApiError::permanent(format!("NDJSON encode error: {e}")))?;
+            writer
+                .finish()
+                .map_err(|e| ApiError::permanent(format!("NDJSON encode error: {e}")))?;
         }
         let mut lines: Vec<String> = String::from_utf8(buf)
-            .map_err(|e| ApiError {
-                message: format!("NDJSON encode produced invalid UTF-8: {e}"),
-                transient: false,
-                kind: ApiErrorKind::Other,
-            })?
+            .map_err(|e| ApiError::permanent(format!("NDJSON encode produced invalid UTF-8: {e}")))?
             .lines()
             .map(|line| line.to_string())
             .collect();

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/pipeline/reorder.rs

@@ -15,7 +15,7 @@ use arrow_schema::SchemaRef;
 use bytes::Bytes;
 use tokio::sync::mpsc;
 
-use crate::client::{ApiError, ApiErrorKind, ChunkItem};
+use crate::client::{ApiError, ChunkItem};
 
 /// Same dict-of-lists-keyed-by-index shape as `_ResultSet._pending`: a
 /// `chunk_index` can carry more than one blob (multiple `external_links` per
@@ -123,11 +123,9 @@ fn decode_chunk(blob: &Bytes) -> Result<Vec<RecordBatch>, ApiError> {
 /// schema-only streams that contain no record batches.
 pub(crate) fn decode_ipc_stream(blob: &Bytes) -> Result<(Vec<RecordBatch>, SchemaRef), ApiError> {
     if blob.is_empty() {
-        return Err(ApiError {
-            message: "empty Arrow IPC chunk: expected at least a schema message".to_string(),
-            transient: false,
-            kind: ApiErrorKind::Other,
-        });
+        return Err(ApiError::permanent(
+            "empty Arrow IPC chunk: expected at least a schema message",
+        ));
     }
     let mut buffer = ArrowBuffer::from(blob.clone());
     let mut decoder = StreamDecoder::new();
@@ -137,19 +135,13 @@ pub(crate) fn decode_ipc_stream(blob: &Bytes) -> Result<(Vec<RecordBatch>, Schem
             Ok(Some(batch)) => batches.push(batch),
             Ok(None) => {}
             Err(e) => {
-                return Err(ApiError {
-                    message: format!("Arrow IPC decode error: {e}"),
-                    transient: false,
-                    kind: ApiErrorKind::Other,
-                });
+                return Err(ApiError::permanent(format!("Arrow IPC decode error: {e}")));
             }
         }
     }
-    decoder.finish().map_err(|e| ApiError {
-        message: format!("bad Arrow IPC stream: {e}"),
-        transient: false,
-        kind: ApiErrorKind::Other,
-    })?;
+    decoder
+        .finish()
+        .map_err(|e| ApiError::permanent(format!("bad Arrow IPC stream: {e}")))?;
     let schema = decoder
         .schema()
         .ok_or_else(|| ApiError::permanent("Arrow IPC stream has no schema"))?;
@@ -299,15 +291,7 @@ mod tests {
 
     #[tokio::test]
     async fn error_surfaces_after_already_yielded_items() {
-        let sent = vec![
-            Ok(item(0)),
-            Ok(item(1)),
-            Err(ApiError {
-                message: "boom".into(),
-                transient: false,
-                kind: ApiErrorKind::Other,
-            }),
-        ];
+        let sent = vec![Ok(item(0)), Ok(item(1)), Err(ApiError::permanent("boom"))];
         let (tx, rx) = mpsc::channel(sent.len());
         for r in sent {
             tx.send(r).await.unwrap();

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/pipeline/sea.rs

@@ -17,8 +17,8 @@ use serde_json::Value;
 use tokio::sync::mpsc;
 
 use crate::client::{
-    ApiError, ApiErrorKind, CancelHandle, ColumnDescription, DbClient, InlineOrExternal, QueryStatsAccumulator,
-    QueryStatsData, StatementSubmitResult, join_error,
+    ApiError, CancelHandle, ColumnDescription, DbClient, InlineOrExternal, QueryStatsAccumulator, QueryStatsData,
+    StatementSubmitResult, join_error,
 };
 
 use super::reorder::{ReorderBuffer, decode_chunk_item};
@@ -100,13 +100,10 @@ impl ResultStream {
     /// single capped batch isn't enough).
     async fn fetch_at_least(&mut self, want_rows: usize) -> Result<(), ApiError> {
         if self.poisoned {
-            return Err(ApiError {
-                message: "this result was left incomplete by a previous cancelled, timed-out, or failed fetch -- \
-                          re-run the query instead of continuing to use this cursor/result"
-                    .to_string(),
-                transient: false,
-                kind: ApiErrorKind::Other,
-            });
+            return Err(ApiError::permanent(
+                "this result was left incomplete by a previous cancelled, timed-out, or failed fetch -- \
+                          re-run the query instead of continuing to use this cursor/result",
+            ));
         }
         let mut guard = PoisonOnDrop::new(&mut self.poisoned, &mut self.reporter, self.stats.as_ref());
         while self.pending_rows < want_rows && !self.exhausted {

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/rust/arrowbricks_core/src/pipeline/thrift_exec.rs

@@ -21,6 +21,7 @@ use tokio::sync::mpsc;
 
 use crate::client::{
     ApiError, CancelHandle, ChunkItem, ColumnDescription, DbClient, QueryStatsAccumulator, join_error,
+    lz4_frame_decode_into,
 };
 use crate::thrift;
 
@@ -854,16 +855,7 @@ fn build_inline_blob(
     for b in batches {
         row_count += b.row_count;
         if lz4_compressed {
-            // Decompress directly into the output buffer using FrameDecoder,
-            // avoiding a separate intermediate allocation and copy. The
-            // FrameDecoder's read_to_end appends to the existing buffer.
-            use std::io::Read;
-            let mut decoder = lz4_flex::frame::FrameDecoder::new(&b.batch[..]);
-            while !decoder.get_ref().is_empty() {
-                decoder
-                    .read_to_end(&mut out)
-                    .map_err(|e| ApiError::permanent(format!("LZ4 frame decompress failed: {e}")))?;
-            }
+            lz4_frame_decode_into(&b.batch, &mut out)?;
         } else {
             out.extend_from_slice(&b.batch);
         }

{arrowbricks-5.0.2 → arrowbricks-5.0.3}/src/arrowbricks/cursor.py

@@ -48,6 +48,11 @@ Description = tuple[str, str | None, None, None, None, None, None]
 _ROW_BATCH = 1000
 
 
+def _join_rows(buffered: list[Row], rows: list[Row]) -> list[Row]:
+    # `buffered + rows` copies every row; usually there's nothing buffered.
+    return buffered + rows if buffered else rows
+
+
 def _table_to_rows(table: Any) -> list[Row]:
     if table.num_rows == 0:
         return []
@@ -223,8 +228,7 @@ class Cursor:
         if self._row_pos >= len(self._row_buffer):
             result = self._require_result()
             table = await result.fetchmany_arrow(_ROW_BATCH)
-            if self._schema is None:
-                self._schema = await result.schema()
+            await self._ensure_schema(result)
             self._row_buffer = _table_to_rows(table)
             self._row_pos = 0
         if self._row_pos >= len(self._row_buffer):
@@ -237,26 +241,28 @@ class Cursor:
         buffered = self._take_buffered(size)
         if len(buffered) == size:
             return buffered
-        return buffered + _table_to_rows(await self.fetchmany_arrow(size - len(buffered)))
+        return _join_rows(buffered, _table_to_rows(await self.fetchmany_arrow(size - len(buffered))))
 
     async def fetchall(self) -> list[Row]:
         buffered = self._take_buffered(len(self._row_buffer))
-        return buffered + _table_to_rows(await self.fetchall_arrow())
+        return _join_rows(buffered, _table_to_rows(await self.fetchall_arrow()))
+
+    async def _ensure_schema(self, result: _core.ResultSet) -> None:
+        if self._schema is None:
+            self._schema = await result.schema()
 
     async def fetchmany_arrow(self, size: int) -> core.Table:
         self._require_empty_row_buffer("fetchmany_arrow")
         result = self._require_result()
         table = await result.fetchmany_arrow(size)
-        if self._schema is None:
-            self._schema = await result.schema()
+        await self._ensure_schema(result)
         return table
 
     async def fetchall_arrow(self) -> core.Table:
         self._require_empty_row_buffer("fetchall_arrow")
         result = self._require_result()
         table = await result.fetchall_arrow()
-        if self._schema is None:
-            self._schema = await result.schema()
+        await self._ensure_schema(result)
         return table
 
     def fetchall_streamed(self, *, total_timeout_s: float | None = None) -> AsyncIterator[Any]:
@@ -317,8 +323,7 @@ class Cursor:
                     if item is _core.HEARTBEAT:
                         yield HEARTBEAT
                     else:
-                        if self._schema is None:
-                            self._schema = await result.schema()
+                        await self._ensure_schema(result)
                         yield item
 
         return _gen()