1use 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
11pub 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
37pub 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}