Skip to main content

shape_runtime/
arrow_c.rs

1//! Arrow C Data Interface import helpers.
2//!
3//! These helpers provide a narrow bridge from raw Arrow C pointers to the
4//! runtime `DataTable` representation.
5
6use arrow_array::{RecordBatch, StructArray, ffi::FFI_ArrowArray, ffi::from_ffi};
7use arrow_schema::{Schema, ffi::FFI_ArrowSchema};
8use shape_value::DataTable;
9use std::sync::Arc;
10
11/// Build a [`DataTable`] from owned Arrow C interface structs.
12///
13/// Ownership of `schema` and `array` is transferred to this function.
14pub fn datatable_from_arrow_ffi(
15    schema: FFI_ArrowSchema,
16    array: FFI_ArrowArray,
17) -> Result<DataTable, String> {
18    let arrow_schema = Schema::try_from(&schema)
19        .map_err(|e| format!("failed to decode Arrow schema from C interface: {e}"))?;
20
21    let array_data = unsafe { from_ffi(array, &schema) }
22        .map_err(|e| format!("failed to decode Arrow array from C interface: {e}"))?;
23    let struct_array = StructArray::from(array_data);
24    let (_fields, columns, nulls) = struct_array.into_parts();
25    if nulls.is_some() {
26        return Err(
27            "row-level null mask on top-level Arrow struct is not supported for table import"
28                .to_string(),
29        );
30    }
31
32    let batch = RecordBatch::try_new(Arc::new(arrow_schema), columns)
33        .map_err(|e| format!("failed to build RecordBatch from Arrow C data: {e}"))?;
34    Ok(DataTable::new(batch))
35}
36
37/// Build a [`DataTable`] from raw pointers to Arrow C interface structs.
38///
39/// # Safety
40///
41/// - `schema_ptr` must point to a valid initialized `FFI_ArrowSchema`.
42/// - `array_ptr` must point to a valid initialized `FFI_ArrowArray`.
43/// - This function takes ownership of both values via `ptr::read`.
44pub unsafe fn datatable_from_arrow_c_ptrs(
45    schema_ptr: usize,
46    array_ptr: usize,
47) -> Result<DataTable, String> {
48    if schema_ptr == 0 {
49        return Err("schema_ptr must not be null".to_string());
50    }
51    if array_ptr == 0 {
52        return Err("array_ptr must not be null".to_string());
53    }
54
55    let schema = unsafe { std::ptr::read(schema_ptr as *const FFI_ArrowSchema) };
56    let array = unsafe { std::ptr::read(array_ptr as *const FFI_ArrowArray) };
57    datatable_from_arrow_ffi(schema, array)
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63    use arrow_array::{Array, ArrayRef, Float64Array, Int64Array, StructArray, ffi::to_ffi};
64    use arrow_schema::{DataType, Field, Fields};
65
66    #[test]
67    fn import_arrow_ffi_struct_to_datatable() {
68        let fields = Fields::from(vec![
69            Field::new("id", DataType::Int64, false),
70            Field::new("price", DataType::Float64, false),
71        ]);
72        let struct_arr = StructArray::new(
73            fields,
74            vec![
75                Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef,
76                Arc::new(Float64Array::from(vec![10.0, 11.5, 12.25])) as ArrayRef,
77            ],
78            None,
79        );
80
81        let (ffi_array, ffi_schema) = to_ffi(&struct_arr.to_data()).expect("to_ffi should work");
82        let dt = datatable_from_arrow_ffi(ffi_schema, ffi_array).expect("import should succeed");
83
84        assert_eq!(dt.row_count(), 3);
85        assert_eq!(dt.column_count(), 2);
86        assert_eq!(
87            dt.column_names(),
88            vec!["id".to_string(), "price".to_string()]
89        );
90    }
91}