datafusion_ffi/
arrow_wrappers.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use abi_stable::StableAbi;
21use arrow::{
22    array::{make_array, ArrayRef},
23    datatypes::{Schema, SchemaRef},
24    error::ArrowError,
25    ffi::{from_ffi, to_ffi, FFI_ArrowArray, FFI_ArrowSchema},
26};
27use log::error;
28
29/// This is a wrapper struct around FFI_ArrowSchema simply to indicate
30/// to the StableAbi macros that the underlying struct is FFI safe.
31#[repr(C)]
32#[derive(Debug, StableAbi)]
33pub struct WrappedSchema(#[sabi(unsafe_opaque_field)] pub FFI_ArrowSchema);
34
35impl From<SchemaRef> for WrappedSchema {
36    fn from(value: SchemaRef) -> Self {
37        let ffi_schema = match FFI_ArrowSchema::try_from(value.as_ref()) {
38            Ok(s) => s,
39            Err(e) => {
40                error!("Unable to convert DataFusion Schema to FFI_ArrowSchema in FFI_PlanProperties. {e}");
41                FFI_ArrowSchema::empty()
42            }
43        };
44
45        WrappedSchema(ffi_schema)
46    }
47}
48/// Some functions are expected to always succeed, like getting the schema from a TableProvider.
49/// Since going through the FFI always has the potential to fail, we need to catch these errors,
50/// give the user a warning, and return some kind of result. In this case we default to an
51/// empty schema.
52#[cfg(not(tarpaulin_include))]
53fn catch_df_schema_error(e: ArrowError) -> Schema {
54    error!("Unable to convert from FFI_ArrowSchema to DataFusion Schema in FFI_PlanProperties. {e}");
55    Schema::empty()
56}
57
58impl From<WrappedSchema> for SchemaRef {
59    fn from(value: WrappedSchema) -> Self {
60        let schema = Schema::try_from(&value.0).unwrap_or_else(catch_df_schema_error);
61        Arc::new(schema)
62    }
63}
64
65/// This is a wrapper struct for FFI_ArrowArray to indicate to StableAbi
66/// that the struct is FFI Safe. For convenience, we also include the
67/// schema needed to create a record batch from the array.
68#[repr(C)]
69#[derive(Debug, StableAbi)]
70pub struct WrappedArray {
71    #[sabi(unsafe_opaque_field)]
72    pub array: FFI_ArrowArray,
73
74    pub schema: WrappedSchema,
75}
76
77impl TryFrom<WrappedArray> for ArrayRef {
78    type Error = ArrowError;
79
80    fn try_from(value: WrappedArray) -> Result<Self, Self::Error> {
81        let data = unsafe { from_ffi(value.array, &value.schema.0)? };
82
83        Ok(make_array(data))
84    }
85}
86
87impl TryFrom<&ArrayRef> for WrappedArray {
88    type Error = ArrowError;
89
90    fn try_from(array: &ArrayRef) -> Result<Self, Self::Error> {
91        let (array, schema) = to_ffi(&array.to_data())?;
92        let schema = WrappedSchema(schema);
93
94        Ok(WrappedArray { array, schema })
95    }
96}