|
| 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 | + |
| 18 | +use abi_stable::StableAbi; |
| 19 | +use datafusion_common::{DataFusionError, ScalarValue}; |
| 20 | +use datafusion_expr::ColumnarValue; |
| 21 | + |
| 22 | +use crate::arrow_wrappers::WrappedArray; |
| 23 | + |
| 24 | +/// A stable struct for sharing [`ColumnarValue`] across FFI boundaries. |
| 25 | +/// Scalar values are passed as an Arrow array of length 1. |
| 26 | +#[repr(C)] |
| 27 | +#[derive(Debug, StableAbi)] |
| 28 | +#[allow(non_camel_case_types)] |
| 29 | +pub enum FFI_ColumnarValue { |
| 30 | + Array(WrappedArray), |
| 31 | + Scalar(WrappedArray), |
| 32 | +} |
| 33 | + |
| 34 | +impl TryFrom<ColumnarValue> for FFI_ColumnarValue { |
| 35 | + type Error = DataFusionError; |
| 36 | + fn try_from(value: ColumnarValue) -> Result<Self, Self::Error> { |
| 37 | + Ok(match value { |
| 38 | + ColumnarValue::Array(v) => { |
| 39 | + FFI_ColumnarValue::Array(WrappedArray::try_from(&v)?) |
| 40 | + } |
| 41 | + ColumnarValue::Scalar(v) => { |
| 42 | + FFI_ColumnarValue::Scalar(WrappedArray::try_from(&v)?) |
| 43 | + } |
| 44 | + }) |
| 45 | + } |
| 46 | +} |
| 47 | + |
| 48 | +impl TryFrom<FFI_ColumnarValue> for ColumnarValue { |
| 49 | + type Error = DataFusionError; |
| 50 | + fn try_from(value: FFI_ColumnarValue) -> Result<Self, Self::Error> { |
| 51 | + Ok(match value { |
| 52 | + FFI_ColumnarValue::Array(v) => ColumnarValue::Array(v.try_into()?), |
| 53 | + FFI_ColumnarValue::Scalar(v) => { |
| 54 | + ColumnarValue::Scalar(ScalarValue::try_from(v)?) |
| 55 | + } |
| 56 | + }) |
| 57 | + } |
| 58 | +} |
| 59 | + |
| 60 | +#[cfg(test)] |
| 61 | +mod tests { |
| 62 | + use arrow::array::create_array; |
| 63 | + use datafusion_common::{DataFusionError, ScalarValue}; |
| 64 | + use datafusion_expr::ColumnarValue; |
| 65 | + |
| 66 | + use crate::expr::columnar_value::FFI_ColumnarValue; |
| 67 | + |
| 68 | + #[test] |
| 69 | + fn ffi_columnar_value_round_trip() -> Result<(), DataFusionError> { |
| 70 | + let array = create_array!(Int32, [1, 2, 3, 4, 5]); |
| 71 | + |
| 72 | + for original in [ |
| 73 | + ColumnarValue::Array(array), |
| 74 | + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), |
| 75 | + ] { |
| 76 | + let ffi_variant = FFI_ColumnarValue::try_from(original.clone())?; |
| 77 | + |
| 78 | + let returned_value = ColumnarValue::try_from(ffi_variant)?; |
| 79 | + |
| 80 | + assert_eq!(format!("{returned_value:?}"), format!("{original:?}")); |
| 81 | + } |
| 82 | + |
| 83 | + Ok(()) |
| 84 | + } |
| 85 | +} |
0 commit comments