spatialrust_arrow/
stream.rs

1//! Arrow C Stream Interface over SpatialRust record sources.
2
3use std::{
4    ffi::{c_char, c_void, CString},
5    ptr,
6};
7
8use spatialrust_records::SpatialRecordSource;
9
10use crate::{
11    cdata::{export_point_cloud_c_data, ArrowArray, ArrowSchema},
12    ArrowBridgeResult,
13};
14
15/// Arrow C Stream Interface object.
16#[repr(C)]
17pub struct ArrowArrayStream {
18    /// Fills `out` with the stream schema.
19    pub get_schema:
20        Option<unsafe extern "C" fn(stream: *mut ArrowArrayStream, out: *mut ArrowSchema) -> i32>,
21    /// Fills `out` with the next array (`release=null` when exhausted).
22    pub get_next:
23        Option<unsafe extern "C" fn(stream: *mut ArrowArrayStream, out: *mut ArrowArray) -> i32>,
24    /// Optional last-error message.
25    pub get_last_error:
26        Option<unsafe extern "C" fn(stream: *mut ArrowArrayStream) -> *const c_char>,
27    /// Release callback.
28    pub release: Option<unsafe extern "C" fn(stream: *mut ArrowArrayStream)>,
29    /// Implementation private data.
30    pub private_data: *mut c_void,
31}
32
33/// Owned Arrow C Stream export.
34pub struct ExportedArrowArrayStream {
35    raw: Box<ArrowArrayStream>,
36}
37
38impl ExportedArrowArrayStream {
39    /// Returns a mutable raw stream pointer for FFI handoff.
40    pub fn as_mut_ptr(&mut self) -> *mut ArrowArrayStream {
41        self.raw.as_mut()
42    }
43}
44
45impl Drop for ExportedArrowArrayStream {
46    fn drop(&mut self) {
47        if let Some(release) = self.raw.release {
48            unsafe { release(self.raw.as_mut()) };
49        }
50    }
51}
52
53struct StreamPrivate {
54    source: Box<dyn SpatialRecordSource + Send>,
55    last_error: Option<CString>,
56}
57
58/// Exports a [`SpatialRecordSource`] as an Arrow C Stream of point-cloud structs.
59pub fn export_record_source_c_stream(
60    source: Box<dyn SpatialRecordSource + Send>,
61) -> ArrowBridgeResult<ExportedArrowArrayStream> {
62    let private = Box::new(StreamPrivate { source, last_error: None });
63    let raw = Box::new(ArrowArrayStream {
64        get_schema: Some(stream_get_schema),
65        get_next: Some(stream_get_next),
66        get_last_error: Some(stream_get_last_error),
67        release: Some(stream_release),
68        private_data: Box::into_raw(private) as *mut c_void,
69    });
70    Ok(ExportedArrowArrayStream { raw })
71}
72
73unsafe extern "C" fn stream_get_schema(
74    stream: *mut ArrowArrayStream,
75    out: *mut ArrowSchema,
76) -> i32 {
77    if stream.is_null() || out.is_null() {
78        return EINVAL;
79    }
80    let private = match private_mut(stream) {
81        Ok(value) => value,
82        Err(code) => return code,
83    };
84    let cloud = match empty_cloud(private.source.schema().point_schema()) {
85        Ok(cloud) => cloud,
86        Err(error) => {
87            set_error(private, error.to_string());
88            return EIO;
89        }
90    };
91    match export_point_cloud_c_data(&cloud) {
92        Ok((mut exported_schema, exported_array)) => {
93            ptr::write(out, ptr::read(exported_schema.as_mut_ptr()));
94            unsafe {
95                (*exported_schema.as_mut_ptr()).release = None;
96                (*exported_schema.as_mut_ptr()).private_data = ptr::null_mut();
97                (*exported_schema.as_mut_ptr()).children = ptr::null_mut();
98            }
99            drop(exported_array);
100            0
101        }
102        Err(error) => {
103            set_error(private, error.to_string());
104            EIO
105        }
106    }
107}
108
109unsafe extern "C" fn stream_get_next(stream: *mut ArrowArrayStream, out: *mut ArrowArray) -> i32 {
110    if stream.is_null() || out.is_null() {
111        return EINVAL;
112    }
113    let private = match private_mut(stream) {
114        Ok(value) => value,
115        Err(code) => return code,
116    };
117    match private.source.next_record() {
118        None => {
119            ptr::write(out, null_array());
120            0
121        }
122        Some(Ok(record)) => match export_point_cloud_c_data(record.cloud()) {
123            Ok((_schema, mut array)) => {
124                ptr::write(out, ptr::read(array.as_mut_ptr()));
125                unsafe {
126                    (*array.as_mut_ptr()).release = None;
127                    (*array.as_mut_ptr()).private_data = ptr::null_mut();
128                    (*array.as_mut_ptr()).buffers = ptr::null_mut();
129                    (*array.as_mut_ptr()).children = ptr::null_mut();
130                }
131                0
132            }
133            Err(error) => {
134                set_error(private, error.to_string());
135                EIO
136            }
137        },
138        Some(Err(error)) => {
139            set_error(private, error.to_string());
140            EIO
141        }
142    }
143}
144
145unsafe extern "C" fn stream_get_last_error(stream: *mut ArrowArrayStream) -> *const c_char {
146    if stream.is_null() {
147        return ptr::null();
148    }
149    match private_mut(stream) {
150        Ok(private) => {
151            private.last_error.as_ref().map(|value| value.as_ptr()).unwrap_or(ptr::null())
152        }
153        Err(_) => ptr::null(),
154    }
155}
156
157unsafe extern "C" fn stream_release(stream: *mut ArrowArrayStream) {
158    if stream.is_null() {
159        return;
160    }
161    let stream = &mut *stream;
162    if stream.release.is_none() {
163        return;
164    }
165    if !stream.private_data.is_null() {
166        drop(Box::from_raw(stream.private_data as *mut StreamPrivate));
167    }
168    stream.get_schema = None;
169    stream.get_next = None;
170    stream.get_last_error = None;
171    stream.release = None;
172    stream.private_data = ptr::null_mut();
173}
174
175unsafe fn private_mut(stream: *mut ArrowArrayStream) -> Result<&'static mut StreamPrivate, i32> {
176    let stream = &mut *stream;
177    if stream.private_data.is_null() {
178        return Err(EINVAL);
179    }
180    Ok(&mut *(stream.private_data as *mut StreamPrivate))
181}
182
183fn empty_cloud(
184    schema: &spatialrust_core::PointSchema,
185) -> ArrowBridgeResult<spatialrust_core::PointCloud> {
186    use spatialrust_core::{PointBuffer, PointBufferSet, PointCloud, SpatialMetadata};
187    let mut buffers = PointBufferSet::new();
188    for field in schema.fields() {
189        buffers.insert(field.name.clone(), PointBuffer::with_capacity(field.dtype, 0));
190    }
191    Ok(PointCloud::try_from_parts(schema.clone(), buffers, SpatialMetadata::default())?)
192}
193
194fn set_error(private: &mut StreamPrivate, message: String) {
195    private.last_error = CString::new(message).ok();
196}
197
198fn null_array() -> ArrowArray {
199    ArrowArray {
200        length: 0,
201        null_count: 0,
202        offset: 0,
203        n_buffers: 0,
204        n_children: 0,
205        buffers: ptr::null_mut(),
206        children: ptr::null_mut(),
207        dictionary: ptr::null_mut(),
208        release: None,
209        private_data: ptr::null_mut(),
210    }
211}
212
213const EINVAL: i32 = 22;
214const EIO: i32 = 5;
215
216#[cfg(test)]
217mod tests {
218    use super::export_record_source_c_stream;
219    use crate::cdata::import_point_cloud_c_data;
220    use spatialrust_core::{
221        PointBuffer, PointBufferSet, PointCloud, SpatialMetadata, StandardSchemas,
222    };
223    use spatialrust_records::{MemoryChunkSource, SchemaDescriptor, SchemaVersion};
224
225    #[test]
226    fn stream_yields_chunked_clouds() {
227        let mut buffers = PointBufferSet::new();
228        buffers.insert("x", PointBuffer::from_f32(vec![0.0, 1.0, 2.0]));
229        buffers.insert("y", PointBuffer::from_f32(vec![0.0; 3]));
230        buffers.insert("z", PointBuffer::from_f32(vec![1.0; 3]));
231        let cloud = PointCloud::try_from_parts(
232            StandardSchemas::point_xyz(),
233            buffers,
234            SpatialMetadata::default(),
235        )
236        .unwrap();
237        let schema =
238            SchemaDescriptor::try_new("point", SchemaVersion::new(1, 0), cloud.schema().clone())
239                .unwrap();
240        let source = MemoryChunkSource::try_new(schema, cloud, 2).unwrap();
241        let mut stream = export_record_source_c_stream(Box::new(source)).unwrap();
242        let stream_ptr = stream.as_mut_ptr();
243        unsafe {
244            let get_next = (*stream_ptr).get_next.expect("get_next");
245            let mut first = std::mem::zeroed();
246            assert_eq!(get_next(stream_ptr, &mut first), 0);
247            assert!(first.release.is_some());
248            assert_eq!(first.length, 2);
249            let mut second = std::mem::zeroed();
250            assert_eq!(get_next(stream_ptr, &mut second), 0);
251            assert_eq!(second.length, 1);
252            let mut done = std::mem::zeroed();
253            assert_eq!(get_next(stream_ptr, &mut done), 0);
254            assert!(done.release.is_none());
255            if let Some(release) = first.release {
256                release(&mut first);
257            }
258            if let Some(release) = second.release {
259                release(&mut second);
260            }
261        }
262        let _ = import_point_cloud_c_data;
263    }
264}
spatialrust_arrow/stream.rs

spatialrust_arrow/
stream.rs
