xlog_cuda/
arrow_device.rs

1use std::ffi::c_void;
2use std::sync::Arc;
3
4use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
5
6use crate::memory::{CudaBuffer, TrackedCudaSlice};
7
8pub const ARROW_DEVICE_CUDA: i32 = 2;
9
10pub(crate) struct ArrowCudaAllocation {
11    _buffer: Arc<CudaBuffer>,
12    _extra: Vec<TrackedCudaSlice<u8>>,
13}
14
15impl ArrowCudaAllocation {
16    pub(crate) fn new(buffer: Arc<CudaBuffer>, extra: Vec<TrackedCudaSlice<u8>>) -> Self {
17        Self {
18            _buffer: buffer,
19            _extra: extra,
20        }
21    }
22}
23
24// SAFETY: ArrowCudaAllocation is used only as a keepalive handle for GPU buffers.
25// Dropping CUDA allocations is thread-safe, and no device memory is accessed.
26unsafe impl Send for ArrowCudaAllocation {}
27unsafe impl Sync for ArrowCudaAllocation {}
28
29#[repr(C)]
30pub struct ArrowDeviceArray {
31    pub device_type: i32,
32    pub device_id: i32,
33    pub array: *mut FFI_ArrowArray,
34    pub schema: *mut FFI_ArrowSchema,
35    pub release: Option<unsafe extern "C" fn(*mut ArrowDeviceArray)>,
36    pub private_data: *mut c_void,
37}
38
39unsafe extern "C" fn release_arrow_device_array(ptr: *mut ArrowDeviceArray) {
40    if ptr.is_null() {
41        return;
42    }
43    let dev = &mut *ptr;
44    if !dev.array.is_null() {
45        // SAFETY: dev.array is non-null (checked); was originally created via Box::into_raw; we are the sole owner
46        unsafe {
47            drop(Box::from_raw(dev.array));
48        }
49    }
50    if !dev.schema.is_null() {
51        // SAFETY: dev.schema is non-null (checked); was originally created via Box::into_raw; we are the sole owner
52        unsafe {
53            drop(Box::from_raw(dev.schema));
54        }
55    }
56    if !dev.private_data.is_null() {
57        // SAFETY: memory layout is guaranteed by the Arrow C Data Interface specification
58        unsafe {
59            drop(Box::from_raw(
60                dev.private_data.cast::<ArrowCudaAllocation>(),
61            ));
62        }
63    }
64    dev.release = None;
65}
66
67pub struct ArrowDeviceArrayOwned {
68    ptr: *mut ArrowDeviceArray,
69}
70
71// SAFETY: both constructors require thread-safe ownership transfer and release.
72// The wrapper only reads metadata or invokes the unique release callback; raw
73// producer mutation while the wrapper is live violates that import contract.
74unsafe impl Send for ArrowDeviceArrayOwned {}
75unsafe impl Sync for ArrowDeviceArrayOwned {}
76
77impl ArrowDeviceArrayOwned {
78    pub fn as_ptr(&self) -> *mut ArrowDeviceArray {
79        self.ptr
80    }
81
82    pub fn into_raw(self) -> *mut ArrowDeviceArray {
83        let ptr = self.ptr;
84        std::mem::forget(self);
85        ptr
86    }
87
88    /// Rebuild an owned wrapper from a raw `ArrowDeviceArray` pointer.
89    ///
90    /// # Safety
91    /// `ptr` must be a valid, uniquely owned pointer produced by
92    /// `ArrowDeviceArrayOwned::into_raw` or an equivalent allocation that
93    /// transfers ownership of the underlying `ArrowDeviceArray`.
94    /// All referenced device ranges must be live and producer-ready at handoff.
95    /// The producer must coordinate further access, and metadata plus release
96    /// must be usable from any host thread for the entire transferred lifetime.
97    pub unsafe fn from_raw(ptr: *mut ArrowDeviceArray) -> Self {
98        Self { ptr }
99    }
100}
101
102impl Drop for ArrowDeviceArrayOwned {
103    fn drop(&mut self) {
104        // SAFETY: ptr is non-null (checked); release was set by the Arrow producer and is valid for the array lifetime
105        unsafe {
106            if !self.ptr.is_null() {
107                if let Some(release) = (*self.ptr).release {
108                    release(self.ptr);
109                }
110            }
111        }
112    }
113}
114
115impl ArrowDeviceArray {
116    #[expect(
117        clippy::new_ret_no_self,
118        reason = "the Arrow C device ABI constructor returns its paired FFI schema and array handles"
119    )]
120    /// Adopt heap-owned Arrow FFI objects containing ready device allocations.
121    ///
122    /// # Safety
123    /// `array` and `schema` must each come from `Box::into_raw`, be uniquely
124    /// transferred, and describe live producer-ready device memory. Their
125    /// buffers and callbacks must support use and release from any host thread;
126    /// further producer accesses must be coordinated with the recipient.
127    pub unsafe fn new(
128        device_type: i32,
129        device_id: i32,
130        array: *mut FFI_ArrowArray,
131        schema: *mut FFI_ArrowSchema,
132    ) -> ArrowDeviceArrayOwned {
133        let dev = ArrowDeviceArray {
134            device_type,
135            device_id,
136            array,
137            schema,
138            release: Some(release_arrow_device_array),
139            private_data: std::ptr::null_mut(),
140        };
141        ArrowDeviceArrayOwned {
142            ptr: Box::into_raw(Box::new(dev)),
143        }
144    }
145}
146
147/// Keepalive wrapper for imported Arrow device arrays.
148///
149/// Holds the Arrow ArrayData so the FFI buffers remain alive until all
150/// device-backed columns are dropped.
151pub struct ArrowDeviceImport {
152    _data: arrow::array::ArrayData,
153    // Release the parsed FFI buffers before their producer's outer keepalive.
154    _outer: Option<ArrowDeviceArrayOwned>,
155}
156
157impl ArrowDeviceImport {
158    pub fn new(data: arrow::array::ArrayData) -> Self {
159        Self {
160            _data: data,
161            _outer: None,
162        }
163    }
164
165    #[cfg(any(test, feature = "arrow-device-import"))]
166    pub(crate) fn from_ffi(data: arrow::array::ArrayData, outer: ArrowDeviceArrayOwned) -> Self {
167        Self {
168            _data: data,
169            _outer: Some(outer),
170        }
171    }
172
173    #[cfg(feature = "arrow-device-import")]
174    pub(crate) fn data(&self) -> &arrow::array::ArrayData {
175        &self._data
176    }
177}
178
179impl ArrowDeviceArrayOwned {
180    /// Take ownership of the underlying FFI array + schema.
181    ///
182    /// # Safety
183    /// The returned outer owner must outlive both FFI objects and every parsed
184    /// buffer derived from them. All three objects must eventually be released.
185    pub unsafe fn into_ffi_parts(self) -> (i32, i32, Self, FFI_ArrowArray, FFI_ArrowSchema) {
186        let dev = &mut *self.ptr;
187        let device_type = dev.device_type;
188        let device_id = dev.device_id;
189
190        // Arrow's move operation leaves empty FFI objects in their original
191        // locations. Their containing allocations still belong to the producer;
192        // a foreign producer need not have allocated either struct with Box.
193        let array = FFI_ArrowArray::from_raw(dev.array);
194        let schema = FFI_ArrowSchema::from_raw(dev.schema);
195        (device_type, device_id, self, array, schema)
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use std::sync::atomic::{AtomicUsize, Ordering};
203
204    unsafe extern "C" fn release_outer(ptr: *mut ArrowDeviceArray) {
205        let owner = unsafe { Box::from_raw(ptr) };
206        // The moved FFI values are released by ArrayData, not by these empty
207        // containers. The producer still owns each original container.
208        drop(unsafe { Box::from_raw(owner.array) });
209        drop(unsafe { Box::from_raw(owner.schema) });
210        let drops = unsafe { Box::from_raw(owner.private_data.cast::<Arc<AtomicUsize>>()) };
211        drops.fetch_add(1, Ordering::SeqCst);
212    }
213
214    #[test]
215    fn ffi_import_keeps_outer_producer_alive_until_last_import_owner() {
216        let drops = Arc::new(AtomicUsize::new(0));
217        let data = arrow::array::ArrayData::new_empty(&arrow::datatypes::DataType::Int32);
218        let (array, schema) = arrow::ffi::to_ffi(&data).unwrap();
219        let raw = Box::into_raw(Box::new(ArrowDeviceArray {
220            device_type: ARROW_DEVICE_CUDA,
221            device_id: 0,
222            array: Box::into_raw(Box::new(array)),
223            schema: Box::into_raw(Box::new(schema)),
224            release: Some(release_outer),
225            private_data: Box::into_raw(Box::new(Arc::clone(&drops))).cast(),
226        }));
227        // SAFETY: this zero-length host ownership fixture transfers valid FFI
228        // objects and a thread-safe producer release; it performs no GPU work.
229        let device = unsafe { ArrowDeviceArrayOwned::from_raw(raw) };
230        let (_, _, outer, array, schema) = unsafe { device.into_ffi_parts() };
231        assert_eq!(drops.load(Ordering::SeqCst), 0);
232        let parsed = unsafe { arrow::ffi::from_ffi(array, &schema) }.unwrap();
233        drop(schema);
234        let imported = Arc::new(ArrowDeviceImport::from_ffi(parsed, outer));
235        let alias = Arc::clone(&imported);
236        drop(imported);
237        assert_eq!(drops.load(Ordering::SeqCst), 0);
238        drop(alias);
239        assert_eq!(drops.load(Ordering::SeqCst), 1);
240    }
241}
xlog_cuda/arrow_device.rs

xlog_cuda/
arrow_device.rs
