spatialrust_arrow/
cdata.rs

1//! Arrow C Data Interface export/import for `PointCloud` columns.
2
3use std::{
4    ffi::{c_char, c_void, CString},
5    ptr,
6};
7
8use spatialrust_core::{
9    DType, PointBuffer, PointBufferSet, PointCloud, PointField, PointSchema, SpatialMetadata,
10};
11
12use crate::{ArrowBridgeError, ArrowBridgeResult};
13
14/// Arrow C Data Interface schema object.
15#[repr(C)]
16#[derive(Debug)]
17pub struct ArrowSchema {
18    /// Format string (`f`, `g`, `i`, `+s`, ...).
19    pub format: *const c_char,
20    /// Field name.
21    pub name: *const c_char,
22    /// Optional metadata key/value blob.
23    pub metadata: *const c_char,
24    /// Dictionary / nullable flags.
25    pub flags: i64,
26    /// Child count for nested types.
27    pub n_children: i64,
28    /// Child schemas.
29    pub children: *mut *mut ArrowSchema,
30    /// Dictionary schema.
31    pub dictionary: *mut ArrowSchema,
32    /// Release callback.
33    pub release: Option<unsafe extern "C" fn(schema: *mut ArrowSchema)>,
34    /// Implementation private data.
35    pub private_data: *mut c_void,
36}
37
38/// Arrow C Data Interface array object.
39#[repr(C)]
40#[derive(Debug)]
41pub struct ArrowArray {
42    /// Logical length.
43    pub length: i64,
44    /// Null count (`-1` if unknown).
45    pub null_count: i64,
46    /// Buffer offset.
47    pub offset: i64,
48    /// Number of buffers.
49    pub n_buffers: i64,
50    /// Nested child count.
51    pub n_children: i64,
52    /// Buffer pointer table.
53    pub buffers: *mut *const c_void,
54    /// Child arrays.
55    pub children: *mut *mut ArrowArray,
56    /// Dictionary array.
57    pub dictionary: *mut ArrowArray,
58    /// Release callback.
59    pub release: Option<unsafe extern "C" fn(array: *mut ArrowArray)>,
60    /// Implementation private data.
61    pub private_data: *mut c_void,
62}
63
64/// Owned export that keeps `ArrowSchema` alive until dropped or released.
65pub struct ExportedArrowSchema {
66    raw: Box<ArrowSchema>,
67}
68
69impl ExportedArrowSchema {
70    /// Returns a mutable raw pointer suitable for FFI handoff.
71    pub fn as_mut_ptr(&mut self) -> *mut ArrowSchema {
72        self.raw.as_mut()
73    }
74
75    /// Borrows the schema.
76    #[must_use]
77    pub fn schema(&self) -> &ArrowSchema {
78        &self.raw
79    }
80}
81
82impl Drop for ExportedArrowSchema {
83    fn drop(&mut self) {
84        if let Some(release) = self.raw.release {
85            unsafe { release(self.raw.as_mut()) };
86        }
87    }
88}
89
90/// Owned export that keeps `ArrowArray` alive until dropped or released.
91pub struct ExportedArrowArray {
92    raw: Box<ArrowArray>,
93}
94
95impl ExportedArrowArray {
96    /// Returns a mutable raw pointer suitable for FFI handoff.
97    pub fn as_mut_ptr(&mut self) -> *mut ArrowArray {
98        self.raw.as_mut()
99    }
100
101    /// Borrows the array.
102    #[must_use]
103    pub fn array(&self) -> &ArrowArray {
104        &self.raw
105    }
106}
107
108impl Drop for ExportedArrowArray {
109    fn drop(&mut self) {
110        if let Some(release) = self.raw.release {
111            unsafe { release(self.raw.as_mut()) };
112        }
113    }
114}
115
116/// Exports a point cloud as an Arrow struct array plus matching schema.
117pub fn export_point_cloud_c_data(
118    cloud: &PointCloud,
119) -> ArrowBridgeResult<(ExportedArrowSchema, ExportedArrowArray)> {
120    cloud.validate()?;
121    let schema = export_schema(cloud.schema())?;
122    let array = export_array(cloud)?;
123    Ok((schema, array))
124}
125
126/// Imports a SpatialRust point cloud from Arrow C Data struct columns.
127///
128/// # Safety
129///
130/// `schema` and `array` must form a valid exported Arrow C Data pair produced
131/// for a SpatialRust point cloud (or an equivalent struct of primitive columns).
132pub unsafe fn import_point_cloud_c_data(
133    schema: *const ArrowSchema,
134    array: *const ArrowArray,
135    metadata: SpatialMetadata,
136) -> ArrowBridgeResult<PointCloud> {
137    if schema.is_null() || array.is_null() {
138        return Err(ArrowBridgeError::NullPointer("schema/array".into()));
139    }
140    let schema = &*schema;
141    let array = &*array;
142    if schema.format.is_null() {
143        return Err(ArrowBridgeError::NullPointer("schema.format".into()));
144    }
145    let format = std::ffi::CStr::from_ptr(schema.format).to_string_lossy();
146    if format.as_ref() != "+s" {
147        return Err(ArrowBridgeError::SchemaMismatch(format!(
148            "expected Arrow struct (+s), found {format}"
149        )));
150    }
151    if schema.n_children < 0 || array.n_children < 0 || schema.n_children != array.n_children {
152        return Err(ArrowBridgeError::SchemaMismatch("schema/array child counts disagree".into()));
153    }
154    let n = schema.n_children as usize;
155    let mut point_schema = PointSchema::new();
156    let mut buffers = PointBufferSet::new();
157    let length = usize::try_from(array.length).map_err(|_| {
158        ArrowBridgeError::InvalidConfiguration("array length does not fit usize".into())
159    })?;
160    for index in 0..n {
161        let child_schema = *schema.children.add(index);
162        let child_array = *array.children.add(index);
163        if child_schema.is_null() || child_array.is_null() {
164            return Err(ArrowBridgeError::NullPointer("child schema/array".into()));
165        }
166        let (field, buffer) = import_child(&*child_schema, &*child_array, length)?;
167        buffers.insert(field.name.clone(), buffer);
168        point_schema = point_schema.with_field(field);
169    }
170    Ok(PointCloud::try_from_parts(point_schema, buffers, metadata)?)
171}
172
173fn export_schema(point_schema: &PointSchema) -> ArrowBridgeResult<ExportedArrowSchema> {
174    let children = point_schema
175        .fields()
176        .iter()
177        .map(export_field_schema)
178        .collect::<ArrowBridgeResult<Vec<_>>>()?;
179    let mut child_ptrs = children.into_iter().map(Box::into_raw).collect::<Vec<*mut ArrowSchema>>();
180    let child_table = child_ptrs.as_mut_ptr();
181    let n_children = child_ptrs.len() as i64;
182    // Leak the vec table into private data; release rebuilds and frees it.
183    std::mem::forget(child_ptrs);
184
185    let format = CString::new("+s").expect("static");
186    let name = CString::new("PointCloud").expect("static");
187    let private = Box::new(SchemaPrivate {
188        format,
189        name,
190        children_table: child_table,
191        n_children: n_children as usize,
192    });
193    let raw = Box::new(ArrowSchema {
194        format: private.format.as_ptr(),
195        name: private.name.as_ptr(),
196        metadata: ptr::null(),
197        flags: 0,
198        n_children,
199        children: child_table,
200        dictionary: ptr::null_mut(),
201        release: Some(release_schema),
202        private_data: Box::into_raw(private) as *mut c_void,
203    });
204    Ok(ExportedArrowSchema { raw })
205}
206
207fn export_field_schema(field: &PointField) -> ArrowBridgeResult<Box<ArrowSchema>> {
208    if field.components != 1 {
209        return Err(ArrowBridgeError::InvalidConfiguration(
210            "Arrow C Data export currently supports scalar fields only".into(),
211        ));
212    }
213    let format = CString::new(dtype_format(field.dtype)?)
214        .map_err(|_| ArrowBridgeError::InvalidConfiguration("field format contained NUL".into()))?;
215    let name = CString::new(field.name.as_str())
216        .map_err(|_| ArrowBridgeError::InvalidConfiguration("field name contained NUL".into()))?;
217    let private =
218        Box::new(SchemaPrivate { format, name, children_table: ptr::null_mut(), n_children: 0 });
219    Ok(Box::new(ArrowSchema {
220        format: private.format.as_ptr(),
221        name: private.name.as_ptr(),
222        metadata: ptr::null(),
223        flags: 0,
224        n_children: 0,
225        children: ptr::null_mut(),
226        dictionary: ptr::null_mut(),
227        release: Some(release_schema),
228        private_data: Box::into_raw(private) as *mut c_void,
229    }))
230}
231
232fn export_array(cloud: &PointCloud) -> ArrowBridgeResult<ExportedArrowArray> {
233    let children = cloud
234        .schema()
235        .fields()
236        .iter()
237        .map(|field| {
238            let buffer = cloud.field(&field.name)?;
239            export_primitive_array(buffer, cloud.len())
240        })
241        .collect::<ArrowBridgeResult<Vec<_>>>()?;
242    let mut child_ptrs = children.into_iter().map(Box::into_raw).collect::<Vec<*mut ArrowArray>>();
243    let child_table = child_ptrs.as_mut_ptr();
244    let n_children = child_ptrs.len() as i64;
245    std::mem::forget(child_ptrs);
246
247    let private = Box::new(ArrayPrivate {
248        buffers_table: ptr::null_mut(),
249        n_buffers: 0,
250        children_table: child_table,
251        n_children: n_children as usize,
252        owned_buffers: Vec::new(),
253    });
254    let raw = Box::new(ArrowArray {
255        length: cloud.len() as i64,
256        null_count: 0,
257        offset: 0,
258        n_buffers: 1, // struct arrays expose a single validity/null bitmap buffer slot
259        n_children,
260        buffers: {
261            // Struct: one null bitmap pointer (null means non-nullable / no bitmap).
262            let mut buffers = vec![ptr::null()];
263            let table = buffers.as_mut_ptr();
264            // Store buffers table in private for release.
265            // Safety: overwrite after private box created - reassign below via raw.
266            std::mem::forget(buffers);
267            table
268        },
269        children: child_table,
270        dictionary: ptr::null_mut(),
271        release: Some(release_array),
272        private_data: ptr::null_mut(),
273    });
274    // Move buffers_table into private and attach private_data.
275    let mut raw = raw;
276    let buffers_table = raw.buffers;
277    let mut private = private;
278    private.buffers_table = buffers_table;
279    private.n_buffers = 1;
280    raw.private_data = Box::into_raw(private) as *mut c_void;
281    Ok(ExportedArrowArray { raw })
282}
283
284fn export_primitive_array(buffer: &PointBuffer, len: usize) -> ArrowBridgeResult<Box<ArrowArray>> {
285    if buffer.len() != len {
286        return Err(ArrowBridgeError::SchemaMismatch(
287            "column length does not match point count".into(),
288        ));
289    }
290    let owned = clone_bytes(buffer);
291    let data_ptr = owned.as_ptr() as *const c_void;
292    let mut buffers = vec![ptr::null(), data_ptr];
293    let buffers_table = buffers.as_mut_ptr();
294    std::mem::forget(buffers);
295    let private = Box::new(ArrayPrivate {
296        buffers_table,
297        n_buffers: 2,
298        children_table: ptr::null_mut(),
299        n_children: 0,
300        owned_buffers: vec![owned],
301    });
302    Ok(Box::new(ArrowArray {
303        length: len as i64,
304        null_count: 0,
305        offset: 0,
306        n_buffers: 2,
307        n_children: 0,
308        buffers: buffers_table,
309        children: ptr::null_mut(),
310        dictionary: ptr::null_mut(),
311        release: Some(release_array),
312        private_data: Box::into_raw(private) as *mut c_void,
313    }))
314}
315
316unsafe fn import_child(
317    schema: &ArrowSchema,
318    array: &ArrowArray,
319    expected_len: usize,
320) -> ArrowBridgeResult<(PointField, PointBuffer)> {
321    if schema.format.is_null() || schema.name.is_null() {
322        return Err(ArrowBridgeError::NullPointer("child format/name".into()));
323    }
324    let format = std::ffi::CStr::from_ptr(schema.format).to_string_lossy();
325    let name = std::ffi::CStr::from_ptr(schema.name)
326        .to_str()
327        .map_err(|_| ArrowBridgeError::InvalidConfiguration("field name is not UTF-8".into()))?
328        .to_owned();
329    let dtype = format_dtype(format.as_ref())?;
330    let length = usize::try_from(array.length).map_err(|_| {
331        ArrowBridgeError::InvalidConfiguration("child length does not fit usize".into())
332    })?;
333    if length != expected_len {
334        return Err(ArrowBridgeError::SchemaMismatch(format!(
335            "child `{name}` length {length} != parent {expected_len}"
336        )));
337    }
338    if array.n_buffers < 2 || array.buffers.is_null() {
339        return Err(ArrowBridgeError::InvalidConfiguration(
340            "primitive arrays require validity + data buffers".into(),
341        ));
342    }
343    let data_ptr = *array.buffers.add(1);
344    if length > 0 && data_ptr.is_null() {
345        return Err(ArrowBridgeError::NullPointer(format!("data for `{name}`")));
346    }
347    let buffer = copy_buffer(dtype, data_ptr, length)?;
348    let semantic = semantic_for_name(&name);
349    Ok((PointField::scalar(name, semantic, dtype), buffer))
350}
351
352fn clone_bytes(buffer: &PointBuffer) -> Vec<u8> {
353    match buffer {
354        PointBuffer::F32(values) => bytemuck_bytes(values),
355        PointBuffer::F64(values) => bytemuck_bytes(values),
356        PointBuffer::U8(values) => values.clone(),
357        PointBuffer::U16(values) => bytemuck_bytes(values),
358        PointBuffer::U32(values) => bytemuck_bytes(values),
359        PointBuffer::I32(values) => bytemuck_bytes(values),
360    }
361}
362
363fn bytemuck_bytes<T: Copy>(values: &[T]) -> Vec<u8> {
364    let mut bytes = Vec::with_capacity(std::mem::size_of_val(values));
365    for value in values {
366        let ptr = value as *const T as *const u8;
367        let slice = unsafe { std::slice::from_raw_parts(ptr, std::mem::size_of::<T>()) };
368        bytes.extend_from_slice(slice);
369    }
370    bytes
371}
372
373unsafe fn copy_buffer(
374    dtype: DType,
375    data_ptr: *const c_void,
376    length: usize,
377) -> ArrowBridgeResult<PointBuffer> {
378    if length == 0 {
379        return Ok(PointBuffer::with_capacity(dtype, 0));
380    }
381    Ok(match dtype {
382        DType::F32 => PointBuffer::F32(copy_typed(data_ptr, length)),
383        DType::F64 => PointBuffer::F64(copy_typed(data_ptr, length)),
384        DType::U8 => PointBuffer::U8(copy_typed(data_ptr, length)),
385        DType::U16 => PointBuffer::U16(copy_typed(data_ptr, length)),
386        DType::U32 => PointBuffer::U32(copy_typed(data_ptr, length)),
387        DType::I32 => PointBuffer::I32(copy_typed(data_ptr, length)),
388        DType::F16 => {
389            return Err(ArrowBridgeError::InvalidConfiguration(
390                "F16 Arrow import is not implemented".into(),
391            ))
392        }
393    })
394}
395
396unsafe fn copy_typed<T: Copy>(data_ptr: *const c_void, length: usize) -> Vec<T> {
397    let slice = std::slice::from_raw_parts(data_ptr as *const T, length);
398    slice.to_vec()
399}
400
401fn dtype_format(dtype: DType) -> ArrowBridgeResult<&'static str> {
402    match dtype {
403        DType::F32 => Ok("f"),
404        DType::F64 => Ok("g"),
405        DType::U8 => Ok("C"),
406        DType::U16 => Ok("S"),
407        DType::U32 => Ok("I"),
408        DType::I32 => Ok("i"),
409        DType::F16 => Err(ArrowBridgeError::InvalidConfiguration(
410            "F16 Arrow export is not implemented".into(),
411        )),
412    }
413}
414
415fn format_dtype(format: &str) -> ArrowBridgeResult<DType> {
416    match format {
417        "f" => Ok(DType::F32),
418        "g" => Ok(DType::F64),
419        "C" => Ok(DType::U8),
420        "S" => Ok(DType::U16),
421        "I" => Ok(DType::U32),
422        "i" => Ok(DType::I32),
423        other => {
424            Err(ArrowBridgeError::SchemaMismatch(format!("unsupported Arrow format `{other}`")))
425        }
426    }
427}
428
429fn semantic_for_name(name: &str) -> spatialrust_core::FieldSemantic {
430    use spatialrust_core::FieldSemantic;
431    match name {
432        "x" => FieldSemantic::PositionX,
433        "y" => FieldSemantic::PositionY,
434        "z" => FieldSemantic::PositionZ,
435        "intensity" => FieldSemantic::Intensity,
436        "nx" | "normal_x" => FieldSemantic::NormalX,
437        "ny" | "normal_y" => FieldSemantic::NormalY,
438        "nz" | "normal_z" => FieldSemantic::NormalZ,
439        "r" | "red" => FieldSemantic::ColorR,
440        "g" | "green" => FieldSemantic::ColorG,
441        "b" | "blue" => FieldSemantic::ColorB,
442        _ => FieldSemantic::Unknown,
443    }
444}
445
446struct SchemaPrivate {
447    format: CString,
448    name: CString,
449    children_table: *mut *mut ArrowSchema,
450    n_children: usize,
451}
452
453struct ArrayPrivate {
454    buffers_table: *mut *const c_void,
455    n_buffers: usize,
456    children_table: *mut *mut ArrowArray,
457    n_children: usize,
458    /// Keeps exported column bytes alive for the Arrow buffer pointers.
459    #[allow(dead_code)]
460    owned_buffers: Vec<Vec<u8>>,
461}
462
463unsafe extern "C" fn release_schema(schema: *mut ArrowSchema) {
464    if schema.is_null() {
465        return;
466    }
467    let schema = &mut *schema;
468    if schema.release.is_none() {
469        return;
470    }
471    if !schema.private_data.is_null() {
472        let private = Box::from_raw(schema.private_data as *mut SchemaPrivate);
473        if !private.children_table.is_null() && private.n_children > 0 {
474            let children =
475                Vec::from_raw_parts(private.children_table, private.n_children, private.n_children);
476            for child in children {
477                if !child.is_null() {
478                    if let Some(release) = (*child).release {
479                        release(child);
480                    }
481                    drop(Box::from_raw(child));
482                }
483            }
484        }
485        drop(private);
486    }
487    schema.release = None;
488    schema.private_data = ptr::null_mut();
489    schema.children = ptr::null_mut();
490    schema.format = ptr::null();
491    schema.name = ptr::null();
492}
493
494unsafe extern "C" fn release_array(array: *mut ArrowArray) {
495    if array.is_null() {
496        return;
497    }
498    let array = &mut *array;
499    if array.release.is_none() {
500        return;
501    }
502    if !array.private_data.is_null() {
503        let private = Box::from_raw(array.private_data as *mut ArrayPrivate);
504        if !private.children_table.is_null() && private.n_children > 0 {
505            let children =
506                Vec::from_raw_parts(private.children_table, private.n_children, private.n_children);
507            for child in children {
508                if !child.is_null() {
509                    if let Some(release) = (*child).release {
510                        release(child);
511                    }
512                    drop(Box::from_raw(child));
513                }
514            }
515        }
516        if !private.buffers_table.is_null() && private.n_buffers > 0 {
517            let _ =
518                Vec::from_raw_parts(private.buffers_table, private.n_buffers, private.n_buffers);
519        }
520        drop(private);
521    }
522    array.release = None;
523    array.private_data = ptr::null_mut();
524    array.buffers = ptr::null_mut();
525    array.children = ptr::null_mut();
526}
527
528#[cfg(test)]
529mod tests {
530    use super::{export_point_cloud_c_data, import_point_cloud_c_data};
531    use spatialrust_core::{
532        PointBuffer, PointBufferSet, PointCloud, SpatialMetadata, StandardSchemas,
533    };
534
535    #[test]
536    fn roundtrip_xyz_cloud() {
537        let mut buffers = PointBufferSet::new();
538        buffers.insert("x", PointBuffer::from_f32(vec![1.0, 2.0]));
539        buffers.insert("y", PointBuffer::from_f32(vec![3.0, 4.0]));
540        buffers.insert("z", PointBuffer::from_f32(vec![5.0, 6.0]));
541        let cloud = PointCloud::try_from_parts(
542            StandardSchemas::point_xyz(),
543            buffers,
544            SpatialMetadata::default(),
545        )
546        .unwrap();
547        let (mut schema, mut array) = export_point_cloud_c_data(&cloud).unwrap();
548        let imported = unsafe {
549            import_point_cloud_c_data(
550                schema.as_mut_ptr(),
551                array.as_mut_ptr(),
552                SpatialMetadata::default(),
553            )
554        }
555        .unwrap();
556        assert_eq!(imported.len(), 2);
557        assert_eq!(imported.field("x").unwrap().as_f32().unwrap(), &[1.0, 2.0]);
558        assert_eq!(imported.field("z").unwrap().as_f32().unwrap(), &[5.0, 6.0]);
559    }
560}
spatialrust_arrow/cdata.rs

spatialrust_arrow/
cdata.rs
