spatialrust_arrow/
device.rs

1//! Arrow C Device Data Interface (CPU-only v1).
2
3use std::ptr;
4
5use spatialrust_core::{PointCloud, SpatialMetadata};
6
7use crate::{
8    cdata::{
9        export_point_cloud_c_data, import_point_cloud_c_data, ArrowArray, ExportedArrowSchema,
10    },
11    ArrowBridgeError, ArrowBridgeResult,
12};
13
14/// Arrow device type codes used by the C Device Data Interface.
15#[repr(i32)]
16#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
17pub enum ArrowDeviceType {
18    /// Host CPU memory.
19    Cpu = 1,
20    /// CUDA device memory (recognized but not exported by this crate yet).
21    Cuda = 2,
22}
23
24/// Arrow C Device Array wrapper around an [`ArrowArray`].
25#[repr(C)]
26pub struct ArrowDeviceArray {
27    /// Device-resident or host array payload.
28    pub array: ArrowArray,
29    /// Device type.
30    pub device_id: i64,
31    /// Device type code.
32    pub device_type: i32,
33    /// Sync event handle; unused for CPU.
34    pub sync_event: *mut std::ffi::c_void,
35    /// Reserved for ABI growth.
36    pub reserved: [i64; 3],
37}
38
39/// Owned CPU device-array export.
40pub struct ExportedArrowDeviceArray {
41    schema: ExportedArrowSchema,
42    device: Box<ArrowDeviceArray>,
43}
44
45impl ExportedArrowDeviceArray {
46    /// Returns the exported schema.
47    #[must_use]
48    pub fn schema(&mut self) -> &mut ExportedArrowSchema {
49        &mut self.schema
50    }
51
52    /// Returns a mutable device-array pointer for FFI handoff.
53    pub fn as_mut_ptr(&mut self) -> *mut ArrowDeviceArray {
54        self.device.as_mut()
55    }
56}
57
58impl Drop for ExportedArrowDeviceArray {
59    fn drop(&mut self) {
60        if let Some(release) = self.device.array.release {
61            unsafe { release(&mut self.device.array) };
62            self.device.array.release = None;
63        }
64    }
65}
66
67/// Exports a point cloud as an Arrow device array on CPU.
68pub fn export_point_cloud_device_array(
69    cloud: &PointCloud,
70) -> ArrowBridgeResult<ExportedArrowDeviceArray> {
71    let (schema, mut array) = export_point_cloud_c_data(cloud)?;
72    let moved = unsafe { ptr::read(array.as_mut_ptr()) };
73    unsafe {
74        (*array.as_mut_ptr()).release = None;
75        (*array.as_mut_ptr()).private_data = ptr::null_mut();
76        (*array.as_mut_ptr()).buffers = ptr::null_mut();
77        (*array.as_mut_ptr()).children = ptr::null_mut();
78    }
79    drop(array);
80    let device = Box::new(ArrowDeviceArray {
81        array: moved,
82        device_id: -1,
83        device_type: ArrowDeviceType::Cpu as i32,
84        sync_event: ptr::null_mut(),
85        reserved: [0; 3],
86    });
87    Ok(ExportedArrowDeviceArray { schema, device })
88}
89
90/// Imports a CPU Arrow device array into a point cloud.
91///
92/// Non-CPU device types are rejected until an explicit device copy path exists.
93///
94/// # Safety
95///
96/// `device_array` and `schema` must form a valid exported pair.
97pub unsafe fn import_point_cloud_device_array(
98    schema: *const crate::cdata::ArrowSchema,
99    device_array: *const ArrowDeviceArray,
100    metadata: SpatialMetadata,
101) -> ArrowBridgeResult<PointCloud> {
102    if device_array.is_null() {
103        return Err(ArrowBridgeError::NullPointer("device array".into()));
104    }
105    let device_array = &*device_array;
106    if device_array.device_type != ArrowDeviceType::Cpu as i32 {
107        return Err(ArrowBridgeError::InvalidConfiguration(format!(
108            "Arrow device type {} is not supported without an explicit device copy",
109            device_array.device_type
110        )));
111    }
112    import_point_cloud_c_data(schema, &device_array.array, metadata)
113}
114
115#[cfg(test)]
116mod tests {
117    use super::{
118        export_point_cloud_device_array, import_point_cloud_device_array, ArrowDeviceType,
119    };
120    use spatialrust_core::{
121        PointBuffer, PointBufferSet, PointCloud, SpatialMetadata, StandardSchemas,
122    };
123
124    #[test]
125    fn cpu_device_array_roundtrip() {
126        let mut buffers = PointBufferSet::new();
127        buffers.insert("x", PointBuffer::from_f32(vec![1.0]));
128        buffers.insert("y", PointBuffer::from_f32(vec![2.0]));
129        buffers.insert("z", PointBuffer::from_f32(vec![3.0]));
130        let cloud = PointCloud::try_from_parts(
131            StandardSchemas::point_xyz(),
132            buffers,
133            SpatialMetadata::default(),
134        )
135        .unwrap();
136        let mut exported = export_point_cloud_device_array(&cloud).unwrap();
137        let schema_ptr = exported.schema().as_mut_ptr();
138        let device_ptr = exported.as_mut_ptr();
139        assert_eq!(unsafe { (*device_ptr).device_type }, ArrowDeviceType::Cpu as i32);
140        let imported = unsafe {
141            import_point_cloud_device_array(schema_ptr, device_ptr, SpatialMetadata::default())
142        }
143        .unwrap();
144        assert_eq!(imported.field("x").unwrap().as_f32().unwrap(), &[1.0]);
145    }
146}
spatialrust_arrow/device.rs

spatialrust_arrow/
device.rs
