use std::sync::Arc;
use abi_stable::StableAbi;
use arrow::array::{ArrayRef, make_array};
use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::ArrowError;
use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema, from_ffi, to_ffi};
use datafusion_common::{DataFusionError, ScalarValue};
use log::error;
#[repr(C)]
#[derive(Debug, StableAbi)]
pub struct WrappedSchema(#[sabi(unsafe_opaque_field)] pub FFI_ArrowSchema);
impl From<SchemaRef> for WrappedSchema {
fn from(value: SchemaRef) -> Self {
let ffi_schema = match FFI_ArrowSchema::try_from(value.as_ref()) {
Ok(s) => s,
Err(e) => {
error!(
"Unable to convert DataFusion Schema to FFI_ArrowSchema in FFI_PlanProperties. {e}"
);
FFI_ArrowSchema::empty()
}
};
WrappedSchema(ffi_schema)
}
}
#[cfg(not(tarpaulin_include))]
fn catch_df_schema_error(e: &ArrowError) -> Schema {
error!(
"Unable to convert from FFI_ArrowSchema to DataFusion Schema in FFI_PlanProperties. {e}"
);
Schema::empty()
}
impl From<WrappedSchema> for SchemaRef {
fn from(value: WrappedSchema) -> Self {
let schema =
Schema::try_from(&value.0).unwrap_or_else(|e| catch_df_schema_error(&e));
Arc::new(schema)
}
}
#[repr(C)]
#[derive(Debug, StableAbi)]
pub struct WrappedArray {
#[sabi(unsafe_opaque_field)]
pub array: FFI_ArrowArray,
pub schema: WrappedSchema,
}
impl TryFrom<WrappedArray> for ArrayRef {
type Error = ArrowError;
fn try_from(value: WrappedArray) -> Result<Self, Self::Error> {
let data = unsafe { from_ffi(value.array, &value.schema.0)? };
Ok(make_array(data))
}
}
impl TryFrom<&ArrayRef> for WrappedArray {
type Error = ArrowError;
fn try_from(array: &ArrayRef) -> Result<Self, Self::Error> {
let (array, schema) = to_ffi(&array.to_data())?;
let schema = WrappedSchema(schema);
Ok(WrappedArray { array, schema })
}
}
impl TryFrom<&ScalarValue> for WrappedArray {
type Error = DataFusionError;
fn try_from(value: &ScalarValue) -> Result<Self, Self::Error> {
let array = value.to_array()?;
WrappedArray::try_from(&array).map_err(Into::into)
}
}
impl TryFrom<WrappedArray> for ScalarValue {
type Error = DataFusionError;
fn try_from(value: WrappedArray) -> Result<Self, Self::Error> {
let array: ArrayRef = value.try_into()?;
ScalarValue::try_from_array(array.as_ref(), 0)
}
}