arrowbricks 3.1.3__tar.gz → 3.1.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/PKG-INFO +12 -2
  2. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/README.md +11 -1
  3. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/pyproject.toml +1 -1
  4. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/Cargo.lock +2 -1
  5. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/Cargo.toml +3 -1
  6. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/README.md +2 -2
  7. arrowbricks-3.1.4/rust/arrowbricks_core/proptest-regressions/json_convert.txt +7 -0
  8. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/client/download.rs +66 -10
  9. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/json_convert.rs +11 -1
  10. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/lib.rs +4 -1
  11. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/pipeline/ndjson.rs +155 -11
  12. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/pipeline/thrift_exec.rs +7 -2
  13. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests/wiremock_thrift.rs +55 -0
  14. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/test_ipc_stream.py +37 -0
  15. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/LICENSE +0 -0
  16. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/.gitignore +0 -0
  17. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/examples/duckdb_query.py +0 -0
  18. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/examples/fastapi_sse.py +0 -0
  19. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/rustfmt.toml +0 -0
  20. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/client/error.rs +0 -0
  21. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/client/model.rs +0 -0
  22. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/client/sea.rs +0 -0
  23. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/client/thrift_rpc.rs +0 -0
  24. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/client/volume.rs +0 -0
  25. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/client.rs +0 -0
  26. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/heartbeat.rs +0 -0
  27. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/pipeline/reorder.rs +0 -0
  28. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/pipeline/sea.rs +0 -0
  29. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/pipeline/stats.rs +0 -0
  30. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/pipeline/test_support.rs +0 -0
  31. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/pipeline.rs +0 -0
  32. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/src/thrift.rs +0 -0
  33. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests/common/mod.rs +0 -0
  34. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests/wiremock_pipeline.rs +0 -0
  35. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests/wiremock_volume_files.rs +0 -0
  36. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/conftest.py +0 -0
  37. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/test_parameters.py +0 -0
  38. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/test_stream_ndjson_lines.py +0 -0
  39. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/test_streaming.py +0 -0
  40. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/test_thrift.py +0 -0
  41. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/test_token_provider.py +0 -0
  42. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/test_volume_files.py +0 -0
  43. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/rust/arrowbricks_core/tests_py/thrift_mock.py +0 -0
  44. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/src/arrowbricks/__init__.py +0 -0
  45. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/src/arrowbricks/_core.pyi +0 -0
  46. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/src/arrowbricks/_streaming.py +0 -0
  47. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/src/arrowbricks/client.py +0 -0
  48. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/src/arrowbricks/cursor.py +0 -0
  49. {arrowbricks-3.1.3 → arrowbricks-3.1.4}/src/arrowbricks/py.typed +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: arrowbricks
3
- Version: 3.1.3
3
+ Version: 3.1.4
4
4
  Requires-Dist: arro3-core>=0.8 ; extra == 'arro3'
5
5
  Provides-Extra: arro3
6
6
  License-File: LICENSE
@@ -300,11 +300,19 @@ python examples/benchmark_vs_connector.py
300
300
  For comparing arrowbricks versions with identical APIs, use
301
301
  [`examples/benchmark_versions.py`](examples/benchmark_versions.py). It keeps
302
302
  one connection per version, discards warm-ups, alternates execution order,
303
- passes concurrency explicitly, and checks row counts. See
303
+ passes concurrency explicitly, and checks row counts. Add `--verify-ipc`
304
+ (requires `arro3-core`) to compare serialized Arrow results in memory after
305
+ timing. Results must have stable values, schema metadata, row order, and batch boundaries;
306
+ verification contributes to peak process memory, and checksums are not printed. See
304
307
  [`benchmarks/2026-09-06.md`](benchmarks/2026-09-06.md) for replay/cache measurements
305
308
  and [`benchmarks/2026-09-06-downloads.md`](benchmarks/2026-09-06-downloads.md)
306
309
  for the subsequent cloud-fetch scheduling measurements and limits.
307
310
 
