antecedent_data/
arrow_ffi.rs

1//! Isolated Arrow C Data Interface FFI.
2//!
3//! All `unsafe` for Arrow CDI lives in this module. Safe wrappers validate
4//! types and lengths before exposing library-owned views.
5//!
6//! SPDX-License-Identifier: MIT OR Apache-2.0
7
8#![allow(unsafe_code)]
9
10use std::sync::Arc;
11
12use antecedent_core::VariableId;
13use arrow_array::ffi::{FFI_ArrowArray, from_ffi};
14use arrow_array::{Array, ArrayRef, Float64Array};
15use arrow_schema::ffi::FFI_ArrowSchema;
16
17use crate::buffer::{F64Buffer, ForeignF64Buffer};
18use crate::column::{Float64Column, OwnedColumn, ValidityBitmap};
19use crate::error::DataError;
20use crate::materialize::{MaterializationReason, materialization_diagnostic};
21
22/// Keeps an Arrow array alive for foreign buffer borrows.
23#[derive(Debug)]
24struct ArrayOwner(#[allow(dead_code)] ArrayRef);
25
26/// One column imported from the Arrow C Data Interface.
27pub struct ArrowCColumn {
28    /// Column name.
29    pub name: String,
30    /// Owned FFI array (moved into [`from_ffi`]).
31    pub array: FFI_ArrowArray,
32    /// Owned FFI schema.
33    pub schema: FFI_ArrowSchema,
34}
35
36impl ArrowCColumn {
37    /// Import this CDI column into an Arrow array (takes ownership of FFI structs).
38    ///
39    /// # Errors
40    ///
41    /// When the CDI pair is malformed or unsupported.
42    pub fn into_array(self) -> Result<ArrayRef, DataError> {
43        // SAFETY: ArrowCColumn is only constructed at the FFI boundary with a
44        // compliant CDI pair (see module docs).
45        let Self { array, schema, .. } = self;
46        unsafe { import_array(array, &schema) }
47    }
48}
49
50/// Import a single Arrow array via CDI.
51///
52/// # Safety
53///
54/// `array` and `schema` must form a valid Arrow C Data Interface pair produced
55/// by a compliant exporter. This function takes ownership and will release them.
56///
57/// # Errors
58///
59/// Returns [`DataError::InvalidArgument`] when the FFI pair cannot be decoded.
60pub unsafe fn import_array(
61    array: FFI_ArrowArray,
62    schema: &FFI_ArrowSchema,
63) -> Result<ArrayRef, DataError> {
64    // SAFETY: caller guarantees a valid CDI pair; from_ffi releases on drop/error.
65    let data = unsafe { from_ffi(array, schema) }.map_err(|e| DataError::InvalidArgument {
66        message: format!("Arrow CDI import failed: {e}"),
67    })?;
68    Ok(arrow_array::make_array(data))
69}
70
71/// Convert a float64 Arrow array into a library column, preferring zero-copy.
72pub(crate) fn float64_column_from_array(
73    id: VariableId,
74    array: ArrayRef,
75) -> Result<(OwnedColumn, u64, u64, antecedent_core::Diagnostic), DataError> {
76    let floats = array
77        .as_any()
78        .downcast_ref::<Float64Array>()
79        .ok_or(DataError::TypeMismatch { id, expected: "float64" })?;
80    let n = floats.len();
81    let (validity, validity_copied) = validity_from_arrow(floats)?;
82
83    let values_slice = floats.values().as_ref();
84    let aligned = values_slice.as_ptr() as usize % core::mem::align_of::<f64>() == 0;
85    let can_borrow = aligned && values_slice.len() == n;
86
87    if can_borrow {
88        let ptr = values_slice.as_ptr();
89        // SAFETY: `array` owns the buffer for the lifetime of ForeignF64Buffer.
90        let foreign = unsafe { ForeignF64Buffer::new(Arc::new(ArrayOwner(array)), ptr, n) };
91        // Ensure owner holds the array (read through for dead_code).
92        debug_assert!(!foreign.as_slice().is_empty() || n == 0);
93        let _ = foreign.len();
94        let values = F64Buffer::foreign(foreign);
95        let borrowed = values.nbytes();
96        let col = Float64Column::new(id, values, validity)?;
97        let diag = materialization_diagnostic(MaterializationReason::ExplicitCopy, validity_copied);
98        Ok((OwnedColumn::Float64(col), borrowed, validity_copied, diag))
99    } else {
100        let mut values = Vec::with_capacity(n);
101        for i in 0..n {
102            values.push(if floats.is_null(i) { 0.0 } else { floats.value(i) });
103        }
104        let value_bytes = (values.len() * core::mem::size_of::<f64>()) as u64;
105        let copied = value_bytes + validity_copied;
106        let col = Float64Column::new(id, F64Buffer::owned(Arc::from(values)), validity)?;
107        let diag =
108            materialization_diagnostic(MaterializationReason::ForeignBufferIncompatible, copied);
109        Ok((OwnedColumn::Float64(col), 0, copied, diag))
110    }
111}
112
113fn validity_from_arrow(floats: &Float64Array) -> Result<(ValidityBitmap, u64), DataError> {
114    let n = floats.len();
115    if floats.null_count() == 0 {
116        let v = ValidityBitmap::all_valid(n);
117        let bytes = n.div_ceil(8) as u64;
118        return Ok((v, bytes));
119    }
120    let mut validity_bytes = vec![0u8; n.div_ceil(8)];
121    for row in 0..n {
122        if !floats.is_null(row) {
123            validity_bytes[row / 8] |= 1 << (row % 8);
124        }
125    }
126    let copied = validity_bytes.len() as u64;
127    Ok((ValidityBitmap::from_bytes(validity_bytes, n)?, copied))
128}
129
130/// Re-export FFI types for the Python / foreign boundary only.
131pub use arrow_array::ffi::FFI_ArrowArray as FfiArrowArray;
132pub use arrow_schema::ffi::FFI_ArrowSchema as FfiArrowSchema;
antecedent_data/arrow_ffi.rs

antecedent_data/
arrow_ffi.rs
