laminar_db/ffi/
arrow_ffi.rs

1//! Arrow C Data Interface for zero-copy data exchange.
2
3use arrow::array::{Array, RecordBatch, StructArray};
4use arrow::ffi::{from_ffi, to_ffi, FFI_ArrowArray, FFI_ArrowSchema};
5
6use super::error::{
7    clear_last_error, set_last_error, LAMINAR_ERR_INTERNAL, LAMINAR_ERR_NULL_POINTER, LAMINAR_OK,
8};
9use super::query::LaminarRecordBatch;
10use super::schema::LaminarSchema;
11use crate::api::ApiError;
12
13/// Export a `RecordBatch` to the Arrow C Data Interface.
14///
15/// # Safety
16///
17/// `batch`, `out_array`, and `out_schema` must be valid non-null pointers;
18/// caller must call the release callbacks on both output structs.
19#[no_mangle]
20pub unsafe extern "C" fn laminar_batch_export(
21    batch: *mut LaminarRecordBatch,
22    out_array: *mut FFI_ArrowArray,
23    out_schema: *mut FFI_ArrowSchema,
24) -> i32 {
25    clear_last_error();
26
27    if batch.is_null() || out_array.is_null() || out_schema.is_null() {
28        return LAMINAR_ERR_NULL_POINTER;
29    }
30
31    let batch_ref = unsafe { &(*batch) };
32    let record_batch = batch_ref.inner();
33
34    let struct_array: StructArray = record_batch.clone().into();
35    let data = struct_array.into_data();
36
37    match to_ffi(&data) {
38        Ok((array, schema)) => {
39            unsafe {
40                std::ptr::write(out_array, array);
41                std::ptr::write(out_schema, schema);
42            }
43            LAMINAR_OK
44        }
45        Err(e) => {
46            set_last_error(ApiError::internal(format!("Arrow FFI export failed: {e}")));
47            LAMINAR_ERR_INTERNAL
48        }
49    }
50}
51
52/// Export just the schema to the Arrow C Data Interface.
53///
54/// # Safety
55///
56/// `schema` and `out_schema` must be valid non-null pointers; caller must call the release callback.
57#[no_mangle]
58pub unsafe extern "C" fn laminar_schema_export(
59    schema: *mut LaminarSchema,
60    out_schema: *mut FFI_ArrowSchema,
61) -> i32 {
62    clear_last_error();
63
64    if schema.is_null() || out_schema.is_null() {
65        return LAMINAR_ERR_NULL_POINTER;
66    }
67
68    let schema_ref = unsafe { (*schema).schema() };
69
70    match FFI_ArrowSchema::try_from(schema_ref.as_ref()) {
71        Ok(ffi_schema) => {
72            unsafe {
73                std::ptr::write(out_schema, ffi_schema);
74            }
75            LAMINAR_OK
76        }
77        Err(e) => {
78            set_last_error(ApiError::internal(format!(
79                "Arrow FFI schema export failed: {e}"
80            )));
81            LAMINAR_ERR_INTERNAL
82        }
83    }
84}
85
86/// Export a single column from a `RecordBatch` to the Arrow C Data Interface.
87///
88/// # Safety
89///
90/// `batch`, `out_array`, and `out_schema` must be valid non-null pointers;
91/// `column_index` must be in range; caller must call the release callbacks.
92#[no_mangle]
93pub unsafe extern "C" fn laminar_batch_export_column(
94    batch: *mut LaminarRecordBatch,
95    column_index: usize,
96    out_array: *mut FFI_ArrowArray,
97    out_schema: *mut FFI_ArrowSchema,
98) -> i32 {
99    clear_last_error();
100
101    if batch.is_null() || out_array.is_null() || out_schema.is_null() {
102        return LAMINAR_ERR_NULL_POINTER;
103    }
104
105    let batch_ref = unsafe { &(*batch) };
106    let record_batch = batch_ref.inner();
107
108    if column_index >= record_batch.num_columns() {
109        set_last_error(ApiError::internal(format!(
110            "Column index {column_index} out of bounds (batch has {} columns)",
111            record_batch.num_columns()
112        )));
113        return LAMINAR_ERR_NULL_POINTER;
114    }
115
116    let column = record_batch.column(column_index);
117    let data = column.to_data();
118
119    match to_ffi(&data) {
120        Ok((array, schema)) => {
121            unsafe {
122                std::ptr::write(out_array, array);
123                std::ptr::write(out_schema, schema);
124            }
125            LAMINAR_OK
126        }
127        Err(e) => {
128            set_last_error(ApiError::internal(format!(
129                "Arrow FFI column export failed: {e}"
130            )));
131            LAMINAR_ERR_INTERNAL
132        }
133    }
134}
135
136/// Import a `RecordBatch` from the Arrow C Data Interface, taking ownership of both structs.
137///
138/// # Safety
139///
140/// `array`, `schema`, and `out` must be valid non-null pointers;
141/// ownership of `array` and `schema` is transferred and they are released by this function.
142#[no_mangle]
143pub unsafe extern "C" fn laminar_batch_import(
144    array: *mut FFI_ArrowArray,
145    schema: *mut FFI_ArrowSchema,
146    out: *mut *mut LaminarRecordBatch,
147) -> i32 {
148    clear_last_error();
149
150    if array.is_null() || schema.is_null() || out.is_null() {
151        return LAMINAR_ERR_NULL_POINTER;
152    }
153
154    let ffi_array = unsafe { std::ptr::read(array) };
155    let ffi_schema = unsafe { std::ptr::read(schema) };
156
157    // Zero the originals to prevent double-free; callers must not use them after this.
158    unsafe {
159        std::ptr::write_bytes(array, 0, 1);
160        std::ptr::write_bytes(schema, 0, 1);
161    }
162
163    match from_ffi(ffi_array, &ffi_schema) {
164        Ok(data) => {
165            let struct_array = StructArray::from(data);
166            let batch = RecordBatch::from(struct_array);
167
168            let handle = Box::new(LaminarRecordBatch::new(batch));
169            unsafe { *out = Box::into_raw(handle) };
170            LAMINAR_OK
171        }
172        Err(e) => {
173            set_last_error(ApiError::internal(format!("Arrow FFI import failed: {e}")));
174            LAMINAR_ERR_INTERNAL
175        }
176    }
177}
178
179/// Alias for [`laminar_batch_import`] for use when creating data for writing.
180///
181/// # Safety
182///
183/// Same requirements as [`laminar_batch_import`].
184#[no_mangle]
185pub unsafe extern "C" fn laminar_batch_create(
186    array: *mut FFI_ArrowArray,
187    schema: *mut FFI_ArrowSchema,
188    out: *mut *mut LaminarRecordBatch,
189) -> i32 {
190    laminar_batch_import(array, schema, out)
191}
192
193#[cfg(test)]
194#[allow(clippy::borrow_as_ptr)]
195mod tests {
196    use super::*;
197    use arrow::array::{Int64Array, StringArray};
198    use arrow::datatypes::{DataType, Field, Schema};
199    use std::sync::Arc;
200
201    fn create_test_batch() -> RecordBatch {
202        let schema = Arc::new(Schema::new(vec![
203            Field::new("id", DataType::Int64, false),
204            Field::new("name", DataType::Utf8, true),
205        ]));
206
207        RecordBatch::try_new(
208            schema,
209            vec![
210                Arc::new(Int64Array::from(vec![1, 2, 3])),
211                Arc::new(StringArray::from(vec![Some("Alice"), Some("Bob"), None])),
212            ],
213        )
214        .unwrap()
215    }
216
217    #[test]
218    fn test_export_import_roundtrip() {
219        let batch = create_test_batch();
220        let mut ffi_batch = LaminarRecordBatch::new(batch.clone());
221
222        // Export
223        let mut out_array = FFI_ArrowArray::empty();
224        let mut out_schema = FFI_ArrowSchema::empty();
225
226        let rc = unsafe { laminar_batch_export(&mut ffi_batch, &mut out_array, &mut out_schema) };
227        assert_eq!(rc, LAMINAR_OK);
228
229        // Import
230        let mut imported: *mut LaminarRecordBatch = std::ptr::null_mut();
231        let rc = unsafe { laminar_batch_import(&mut out_array, &mut out_schema, &mut imported) };
232        assert_eq!(rc, LAMINAR_OK);
233        assert!(!imported.is_null());
234
235        // Verify data matches
236        let imported_batch = unsafe { (*imported).inner() };
237        assert_eq!(batch.num_rows(), imported_batch.num_rows());
238        assert_eq!(batch.num_columns(), imported_batch.num_columns());
239
240        // Clean up
241        unsafe {
242            super::super::query::laminar_batch_free(imported);
243        }
244    }
245
246    #[test]
247    fn test_export_column() {
248        let batch = create_test_batch();
249        let mut ffi_batch = LaminarRecordBatch::new(batch);
250
251        let mut out_array = FFI_ArrowArray::empty();
252        let mut out_schema = FFI_ArrowSchema::empty();
253
254        // Export first column
255        let rc = unsafe {
256            laminar_batch_export_column(&mut ffi_batch, 0, &mut out_array, &mut out_schema)
257        };
258        assert_eq!(rc, LAMINAR_OK);
259
260        // Import and verify
261        let data = unsafe { from_ffi(out_array, &out_schema) }.unwrap();
262        let array = Int64Array::from(data);
263        assert_eq!(array.len(), 3);
264        assert_eq!(array.value(0), 1);
265        assert_eq!(array.value(1), 2);
266        assert_eq!(array.value(2), 3);
267    }
268
269    #[test]
270    fn test_export_column_out_of_bounds() {
271        let batch = create_test_batch();
272        let mut ffi_batch = LaminarRecordBatch::new(batch);
273
274        let mut out_array = FFI_ArrowArray::empty();
275        let mut out_schema = FFI_ArrowSchema::empty();
276
277        // Try to export non-existent column
278        let rc = unsafe {
279            laminar_batch_export_column(&mut ffi_batch, 99, &mut out_array, &mut out_schema)
280        };
281        assert_eq!(rc, LAMINAR_ERR_NULL_POINTER);
282    }
283
284    #[test]
285    fn test_schema_export() {
286        let schema = Arc::new(Schema::new(vec![
287            Field::new("id", DataType::Int64, false),
288            Field::new("value", DataType::Float64, true),
289        ]));
290
291        let mut ffi_schema_handle = LaminarSchema::new(schema);
292        let mut out_schema = FFI_ArrowSchema::empty();
293
294        let rc = unsafe { laminar_schema_export(&mut ffi_schema_handle, &mut out_schema) };
295        assert_eq!(rc, LAMINAR_OK);
296
297        // The schema is released when dropped
298        drop(out_schema);
299    }
300
301    #[test]
302    fn test_null_pointer_checks() {
303        let mut out_array = FFI_ArrowArray::empty();
304        let mut out_schema = FFI_ArrowSchema::empty();
305        let mut out: *mut LaminarRecordBatch = std::ptr::null_mut();
306
307        // Export with null batch
308        let rc =
309            unsafe { laminar_batch_export(std::ptr::null_mut(), &mut out_array, &mut out_schema) };
310        assert_eq!(rc, LAMINAR_ERR_NULL_POINTER);
311
312        // Import with null array
313        let rc = unsafe { laminar_batch_import(std::ptr::null_mut(), &mut out_schema, &mut out) };
314        assert_eq!(rc, LAMINAR_ERR_NULL_POINTER);
315
316        // Schema export with null schema
317        let rc = unsafe { laminar_schema_export(std::ptr::null_mut(), &mut out_schema) };
318        assert_eq!(rc, LAMINAR_ERR_NULL_POINTER);
319    }
320}
laminar_db/ffi/arrow_ffi.rs

laminar_db/ffi/
arrow_ffi.rs
