datafusion_ffi/
util.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 crate::arrow_wrappers::WrappedSchema;
19use abi_stable::std_types::RVec;
20use arrow::datatypes::Field;
21use arrow::{datatypes::DataType, ffi::FFI_ArrowSchema};
22use arrow_schema::FieldRef;
23use std::sync::Arc;
24
25/// This macro is a helpful conversion utility to convert from an abi_stable::RResult to a
26/// DataFusion result.
27#[macro_export]
28macro_rules! df_result {
29    ( $x:expr ) => {
30        match $x {
31            abi_stable::std_types::RResult::ROk(v) => Ok(v),
32            abi_stable::std_types::RResult::RErr(e) => {
33                datafusion_common::exec_err!("FFI error: {}", e)
34            }
35        }
36    };
37}
38
39/// This macro is a helpful conversion utility to convert from a DataFusion Result to an abi_stable::RResult
40#[macro_export]
41macro_rules! rresult {
42    ( $x:expr ) => {
43        match $x {
44            Ok(v) => abi_stable::std_types::RResult::ROk(v),
45            Err(e) => abi_stable::std_types::RResult::RErr(
46                abi_stable::std_types::RString::from(e.to_string()),
47            ),
48        }
49    };
50}
51
52/// This macro is a helpful conversion utility to convert from a DataFusion Result to an abi_stable::RResult
53/// and to also call return when it is an error. Since you cannot use `?` on an RResult, this is designed
54/// to mimic the pattern.
55#[macro_export]
56macro_rules! rresult_return {
57    ( $x:expr ) => {
58        match $x {
59            Ok(v) => v,
60            Err(e) => {
61                return abi_stable::std_types::RResult::RErr(
62                    abi_stable::std_types::RString::from(e.to_string()),
63                )
64            }
65        }
66    };
67}
68
69/// This is a utility function to convert a slice of [`Field`] to its equivalent
70/// FFI friendly counterpart, [`WrappedSchema`]
71pub fn vec_fieldref_to_rvec_wrapped(
72    fields: &[FieldRef],
73) -> Result<RVec<WrappedSchema>, arrow::error::ArrowError> {
74    Ok(fields
75        .iter()
76        .map(FFI_ArrowSchema::try_from)
77        .collect::<Result<Vec<_>, arrow::error::ArrowError>>()?
78        .into_iter()
79        .map(WrappedSchema)
80        .collect())
81}
82
83/// This is a utility function to convert an FFI friendly vector of [`WrappedSchema`]
84/// to their equivalent [`Field`].
85pub fn rvec_wrapped_to_vec_fieldref(
86    fields: &RVec<WrappedSchema>,
87) -> Result<Vec<FieldRef>, arrow::error::ArrowError> {
88    fields
89        .iter()
90        .map(|d| Field::try_from(&d.0).map(Arc::new))
91        .collect()
92}
93
94/// This is a utility function to convert a slice of [`DataType`] to its equivalent
95/// FFI friendly counterpart, [`WrappedSchema`]
96pub fn vec_datatype_to_rvec_wrapped(
97    data_types: &[DataType],
98) -> Result<RVec<WrappedSchema>, arrow::error::ArrowError> {
99    Ok(data_types
100        .iter()
101        .map(FFI_ArrowSchema::try_from)
102        .collect::<Result<Vec<_>, arrow::error::ArrowError>>()?
103        .into_iter()
104        .map(WrappedSchema)
105        .collect())
106}
107
108/// This is a utility function to convert an FFI friendly vector of [`WrappedSchema`]
109/// to their equivalent [`DataType`].
110pub fn rvec_wrapped_to_vec_datatype(
111    data_types: &RVec<WrappedSchema>,
112) -> Result<Vec<DataType>, arrow::error::ArrowError> {
113    data_types
114        .iter()
115        .map(|d| DataType::try_from(&d.0))
116        .collect()
117}
118
119#[cfg(test)]
120mod tests {
121    use abi_stable::std_types::{RResult, RString};
122    use datafusion::error::DataFusionError;
123
124    fn wrap_result(result: Result<String, DataFusionError>) -> RResult<String, RString> {
125        RResult::ROk(rresult_return!(result))
126    }
127
128    #[test]
129    fn test_conversion() {
130        const VALID_VALUE: &str = "valid_value";
131        const ERROR_VALUE: &str = "error_value";
132
133        let ok_r_result: RResult<RString, RString> =
134            RResult::ROk(VALID_VALUE.to_string().into());
135        let err_r_result: RResult<RString, RString> =
136            RResult::RErr(ERROR_VALUE.to_string().into());
137
138        let returned_ok_result = df_result!(ok_r_result);
139        assert!(returned_ok_result.is_ok());
140        assert!(returned_ok_result.unwrap().to_string() == VALID_VALUE);
141
142        let returned_err_result = df_result!(err_r_result);
143        assert!(returned_err_result.is_err());
144        assert!(
145            returned_err_result.unwrap_err().strip_backtrace()
146                == format!("Execution error: FFI error: {ERROR_VALUE}")
147        );
148
149        let ok_result: Result<String, DataFusionError> = Ok(VALID_VALUE.to_string());
150        let err_result: Result<String, DataFusionError> =
151            datafusion_common::exec_err!("{ERROR_VALUE}");
152
153        let returned_ok_r_result = wrap_result(ok_result);
154        assert!(returned_ok_r_result == RResult::ROk(VALID_VALUE.into()));
155
156        let returned_err_r_result = wrap_result(err_result);
157        assert!(returned_err_r_result.is_err());
158        assert!(returned_err_r_result
159            .unwrap_err()
160            .starts_with(format!("Execution error: {ERROR_VALUE}").as_str()));
161    }
162}