311
+ Follow-up experiments cover [spare request slots](benchmarks/2026-09-06-spare-slots.md),
312
+ [concurrent replay and NDJSON encoding](benchmarks/2026-09-06-lowlevel.md),
313
+ and the rejected [LZ4 capacity](benchmarks/2026-09-06-lz4-capacity.md) and
314
+ [direct-fill buffer](benchmarks/2026-09-06-direct-fill.md) changes.
315
+
308
316
  Cached IPC replay can be measured without a warehouse:
309
317
 
310
318
  ```bash
@@ -315,6 +323,8 @@ Replays now share immutable input bytes across decoded tables, reducing
315
323
  repeated copies and memory use. Arrow may still copy misaligned fixed-width
316
324
  buffers or decompress IPC-compressed bodies. Keeping a small slice of a
317
325
  decoded array can retain the full source allocation until that slice is released.
326
+ Decoding releases the Python interpreter lock, so independent replay calls
327
+ can run concurrently on separate Python threads.
318
328
 
319
329
  `BENCHMARK_SQL` overrides the query, `BENCHMARK_RUNS` (default 3) controls how many timed runs to average.
320
330
 
@@ -288,11 +288,19 @@ python examples/benchmark_vs_connector.py
288
288
  For comparing arrowbricks versions with identical APIs, use
289
289
  [`examples/benchmark_versions.py`](examples/benchmark_versions.py). It keeps
290
290
  one connection per version, discards warm-ups, alternates execution order,
291
- passes concurrency explicitly, and checks row counts. See
291
+ passes concurrency explicitly, and checks row counts. Add `--verify-ipc`
292
+ (requires `arro3-core`) to compare serialized Arrow results in memory after
293
+ timing. Results must have stable values, schema metadata, row order, and batch boundaries;
294
+ verification contributes to peak process memory, and checksums are not printed. See
292
295
  [`benchmarks/2026-09-06.md`](benchmarks/2026-09-06.md) for replay/cache measurements
293
296
  and [`benchmarks/2026-09-06-downloads.md`](benchmarks/2026-09-06-downloads.md)
294
297
  for the subsequent cloud-fetch scheduling measurements and limits.
295
298
 
299
+ Follow-up experiments cover [spare request slots](benchmarks/2026-09-06-spare-slots.md),
300
+ [concurrent replay and NDJSON encoding](benchmarks/2026-09-06-lowlevel.md),
301
+ and the rejected [LZ4 capacity](benchmarks/2026-09-06-lz4-capacity.md) and
302
+ [direct-fill buffer](benchmarks/2026-09-06-direct-fill.md) changes.
303
+
296
304
  Cached IPC replay can be measured without a warehouse:
297
305
 
298
306
  ```bash
@@ -303,6 +311,8 @@ Replays now share immutable input bytes across decoded tables, reducing
303
311
  repeated copies and memory use. Arrow may still copy misaligned fixed-width
304
312
  buffers or decompress IPC-compressed bodies. Keeping a small slice of a
305
313
  decoded array can retain the full source allocation until that slice is released.
314
+ Decoding releases the Python interpreter lock, so independent replay calls
315
+ can run concurrently on separate Python threads.
306
316
 
307
317
  `BENCHMARK_SQL` overrides the query, `BENCHMARK_RUNS` (default 3) controls how many timed runs to average.
