arrowbricks 1.3.0__tar.gz → 1.3.2__tar.gz

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "arrowbricks"
-version = "1.3.0"
+version = "1.3.2"
 description = "Runs SQL against a Databricks SQL warehouse via the Statement Execution API and hands you the result as Arrow -- a DB-API-ish Cursor (fetchone/fetchmany/fetchall/fetchall_arrow) or NDJSON streaming. Rust/PyO3 core throughout -- zero required runtime dependencies."
 readme = "README.md"
 license = "MIT"

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

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

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

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

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

@@ -2,12 +2,12 @@
 //! decode of the reordered chunk stream via arrow-rs.
 
 use std::collections::{HashMap, VecDeque};
-use std::io::Cursor as IoCursor;
 use std::sync::Arc;
 
 use arrow::array::{Array, AsArray};
+use arrow::buffer::Buffer as ArrowBuffer;
 use arrow::datatypes::{DataType, Float32Type, Float64Type, SchemaRef};
-use arrow::ipc::reader::StreamReader;
+use arrow::ipc::reader::StreamDecoder;
 use arrow::record_batch::RecordBatch;
 use bytes::Bytes;
 use serde_json::Value;
@@ -97,15 +97,66 @@ impl ExecuteResult {
     }
 }
 
+/// Decodes a chunk's raw Arrow-IPC stream bytes into batches. Uses
+/// `StreamDecoder`'s push-based interface fed by an `arrow::buffer::Buffer`
+/// built directly from `blob` (`Buffer::from(bytes::Bytes)`, confirmed
+/// zero-copy in arrow-buffer's own source -- `bytes.rs`'s
+/// `impl From<bytes::Bytes> for Bytes` stores the original `bytes::Bytes` via
+/// `Deallocation::Custom`, no memcpy) instead of the higher-level
+/// `StreamReader` (reads via `std::io::Read` into freshly allocated buffers,
+/// copying every column's data out of `blob` on every decode -- what this
+/// used before). For properly aligned IPC data (the normal case -- Databricks
+/// writes it, not this crate), decoded batches now slice directly into the
+/// same allocation `blob` already held since the network fetch, cutting out
+/// a second full copy of every chunk's bytes; `require_alignment` stays at
+/// its default `false`, so a misaligned *fixed-width* buffer still falls back
+/// to a copy automatically rather than erroring (arrow-ipc's own documented
+/// behavior) -- variable-width values, null bitmaps, and nested/dictionary
+/// children all stay zero-copy regardless. One exception this doesn't cover:
+/// if a `RecordBatch` message declared IPC *buffer*-level compression (a
+/// different, unrelated feature from this crate's own cloud-fetch
+/// `result_compression` unwrap in `client.rs`, which already ran before this
+/// function ever sees the bytes), `arrow-ipc`'s own reader always
+/// decompresses into fresh buffers there -- not something Databricks has
+/// been observed to use in this format, but not something this crate
+/// controls either.
+///
+/// A zero-length `blob` is rejected explicitly rather than handed to
+/// `StreamDecoder`: found in code review that an empty buffer makes the
+/// `while` loop below a no-op and `decoder.finish()` sees a still-pristine
+/// decoder state, which its own `Ok(())` arm treats as a *clean, empty*
+/// stream -- silently returning zero batches with no error at all, the same
+/// silent-truncation failure mode as the real multi-frame LZ4 bug this crate
+/// already shipped once (see the `result_compression` invariant above). The
+/// old `StreamReader`-based version failed loudly on this input instead
+/// ("Expected schema message, found empty stream"); this restores that.
 fn decode_chunk(blob: &Bytes) -> Result<Vec<RecordBatch>, ApiError> {
-    let reader = StreamReader::try_new(IoCursor::new(&blob[..]), None).map_err(|e| ApiError {
+    if blob.is_empty() {
+        return Err(ApiError {
+            message: "empty Arrow IPC chunk: expected at least a schema message".to_string(),
+            transient: false,
+        });
+    }
+    let mut buffer = ArrowBuffer::from(blob.clone());
+    let mut decoder = StreamDecoder::new();
+    let mut batches = Vec::new();
+    while !buffer.is_empty() {
+        match decoder.decode(&mut buffer) {
+            Ok(Some(batch)) => batches.push(batch),
+            Ok(None) => {}
+            Err(e) => {
+                return Err(ApiError {
+                    message: format!("Arrow IPC decode error: {e}"),
+                    transient: false,
+                });
+            }
+        }
+    }
+    decoder.finish().map_err(|e| ApiError {
         message: format!("bad Arrow IPC stream: {e}"),
         transient: false,
     })?;
-    reader.collect::<Result<Vec<_>, _>>().map_err(|e| ApiError {
-        message: format!("Arrow IPC decode error: {e}"),
-        transient: false,
-    })
+    Ok(batches)
 }
 
 /// Caps how many chunks `ResultStream::fetch_at_least` will pull/decode
@@ -797,4 +848,187 @@ mod tests {
             "a real NULL must stay null, never become a string"
         );
     }
+
+    /// Regression test for the switch from `StreamReader` to the push-based
+    /// `StreamDecoder` (see `decode_chunk`'s doc comment): a single chunk can
+    /// contain more than one Arrow-IPC `RecordBatch` message back to back in
+    /// the same stream, and `StreamDecoder::decode` only ever returns one
+    /// batch per call -- `decode_chunk` must keep calling it until the whole
+    /// buffer is drained, not stop after the first. `StreamReader`'s own
+    /// `Iterator` impl made this automatic; the lower-level API doesn't, so
+    /// this is exactly the kind of thing that regresses silently (a single-
+    /// batch-per-chunk bug would still pass every other test in this suite,
+    /// since none of them writes more than one batch per chunk).
+    #[test]
+    fn decode_chunk_reads_every_record_batch_in_a_multi_batch_stream() {
+        use arrow::ipc::writer::StreamWriter;
+
+        let batch_a = make_batch(vec![1, 2], vec![1.0, 2.0]);
+        let batch_b = make_batch(vec![3, 4, 5], vec![3.0, 4.0, 5.0]);
+
+        let mut buf = Vec::new();
+        {
+            let mut writer = StreamWriter::try_new(&mut buf, &batch_a.schema()).unwrap();
+            writer.write(&batch_a).unwrap();
+            writer.write(&batch_b).unwrap();
+            writer.finish().unwrap();
+        }
+
+        let batches = decode_chunk(&Bytes::from(buf)).unwrap();
+        assert_eq!(
+            batches.len(),
+            2,
+            "both record batches in the stream must be decoded, not just the first"
+        );
+        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+        assert_eq!(total_rows, 5);
+    }
+
+    /// Regression test for a bug caught in code review before it shipped: an
+    /// empty (zero-byte) blob made `StreamDecoder`'s decode loop a no-op and
+    /// `decoder.finish()` saw a still-pristine state, which it treats as a
+    /// clean empty stream -- silently returning zero batches with no error at
+    /// all, instead of the loud failure the old `StreamReader`-based version
+    /// gave on the same input ("Expected schema message, found empty
+    /// stream"). Same silent-truncation shape as the real multi-frame LZ4 bug
+    /// this crate already shipped once (see `client.rs`'s
+    /// `decompress_lz4_frame` doc comment) -- a genuinely empty chunk blob
+    /// must never be mistaken for "legitimately zero rows".
+    #[test]
+    fn decode_chunk_rejects_an_empty_blob() {
+        let err = decode_chunk(&Bytes::new()).expect_err("an empty blob must error, not silently decode to zero rows");
+        assert!(
+            err.message.contains("empty"),
+            "error should mention the blob was empty: {}",
+            err.message
+        );
+    }
+
+    /// Companion to `decode_chunk_rejects_an_empty_blob` -- proves the empty-
+    /// blob check doesn't overcorrect: a *non-empty* stream containing only a
+    /// schema message and no `RecordBatch` at all (a legitimate shape for a
+    /// genuinely empty query result) must still decode successfully to zero
+    /// batches, not error.
+    #[test]
+    fn decode_chunk_accepts_a_schema_only_stream_with_zero_batches() {
+        use arrow::datatypes::{Field, Schema};
+        use arrow::ipc::writer::StreamWriter;
+
+        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
+        let mut buf = Vec::new();
+        {
+            let mut writer = StreamWriter::try_new(&mut buf, &schema).unwrap();
+            writer.finish().unwrap();
+        }
+
+        let batches = decode_chunk(&Bytes::from(buf)).unwrap();
+        assert_eq!(
+            batches.len(),
+            0,
+            "a schema-only stream with no batches is valid, not an error"
+        );
+    }
+
+    /// Regression/documentation test for a real behavior change found in code
+    /// review: `StreamDecoder` (unlike the old `StreamReader`) hard-errors on
+    /// any bytes left over after a stream's own EOS marker, instead of
+    /// silently ignoring them. Locking this in deliberately -- erroring beats
+    /// silently dropping whatever came after the truncation point, same
+    /// reasoning as the empty-blob check above -- even though real Databricks
+    /// chunks have not been observed to have trailing bytes.
+    #[test]
+    fn decode_chunk_errors_on_trailing_bytes_after_a_complete_stream() {
+        use arrow::ipc::writer::StreamWriter;
+
+        let batch = make_batch(vec![1, 2], vec![1.0, 2.0]);
+        let mut buf = Vec::new();
+        {
+            let mut writer = StreamWriter::try_new(&mut buf, &batch.schema()).unwrap();
+            writer.write(&batch).unwrap();
+            writer.finish().unwrap();
+        }
+        buf.extend_from_slice(&[0xAA; 8]);
+
+        decode_chunk(&Bytes::from(buf)).expect_err("trailing bytes after a complete stream's EOS marker must error");
+    }
+
+    /// Diagnostic only, not a correctness check (relative timing is too
+    /// flaky for CI) -- `cargo test --release -- --ignored --nocapture
+    /// decode_chunk_speed` to compare the current `StreamDecoder`-based
+    /// `decode_chunk` against the old `StreamReader`-based approach it
+    /// replaced, on a batch shaped like a real chunk (120 columns, mixed
+    /// types, 50k rows -- this session's own real-table benchmark).
+    #[test]
+    #[ignore]
+    fn decode_chunk_speed_vs_stream_reader() {
+        use arrow::array::{Float64Array, Int64Array, StringArray};
+        use arrow::datatypes::{Field, Schema};
+        use arrow::ipc::writer::StreamWriter;
+        use std::io::Cursor as IoCursor;
+        use std::time::Instant;
+
+        const ROWS: usize = 50_000;
+        const COLS: usize = 120;
+
+        let mut fields = Vec::with_capacity(COLS);
+        let mut columns: Vec<Arc<dyn Array>> = Vec::with_capacity(COLS);
+        for i in 0..COLS {
+            match i % 3 {
+                0 => {
+                    fields.push(Field::new(format!("c{i}"), DataType::Int64, false));
+                    columns.push(Arc::new(Int64Array::from((0..ROWS as i64).collect::<Vec<_>>())));
+                }
+                1 => {
+                    fields.push(Field::new(format!("c{i}"), DataType::Float64, false));
+                    columns.push(Arc::new(Float64Array::from(
+                        (0..ROWS).map(|r| r as f64 * 1.5).collect::<Vec<_>>(),
+                    )));
+                }
+                _ => {
+                    fields.push(Field::new(format!("c{i}"), DataType::Utf8, false));
+                    columns.push(Arc::new(StringArray::from(
+                        (0..ROWS).map(|r| format!("row-{r}")).collect::<Vec<_>>(),
+                    )));
+                }
+            }
+        }
+        let schema = Arc::new(Schema::new(fields));
+        let batch = RecordBatch::try_new(schema.clone(), columns).unwrap();
+
+        let mut buf = Vec::new();
+        {
+            let mut writer = StreamWriter::try_new(&mut buf, &schema).unwrap();
+            writer.write(&batch).unwrap();
+            writer.finish().unwrap();
+        }
+        let blob = Bytes::from(buf);
+
+        const ITERS: u32 = 30;
+
+        // Old approach: StreamReader over an IoCursor -- copies every
+        // column's data into freshly allocated buffers on every decode.
+        let old_start = Instant::now();
+        for _ in 0..ITERS {
+            let reader = arrow::ipc::reader::StreamReader::try_new(IoCursor::new(&blob[..]), None).unwrap();
+            let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().unwrap();
+            assert_eq!(batches[0].num_rows(), ROWS);
+        }
+        let old_elapsed = old_start.elapsed();
+
+        // New approach: this file's actual decode_chunk.
+        let new_start = Instant::now();
+        for _ in 0..ITERS {
+            let batches = decode_chunk(&blob).unwrap();
+            assert_eq!(batches[0].num_rows(), ROWS);
+        }
+        let new_elapsed = new_start.elapsed();
+
+        println!(
+            "decode_chunk speed ({COLS} cols x {ROWS} rows, {ITERS} iters): \
+             StreamReader (old) = {old_elapsed:?} ({:?}/iter), \
+             StreamDecoder (new) = {new_elapsed:?} ({:?}/iter)",
+            old_elapsed / ITERS,
+            new_elapsed / ITERS,
+        );
+    }
 }