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