308
318
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "arrowbricks"
3
- version = "3.1.3"
3
+ version = "3.1.4"
4
4
  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."
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -178,7 +178,7 @@ dependencies = [
178
178
 
179
179
  [[package]]
180
180
  name = "arrowbricks_core"
181
- version = "3.1.3"
181
+ version = "3.1.4"
182
182
  dependencies = [
183
183
  "arrow-array",
184
184
  "arrow-buffer",
@@ -190,6 +190,7 @@ dependencies = [
190
190
  "chrono",
191
191
  "hyper-rustls",
192
192
  "lz4_flex",
193
+ "memchr",
193
194
  "proptest",
194
195
  "pyo3",
195
196
  "pyo3-arrow",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "arrowbricks_core"
3
- version = "3.1.3"
3
+ version = "3.1.4"
4
4
  edition = "2024"
5
5
  readme = "README.md"
6
6
 
@@ -17,6 +17,8 @@ arrow-ipc = "59.1.0"
17
17
  arrow-schema = "59.1.0"
18
18
  arrow-json = "59.1.0"
19
19
  bytes = "1.12.1"
20
+ # Already in the dependency graph; vectorized newline search in NDJSON output.
21
+ memchr = "2"
20
22
  # Already an unconditional transitive dependency of arrow-cast (pulled in by
21
23
  # pyo3-arrow) at this exact version -- declared directly here too so
22
24
  # json_convert.rs can decode BINARY columns without hand-rolling base64. No
@@ -73,7 +73,7 @@ unlike the table object itself.
73
73
  ## API
74
74
 
75
75
  - `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).
76
- - `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.
76
+ - `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.
77
77
  - `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.
78
78
  - `ResultSet.fetchall_arrow() -> Table` -- drains everything remaining.
79
79
  - `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.
82
82
  - `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.
83
83
  - `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.
84
84
  - `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.
85
- - `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.
85
+ - `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.
86
86
  - `HEARTBEAT` -- module-level singleton; compare with `is`, e.g. `if item is _core.HEARTBEAT: ...`.
87
87
 
88
88
  ## With DuckDB
@@ -0,0 +1,7 @@
1
+ # Seeds for failure cases proptest has generated in the past. It is
2
+ # automatically read and these particular cases re-run before any
3
+ # novel cases are generated.
4
+ #
5
+ # It is recommended to check this file in to source control so that
6
+ # everyone who runs the test benefits from these saved cases.
7
+ cc c23ce2dbe22693eff6644d643de77c6c113dc918c876434391fb93a05438d706 # shrinks to type_name = "STRUCT", precision = None, scale = None, type_text = Some(""), values = []
@@ -36,18 +36,15 @@ use super::model::QueryStatsAccumulator;
36
36
  /// after each `EndMark` and picks up the next concatenated frame on a
37
37
  /// subsequent `read_to_end` call against the *same* instance (verified: the
38
38
  /// decoder's position in the underlying byte slice carries over across
39
- /// calls) -- so looping `read_to_end` on one decoder until it stops growing
40
- /// `out` reads every frame without reconstructing a decoder per frame.
39
+ /// calls) -- so looping `read_to_end` on one decoder until its underlying
40
+ /// reader is exhausted reads every frame without reconstructing a decoder.
41
41
  pub(crate) fn decompress_lz4_frame(compressed: &Bytes) -> Result<Bytes, ApiError> {
42
42
  use std::io::Read;
43
- // `compressed.len()` is a real lower bound, but LZ4 on Arrow-IPC data
44
- // (long dictionary/offset-buffer runs, mostly-repeated bytes) typically
45
- // compresses several-fold -- estimating just the lower bound means the
46
- // real decompressed size almost always blows past initial capacity,
47
- // paying for repeated doubling-and-copy growth on every chunk. `* 4` is
48
- // a heuristic, not a guarantee (`Vec` still grows normally if it's wrong
49
- // either way) -- just a better starting point than the guaranteed-too-
50
- // small lower bound.
43
+ // LZ4 on Arrow-IPC data often compresses several-fold. Starting at the
44
+ // compressed size can therefore pay for repeated buffer growth. Neither
45
+ // size is a bound on the other: incompressible input can expand slightly.
46
+ // `* 4` remains a heuristic; Vec grows normally when it underestimates.
47
+ // See benchmark_lz4_capacity_hypotheses for the allocation tradeoffs.
51
48
  let mut out = Vec::with_capacity(compressed.len() * 4);
52
49
  let mut decoder = lz4_flex::frame::FrameDecoder::new(&compressed[..]);
53
50
  // Terminate on the *reader* being exhausted, not on "output stopped
@@ -301,6 +298,65 @@ impl DbClient {
301
298
  mod tests {
302
299
  use super::*;
303
300
 
301
+ #[test]
302
+ #[ignore = "manual release-mode decompression allocation experiment"]
303
+ fn benchmark_lz4_capacity_hypotheses() {
304
+ use std::hint::black_box;
305
+ use std::io::{Read, Write};
306
+ use std::time::Instant;
307
+
308
+ // Deterministic random bytes mixed with repeated bytes cover different
309
+ // compression ratios without depending on warehouse data.
310
+ for random_fraction in [0, 25, 50, 100] {
311
+ let mut state = 0x12345678_u64;
312
+ let input: Vec<u8> = (0..8 * 1024 * 1024)
313
+ .map(|i| {
314
+ state ^= state << 13;
315
+ state ^= state >> 7;
316
+ state ^= state << 17;
317
+ if i % 100 < random_fraction { state as u8 } else { 0 }
318
+ })
319
+ .collect();
320
+ let mut compressed = Vec::new();
321
+ for part in input.chunks(512 * 1024) {
322
+ let mut encoder = lz4_flex::frame::FrameEncoder::new(Vec::new());
323
+ encoder.write_all(part).unwrap();
324
+ compressed.extend(encoder.finish().unwrap());
325
+ }
326
+ let variants = [
327
+ ("compressed_x1", compressed.len()),
328
+ ("compressed_x2", compressed.len() * 2),
329
+ ("compressed_x4", compressed.len() * 4),
330
+ ("exact", input.len()),
331
+ ("short_hint", input.len() - 64),
332
+ ];
333
+ for round in 0..10 {
334
+ for index in 0..variants.len() {
335
+ let (variant, capacity) = variants[(index + round) % variants.len()];
336
+ let start = Instant::now();
337
+ let mut output = Vec::with_capacity(capacity);
338
+ let mut decoder = lz4_flex::frame::FrameDecoder::new(black_box(compressed.as_slice()));
339
+ while !decoder.get_ref().is_empty() {
340
+ decoder.read_to_end(&mut output).unwrap();
341
+ }
342
+ let seconds = start.elapsed().as_secs_f64();
343
+ assert_eq!(output, input);
344
+ if round > 1 {
345
+ println!(
346
+ "{}",
347
+ serde_json::json!({
348
+ "random_percent": random_fraction, "variant": variant,
349
+ "round": round, "seconds": seconds,
350
+ "compressed_bytes": compressed.len(), "decoded_bytes": output.len(),
351
+ "retained_capacity": output.capacity()
352
+ })
353
+ );
354
+ }
355
+ }
356
+ }
357
+ }
358
+ }
359
+
304
360
  #[tokio::test]
305
361
  async fn split_download_starts_tail_requests_before_the_probe_body_finishes() {
306
362
  use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -334,7 +334,8 @@ fn build_column(
334
334
 
335
335
  let validity: Vec<bool> = rows_parsed.iter().map(Option::is_some).collect();
336
336
  let fields = Fields::from(child_fields);
337
- let struct_array = StructArray::new(fields.clone(), child_arrays, Some(NullBuffer::from(validity)));
337
+ let struct_array = StructArray::try_new(fields.clone(), child_arrays, Some(NullBuffer::from(validity)))
338
+ .map_err(|e| ApiError::permanent(format!("column `{name}`: invalid STRUCT array: {e}")))?;
338
339
  Ok((DataType::Struct(fields), Arc::new(struct_array)))
339
340
  }
340
341
  other => Err(ApiError {
@@ -702,6 +703,15 @@ mod tests {
702
703
  assert!(s.is_null(1), "row 1's whole struct must be null");
703
704
  }
704
705
 
706
+ #[test]
707
+ fn empty_struct_returns_a_conversion_error_instead_of_panicking() {
708
+ let columns = vec![struct_col("s", "STRUCT<>")];
709
+ for rows in [vec![], vec![vec![Some("{}".to_string())]], vec![vec![None]]] {
710
+ let err = json_array_to_record_batch(&rows, &columns).unwrap_err();
711
+ assert!(err.message.contains("invalid STRUCT array"));
712
+ }
713
+ }
714
+
705
715
  #[test]
706
716
  fn struct_with_a_nested_unsupported_composite_field_errors_instead_of_guessing() {
707
717
  let columns = vec![struct_col("s", "STRUCT<a: INT, nested: STRUCT<x: INT>>")];
@@ -62,11 +62,14 @@ fn write_ipc_stream(py: Python<'_>, stream: Bound<'_, PyAny>, buf: Bound<'_, PyA
62
62
  #[pyfunction]
63
63
  #[pyo3(signature = (data))]
64
64
  fn read_ipc_stream(data: Bound<'_, PyBytes>) -> PyResult<PyTable> {
65
+ let py = data.py();
65
66
  // The immutable Python bytes own the memory for as long as any decoded
66
67
  // array needs it, including arrays exported through the C Data Interface.
67
68
  // PyBackedBytes provides that ownership without copying or custom unsafe code.
68
69
  let blob = bytes::Bytes::from_owner(pyo3::pybacked::PyBackedBytes::from(data));
69
- let (batches, schema) = pipeline::decode_ipc_stream(&blob).map_err(|e| PyRuntimeError::new_err(e.message))?;
70
+ let (batches, schema) = py
71
+ .detach(|| pipeline::decode_ipc_stream(&blob))
72
+ .map_err(|e| PyRuntimeError::new_err(e.message))?;
70
73
  PyTable::try_new(batches, schema).map_err(|e| PyRuntimeError::new_err(e.to_string()))
71
74
  }
72
75
 
@@ -6,6 +6,7 @@
6
6
  //! `Infinity`/`-Infinity`) string-patching arrow-json's own fixed encoding
7
7
  //! needs help with.
8
8
 
9
+ use std::io::{self, Write};
9
10
  use std::sync::Arc;
10
11
 
11
12
  use arrow_array::Array;
@@ -22,6 +23,37 @@ use super::sea::submit_sea_and_report;
22
23
  use super::stats::{ReportOnDrop, StatsReporter};
23
24
  use super::thrift_exec::{ThriftSubmitResult, submit_thrift_and_start_fetch};
24
25
 
26
+ /// Collect rows as Arrow writes them, avoiding a second, chunk-sized buffer.
27
+ /// Writes may end within a row (or even a UTF-8 character).
28
+ struct LineCollector {
29
+ lines: Vec<String>,
30
+ pending: Vec<u8>,
31
+ }
32
+
33
+ impl Write for LineCollector {
34
+ fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
35
+ let mut start = 0;
36
+ for end in memchr::memchr_iter(b'\n', bytes) {
37
+ let line = &bytes[start..end];
38
+ let row = if self.pending.is_empty() {
39
+ line.to_vec()
40
+ } else {
41
+ self.pending.extend_from_slice(line);
42
+ std::mem::take(&mut self.pending)
43
+ };
44
+ self.lines
45
+ .push(String::from_utf8(row).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?);
46
+ start = end + 1;
47
+ }
48
+ self.pending.extend_from_slice(&bytes[start..]);
49
+ Ok(bytes.len())
50
+ }
51
+
52
+ fn flush(&mut self) -> io::Result<()> {
53
+ Ok(())
54
+ }
55
+ }
56
+
25
57
  /// Converts one chunk's decoded batches into NDJSON lines, one per row, in
26
58
  /// arro3-`write_ndjson(explicit_nulls=True)`-compatible format: null-valued
27
59
  /// keys stay present as JSON `null` rather than being omitted, and a
@@ -43,10 +75,13 @@ fn encode_ndjson_lines(batches: &[RecordBatch], non_finite_as_string: bool) -> R
43
75
  if batches.is_empty() {
44
76
  return Ok(Vec::new());
45
77
  }
46
- let mut buf = Vec::new();
78
+ let mut output = LineCollector {
79
+ lines: Vec::with_capacity(batches.iter().map(RecordBatch::num_rows).sum()),
80
+ pending: Vec::new(),
81
+ };
47
82
  {
48
83
  let builder = arrow_json::WriterBuilder::new().with_explicit_nulls(true);
49
- let mut writer = builder.build::<_, arrow_json::writer::LineDelimited>(&mut buf);
84
+ let mut writer = builder.build::<_, arrow_json::writer::LineDelimited>(&mut output);
50
85
  let refs: Vec<&RecordBatch> = batches.iter().collect();
51
86
  writer.write_batches(&refs).map_err(|e| ApiError {
52
87
  message: format!("NDJSON encode error: {e}"),
@@ -59,15 +94,10 @@ fn encode_ndjson_lines(batches: &[RecordBatch], non_finite_as_string: bool) -> R
59
94
  kind: ApiErrorKind::Other,
60
95
  })?;
61
96
  }
62
- let mut lines: Vec<String> = String::from_utf8(buf)
63
- .map_err(|e| ApiError {
64
- message: format!("NDJSON encode produced invalid UTF-8: {e}"),
65
- transient: false,
66
- kind: ApiErrorKind::Other,
67
- })?
68
- .lines()
69
- .map(|line| line.to_string())
70
- .collect();
97
+ if !output.pending.is_empty() {
98
+ return Err(ApiError::permanent("NDJSON encode produced an unterminated row"));
99
+ }
100
+ let mut lines = output.lines;
71
101
 
72
102
  if non_finite_as_string {
73
103
  patch_non_finite_floats(batches, &mut lines);
@@ -369,6 +399,120 @@ mod tests {
369
399
  use super::*;
370
400
  use crate::pipeline::test_support::make_batch;
371
401
 
402
+ fn encode_ndjson_reference(batches: &[RecordBatch], non_finite_as_string: bool) -> Result<Vec<String>, ApiError> {
403
+ if batches.is_empty() {
404
+ return Ok(Vec::new());
405
+ }
406
+ let mut buf = Vec::new();
407
+ {
408
+ let builder = arrow_json::WriterBuilder::new().with_explicit_nulls(true);
409
+ let mut writer = builder.build::<_, arrow_json::writer::LineDelimited>(&mut buf);
410
+ let refs: Vec<&RecordBatch> = batches.iter().collect();
411
+ writer.write_batches(&refs).map_err(|e| ApiError {
412
+ message: format!("NDJSON encode error: {e}"),
413
+ transient: false,
414
+ kind: ApiErrorKind::Other,
415
+ })?;
416
+ writer.finish().map_err(|e| ApiError {
417
+ message: format!("NDJSON encode error: {e}"),
418
+ transient: false,
419
+ kind: ApiErrorKind::Other,
420
+ })?;
421
+ }
422
+ let mut lines: Vec<String> = String::from_utf8(buf)
423
+ .map_err(|e| ApiError {
424
+ message: format!("NDJSON encode produced invalid UTF-8: {e}"),
425
+ transient: false,
426
+ kind: ApiErrorKind::Other,
427
+ })?
428
+ .lines()
429
+ .map(|line| line.to_string())
430
+ .collect();
431
+
432
+ if non_finite_as_string {
433
+ patch_non_finite_floats(batches, &mut lines);
434
+ }
435
+ Ok(lines)
436
+ }
437
+
438
+ #[test]
439
+ fn line_collector_handles_split_unicode_and_rejects_invalid_utf8() {
440
+ let mut output = LineCollector {
441
+ lines: Vec::new(),
442
+ pending: Vec::new(),
443
+ };
444
+ for byte in "{\"label\":\"café 🦀\"}\n{}\n".as_bytes() {
445
+ output.write_all(&[*byte]).unwrap();
446
+ }
447
+ assert_eq!(output.lines, ["{\"label\":\"café 🦀\"}", "{}"]);
448
+ assert!(output.pending.is_empty());
449
+ assert_eq!(
450
+ output.write_all(&[0xff, b'\n']).unwrap_err().kind(),
451
+ io::ErrorKind::InvalidData
452
+ );
453
+ }
454
+
455
+ fn text_batch(rows: usize, columns: usize, width: usize) -> RecordBatch {
456
+ use arrow_array::StringArray;
457
+ use arrow_schema::{Field, Schema};
458
+ let value = format!("café 🦀 \"quoted\"\n{}", "x".repeat(width));
459
+ let array = Arc::new(StringArray::from_iter(
460
+ (0..rows).map(|i| (i % 7 != 0).then_some(value.as_str())),
461
+ ));
462
+ let fields = (0..columns)
463
+ .map(|i| Field::new(format!("c{i}"), DataType::Utf8, true))
464
+ .collect::<Vec<_>>();
465
+ RecordBatch::try_new(
466
+ Arc::new(Schema::new(fields)),
467
+ (0..columns).map(|_| array.clone() as _).collect(),
468
+ )
469
+ .unwrap()
470
+ }
471
+
472
+ #[test]
473
+ fn collected_ndjson_matches_reference_for_large_unicode_rows_and_multiple_batches() {
474
+ let batches = vec![
475
+ text_batch(20, 3, 12000),
476
+ text_batch(0, 3, 12000),
477
+ text_batch(17, 3, 12000),
478
+ ];
479
+ let expected = encode_ndjson_reference(&batches, false).unwrap();
480
+ let actual = encode_ndjson_lines(&batches, false).unwrap();
481
+ assert_eq!(actual.len(), 37);
482
+ assert_eq!(actual, expected);
483
+ }
484
+
485
+ #[test]
486
+ #[ignore = "manual release-mode allocation/encoding benchmark"]
487
+ fn benchmark_ndjson_collector() {
488
+ use std::hint::black_box;
489
+ use std::time::Instant;
490
+ for (name, rows, columns, width) in [("narrow", 100000, 4, 16), ("wide", 20000, 32, 96)] {
491
+ let batches = vec![text_batch(rows, columns, width)];
492
+ assert_eq!(
493
+ encode_ndjson_lines(&batches, false).unwrap(),
494
+ encode_ndjson_reference(&batches, false).unwrap()
495
+ );
496
+ for round in 0..9 {
497
+ for candidate in if round % 2 == 0 { [false, true] } else { [true, false] } {
498
+ let start = Instant::now();
499
+ let lines = if candidate {
500
+ encode_ndjson_lines(black_box(&batches), false)
501
+ } else {
502
+ encode_ndjson_reference(black_box(&batches), false)
503
+ }
504
+ .unwrap();
505
+ let elapsed = start.elapsed().as_secs_f64() * 1000.0;
506
+ black_box(&lines);
507
+ println!(
508
+ "{}",
509
+ serde_json::json!({"workload":name,"round":round,"candidate":candidate,"ms":elapsed,"rows":lines.len(),"bytes":lines.iter().map(String::len).sum::<usize>()})
510
+ );
511
+ }
512
+ }
513
+ }
514
+ }
515
+
372
516
  #[test]
373
517
  fn non_finite_token_covers_nan_and_both_infinities_only() {
374
518
  assert_eq!(non_finite_token(f64::NAN), Some("\"NaN\""));
@@ -686,8 +686,13 @@ async fn run_thrift_fetch_loop(
686
686
  // Share the budget across all links already known in this batch.
687
687
  // Otherwise the first workers can claim eight slots each while
688
688
  // the remaining files wait without even starting a request.
689
- let split_limit = (concurrency / row_set.result_links.len().max(1)).max(1);
690
- for link in row_set.result_links {
689
+ let link_count = row_set.result_links.len().max(1);
690
+ for (link_position, link) in row_set.result_links.into_iter().enumerate() {
691
+ // Distribute the remainder too: 35 files under a 64-slot budget
692
+ // get 29 two-slot shares and six one-slot shares, rather than
693
+ // leaving 29 slots unused. Every file still gets at least one;
694
+ // the semaphore queues files when there are more than slots.
695
+ let split_limit = (concurrency / link_count + usize::from(link_position < concurrency % link_count)).max(1);
691
696
  let idx = chunk_index;
692
697
  chunk_index += 1;
693
698
  stats.chunks_seen.fetch_add(1, Ordering::Relaxed);
@@ -1456,6 +1456,61 @@ async fn mount_single_link_statement(server: &MockServer, link_path: &'static st
1456
1456
  .await;
1457
1457
  }
1458
1458
 
1459
+ #[tokio::test]
1460
+ async fn thrift_distributes_spare_slots_across_known_links() {
1461
+ let server = MockServer::start().await;
1462
+ mount_open_session_always(&server, b"sess").await;
1463
+ mount_close_operation_ok(&server).await;
1464
+ let direct = DirectResultsSpec {
1465
+ operation_state: Some(operation_state::FINISHED),
1466
+ metadata: Some((false, None)),
1467
+ fetch: Some(FetchSpec {
1468
+ result_links: vec![
1469
+ (format!("{}/_data/first", server.uri()), 60_000),
1470
+ (format!("{}/_data/second", server.uri()), 60_000),
1471
+ ],
1472
+ ..Default::default()
1473
+ }),
1474
+ ..Default::default()
1475
+ };
1476
+ Mock::given(method("POST"))
1477
+ .and(IsThriftRpc("ExecuteStatement"))
1478
+ .respond_with(ResponseTemplate::new(200).set_body_raw(
1479
+ build_execute_statement_resp(b"op", b"secret", Some(direct)),
1480
+ "application/x-thrift",
1481
+ ))
1482
+ .mount(&server)
1483
+ .await;
1484
+ let first = mount_ranged_blob(
1485
+ &server,
1486
+ "/_data/first",
1487
+ build_full_stream_bytes(&test_schema(), 0, 60_000),
1488
+ )
1489
+ .await;
1490
+ let second = mount_ranged_blob(
1491
+ &server,
1492
+ "/_data/second",
1493
+ build_full_stream_bytes(&test_schema(), 60_000, 120_000),
1494
+ )
1495
+ .await;
1496
+ let client = Arc::new(
1497
+ DbClient::new(&server.uri(), WAREHOUSE_ID, "fake-token")
1498
+ .with_protocol(Protocol::Thrift)
1499
+ .with_concurrency(3),
1500
+ );
1501
+ let mut stream = execute_lazy_thrift(client, "SELECT * FROM t", None, None, None)
1502
+ .await
1503
+ .unwrap();
1504
+ let (batches, _) = stream.fetchall_arrow().await.unwrap();
1505
+ assert_ids_in_order(&batches, 120_000);
1506
+ assert_eq!(
1507
+ first.load(Ordering::SeqCst),
1508
+ 2,
1509
+ "the spare slot should split the first file"
1510
+ );
1511
+ assert_eq!(second.load(Ordering::SeqCst), 1, "the second file keeps its own slot");
1512
+ }
1513
+
1459
1514
  #[tokio::test]
1460
1515
  async fn thrift_single_link_downloads_via_parallel_range_requests() {
1461
1516
  let server = MockServer::start().await;
@@ -127,3 +127,40 @@ def test_read_ipc_stream_result_is_callable_more_than_once():
127
127
  second = arrowbricks_core.read_ipc_stream(data)
128
128
  assert first.num_rows == 3
129
129
  assert second.num_rows == 3
130
+
131
+
132
+ def test_read_ipc_stream_allows_another_python_thread_to_run():
133
+ import threading
134
+
135
+ table = core.Table.from_pydict(
136
+ {"label": core.Array(["synthetic text " * 16] * 100_000, type=core.DataType.string())}
137
+ )
138
+ buf = io.BytesIO()
139
+ arrowbricks_core.write_ipc_stream(table, buf)
140
+ data = buf.getvalue()
141
+ ready = threading.Event()
142
+ go = threading.Event()
143
+ progressed = threading.Event()
144
+
145
+ def other_thread():
146
+ ready.set()
147
+ go.wait()
148
+ progressed.set()
149
+
150
+ thread = threading.Thread(target=other_thread)
151
+ thread.start()
152
+ ready.wait()
153
+ previous = sys.getswitchinterval()
154
+ try:
155
+ # Prevent a periodic Python bytecode switch from satisfying the test.
156
+ # The decoder itself must release the interpreter for the worker.
157
+ sys.setswitchinterval(10)
158
+ go.set()
159
+ result = arrowbricks_core.read_ipc_stream(data)
160
+ ran_during_decode = progressed.is_set()
161
+ finally:
162
+ sys.setswitchinterval(previous)
163
+ go.set()
164
+ thread.join(timeout=5)
165
+ assert result.num_rows == 100_000
166
+ assert ran_during_decode, "IPC decoding must release the interpreter lock"
File without changes