feat(ffi): add FFI_ArrowDeviceArray and FFI_ArrowDeviceArrayStream - #10754
Draft
wsulais wants to merge 1 commit into
Draft
feat(ffi): add FFI_ArrowDeviceArray and FFI_ArrowDeviceArrayStream#10754wsulais wants to merge 1 commit into
FFI_ArrowDeviceArray and FFI_ArrowDeviceArrayStream#10754wsulais wants to merge 1 commit into
Conversation
# Which issue does this PR close? - Closes apache#10752. # Rationale for this change arrow-rs binds `ArrowArray`, `ArrowSchema` and `ArrowArrayStream`, but nothing from the [C Device Data Interface](https://arrow.apache.org/docs/format/CDeviceDataInterface.html). A crate holding a device-resident Arrow array has no shared struct to pass it through, so it defines its own `#[repr(C)]` copy of `ArrowDeviceArray`, and two crates that both do that cannot interoperate. Arrow C++, nanoarrow and pyarrow all implement the interface already. This makes arrow-rs able to speak the interface and carry a device array's metadata faithfully. It does not make arrow-rs device-aware: no vendor dependency, no allocation, no copying, no synchronisation, and `sync_event` is never dereferenced. That is the split described in apache#7618 — kernels outside arrow-rs, but arrow-rs "able to receive arrays stored in GPU memory and pass them to kernels that can process them in the GPU". It is also a prerequisite for the async device stream interface raised in apache#7228: `ArrowAsyncTask::extract_data` writes into a `struct ArrowDeviceArray*`, so that interface cannot be bound in Rust until this struct exists. Async additionally raises the `futures` dependency question from apache#7228; the sync structs do not. It is not in this PR. # What changes are included in this PR? `arrow-data::ffi`, next to `FFI_ArrowArray`: - `ArrowDeviceType`, a `#[repr(transparent)]` newtype over `i32` with the constants from `abi.h`. Not an enum: the values track dlpack's `DLDeviceType` upstream, so a producer may send a device type this version does not know, and materialising an unrecognised discriminant into a Rust enum would be UB. - `FFI_ArrowDeviceArray`, with `new_cpu`, an `unsafe new` for callers that own device memory, `from_raw`, `empty`, and accessors. It has no release callback of its own, per the spec — the embedded `FFI_ArrowArray` owns the data. `arrow-array::ffi`: - `to_device_ffi`, `from_device_ffi`, `from_device_ffi_and_data_type`. Import rejects any `device_type` other than `ARROW_DEVICE_CPU` with an `ArrowError::CDataInterface` naming the device, rather than reading a device pointer as host memory. - Both types re-exported, so they reach `arrow::ffi` as `FFI_ArrowArray` does. `arrow-array::ffi_stream`: - `FFI_ArrowDeviceArrayStream`, with `new` from a `RecordBatchReader` (declaring `ARROW_DEVICE_CPU`) and an `unsafe new_unchecked` for a producer supplying its own callbacks. - `ArrowDeviceArrayStreamReader`, a `RecordBatchReader` over one, refusing a non-CPU stream at construction. Export declares a `device_id` of -1, which apache/arrow#41101 added to the spec for device types with no intrinsic device identifier and which Arrow C++ and pyarrow both emit. Import accepts any `device_id` for CPU data, since nanoarrow uses 0. No new dependencies or features; all of it sits behind the existing `ffi` feature. # Are these changes tested? Yes, and no GPU is needed. - `size_of`, `align_of` and `offset_of!` for both structs, gated to 64-bit targets. This is the check that matters: a `#[repr(C)]` struct with its fields in the wrong order compiles cleanly and round-trips fine within Rust. - CPU round trips for the array and the stream, shaped like the existing `to_ffi` / `from_ffi` tests, plus a doctest on `to_device_ffi`. - Non-CPU rejection for both. The stream test's `get_schema` and `get_next` are `unreachable!()`, so it fails if import touches the producer before checking `device_type`, and it asserts the refused stream is released rather than leaked. - An unknown `device_type` of 99 round-trips through the struct and is refused by number. - Both `device_id` conventions import for CPU. - `from_raw` leaves the source released, and a device array can be moved across a thread boundary. Cross-checked locally against two other implementations, which CI cannot do: - apache/arrow's `abi.h`, compiled with gcc 15.3: every `sizeof` and `offsetof` asserted here agrees, for both structs, as do the `ARROW_DEVICE_*` constants. - pyarrow 25.0.1: an int32 array with nulls and a utf8 array exported here and imported with `Array._import_from_c_device`, and the reverse via `Array._export_to_c_device`. pyarrow also emits `device_id == -1` for CPU data. I can add the pyarrow direction to `arrow-pyarrow-integration-testing` if that is wanted; it would need a minimum pyarrow version for the device methods. # Are there any user-facing changes? New public API only, and no breaking changes. arrow-rs still reads only CPU buffers: a non-CPU array is refused with an error, never dereferenced. Not included, each separable: the async device stream interface, device-resident `Buffer`s or kernels that can read them, and `sync_event` semantics beyond passing the pointer through. Assisted-by: Claude Opus 5 (claude-opus-5[1m]) via Claude Code 2.1.233
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
ArrowDeviceArray,ArrowDeviceArrayStream) over FFI #10752.Rationale for this change
arrow-rs binds
ArrowArray,ArrowSchemaandArrowArrayStream, but nothing from theC Device Data Interface. A
crate holding a device-resident Arrow array has no shared struct to pass it through, so it
defines its own
#[repr(C)]copy ofArrowDeviceArray, and two crates that both do that cannotinteroperate. Arrow C++, nanoarrow and pyarrow all implement the interface already.
This makes arrow-rs able to speak the interface and carry a device array's metadata faithfully.
It does not make arrow-rs device-aware: no vendor dependency, no allocation, no copying, no
synchronisation, and
sync_eventis never dereferenced. That is the split described in #7618 —kernels outside arrow-rs, but arrow-rs "able to receive arrays stored in GPU memory and pass them
to kernels that can process them in the GPU".
It is also a prerequisite for the async device stream interface raised in #7228:
ArrowAsyncTask::extract_datawrites into astruct ArrowDeviceArray*, so that interface cannotbe bound in Rust until this struct exists. Async additionally raises the
futuresdependencyquestion from #7228; the sync structs do not. It is not in this PR.
What changes are included in this PR?
arrow-data::ffi, next toFFI_ArrowArray:ArrowDeviceType, a#[repr(transparent)]newtype overi32with the constants fromabi.h. Not an enum: the values track dlpack'sDLDeviceTypeupstream, so a producer maysend a device type this version does not know, and materialising an unrecognised discriminant
into a Rust enum would be UB.
FFI_ArrowDeviceArray, withnew_cpu, anunsafe newfor callers that own device memory,from_raw,empty, and accessors. It has no release callback of its own, per the spec — theembedded
FFI_ArrowArrayowns the data.arrow-array::ffi:to_device_ffi,from_device_ffi,from_device_ffi_and_data_type. Import rejects anydevice_typeother thanARROW_DEVICE_CPUwith anArrowError::CDataInterfacenaming thedevice, rather than reading a device pointer as host memory.
arrow::ffiasFFI_ArrowArraydoes.arrow-array::ffi_stream:FFI_ArrowDeviceArrayStream, withnewfrom aRecordBatchReader(declaringARROW_DEVICE_CPU) and anunsafe new_uncheckedfor a producer supplying its own callbacks.ArrowDeviceArrayStreamReader, aRecordBatchReaderover one, refusing a non-CPU stream atconstruction.
Export declares a
device_idof -1, which apache/arrow#41101 added to the spec for device typeswith no intrinsic device identifier and which Arrow C++ and pyarrow both emit. Import accepts any
device_idfor CPU data, since nanoarrow uses 0.No new dependencies or features; all of it sits behind the existing
ffifeature.Are these changes tested?
Yes, and no GPU is needed.
size_of,align_ofandoffset_of!for both structs, gated to 64-bit targets. This is thecheck that matters: a
#[repr(C)]struct with its fields in the wrong order compiles cleanlyand round-trips fine within Rust.
to_ffi/from_ffitests, plus a doctest on
to_device_ffi.get_schemaandget_nextareunreachable!(),so it fails if import touches the producer before checking
device_type, and it asserts therefused stream is released rather than leaked.
device_typeof 99 round-trips through the struct and is refused by number.device_idconventions import for CPU.from_rawleaves the source released, and a device array can be moved across a threadboundary.
Cross-checked locally against two other implementations, which CI cannot do:
abi.h, compiled with gcc 15.3: everysizeofandoffsetofasserted hereagrees, for both structs, as do the
ARROW_DEVICE_*constants.Array._import_from_c_device, and the reverse viaArray._export_to_c_device. pyarrow alsoemits
device_id == -1for CPU data.I can add the pyarrow direction to
arrow-pyarrow-integration-testingif that is wanted; itwould need a minimum pyarrow version for the device methods.
Are there any user-facing changes?
New public API only, and no breaking changes. arrow-rs still reads only CPU buffers: a non-CPU
array is refused with an error, never dereferenced.
Not included, each separable: the async device stream interface, device-resident
Buffers orkernels that can read them, and
sync_eventsemantics beyond passing the pointer through.Assisted-by: Claude Opus 5 (claude-opus-5[1m]) via Claude Code 2.1.233