[0ax1]

Summary

First of two stacked drafts: expose dictionary-preserving GPU batches so consumers can postpone string expansion until after combining batches. This PR adds the export API only; the companion PR #11422 adds an optional cuDF materialization helper and reports the ~20% end-to-end hot-read gain.

A fixed-schema Arrow device stream cannot describe batches whose dictionary encodings or index widths differ. The new iterator returns an independently owned schema and device array for each batch, using the existing DictionaryExport::Preserve policy.

Changes

• Add vx_cuda_scan_path_device_batches, vx_cuda_device_batches_next, and vx_cuda_device_batches_free.
• Reuse the existing scan options, ordered top-level projection, direct I/O, and batch splitting.
• Keep the existing fixed-schema Arrow stream API and its decoded output unchanged.
• Document ownership, stream synchronization, thread affinity, end-of-stream, and terminal-error behavior; regenerate the C header.
• Add 10 regression tests for varying dictionary/index types, plain batches, nulls, empty batches, projection, independent output lifetimes, invalid inputs, and terminal errors.

[0ax1]

@claude review

[github-actions]

Claude finished @0ax1's task in 2m 1s —— View job

---

Review of PR #11421: dictionary-preserving per-batch device export

• Gather diff and context
• Review device_batches/mod.rs and lib.rs changes
• Review tests, header, README
• Post review

Verdict: I found no correctness bugs. The API is small and follows the existing FFI conventions. The comments are only nits and a consistency question.

What I checked

• Moving the open into open_cuda_scan (lib.rs:342) doesn't change behavior. The existing projected stream uses the same options, projection and batching as before.
• Dictionary policy is set per handle. CudaExecutionCtx::with_dictionary_export replaces the ctx's own copy of CudaSession, so Preserve applies to this handle only. It doesn't change the shared session. Running the fixtures on a Decode session tests exactly this.
• Terminal-error handling holds up (mod.rs:90-110). Setting State::Failed before polling means that an error from the scan, the dtype check or the export leaves the handle failed. A panic caught by try_or's catch_unwind does too. A retry can't skip the batch that failed, and the polls == 2 assertion in assert_terminal covers it.
• Outputs are written only after the export succeeds (mod.rs:204-207), so on error both outputs stay as they were. The null/unaligned pointer checks run before next(), so they don't poll or poison the handle. Tests cover both.
• Freeing the handle with vx_cuda_device_batches_free (a plain Box::from_raw drop) matches the box_wrapper! _free pattern in vortex-ffi/src/lib.rs:123.
• Lifetimes: returned arrays outlive both the handle and the session, and the schema and array are released separately. The first test checks this.

Suggestions (non-blocking)

1. End-of-stream sets no device fields (mod.rs:84-87, 95-98). Both end-of-stream paths return ArrowDeviceArray::empty(), so device_type and device_id are 0. The existing stream's end-of-stream (vortex-cuda/src/arrow/mod.rs:430-434) sets device_type: ARROW_DEVICE_CUDA and the stream's device_id. The doc's end-of-stream check is release == NULL, so nothing breaks. Still, a consumer that reads device_type before checking release would see 0 here. Matching the existing stream would keep the two APIs consistent. A small helper would also remove the duplicated end-of-stream construction:

   fn eos() -> ArrowDeviceArrayWithSchema { /* empty schema + empty array */ }

   Fix this →

2. No device check after export (mod.rs:106-108). The existing stream calls check_device (arrow/mod.rs:490) and releases the array if the export reports a different device_type or device_id. The new path has no such check. A per-batch API doesn't strictly need a single-device rule, since each batch carries its own device_id. But if the handle is meant to promise one device (the ctx is bound to one device), the same check would catch exporter bugs early. Either way, it would help to say in the doc whether device_id can vary between batches.

3. README line wrap (README.md, last paragraph). The line Edit the Rust declarations in `src/` or `cbindgen.toml` goes past the ~100-column wrap used in the rest of the file.

4. Test naming / scope (nit). test_existing_device_stream_still_decodes_mixed_batches tests the old ArrowDeviceArrayStream, not the new API. It is a useful regression guard for "existing API unchanged". It could live next to the existing stream tests in lib.rs or vortex-cuda/src/arrow so that device_batches/tests.rs covers only the new handle.

Tests

The 10 tests cover the important edges:

• dictionary index widths of i16, i32 and i64 within one stream;
• plain batches not gaining identity dictionaries;
• a numeric dictionary next to a string dictionary, with device-resident inputs;
• nulls in both the dictionary values and the codes;
• a zero-row batch versus end-of-stream, and repeated end-of-stream;
• terminal errors from the scan, the dtype check, the export and a panic;
• invalid pointers;
• projection with batch_rows through the path API, including freeing the session early.

The tests read device buffers back through the public ABI rather than Rust helpers, which suits an FFI contract.

I did not run clippy or the tests (they need a CUDA device, and the repo guidance leaves verification to the user).

[codspeed]

Merging this PR will regress 1 benchmark

2 warnings

⚠️ Walltime results may be unreliable

Some benchmarks used the Walltime instrument outside CodSpeed Macro Runners, which can make the results inconsistent.

Learn more about Macro Runners

See affected benchmarks

⚠️ 1 benchmark is dominated by system calls

System calls cannot be simulated consistently and are left out of the measure. The reported value understates the real cost. Switch to the Walltime instrument to measure them.

Learn more about system calls

See affected benchmarks

⚡ 1 improved benchmark
❌ 1 regressed benchmark
✅ 1966 untouched benchmarks
⏩ 3 skipped benchmarks¹

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
❌	Walltime	cuda/bitpacked_u8/unpack/3bw[100M]	301.3 µs	336 µs	-10.33%
⚡	Walltime	scalar_subtract_neon	13.8 µs	12.3 µs	+12.45%

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing ad/cuda-device-batches (f9b64ce) with develop (d81bd4c)}

Footnotes

1. 3 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