datafusion_ffi/expr/
distribution.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::StableAbi;
19use datafusion_common::DataFusionError;
20use datafusion_expr::statistics::{
21    BernoulliDistribution, Distribution, ExponentialDistribution, GaussianDistribution,
22    GenericDistribution, UniformDistribution,
23};
24
25use crate::arrow_wrappers::WrappedArray;
26use crate::expr::interval::FFI_Interval;
27
28/// A stable struct for sharing [`Distribution`] across FFI boundaries.
29/// See ['Distribution'] for the meaning of each variant.
30#[repr(C)]
31#[derive(Debug, StableAbi)]
32#[expect(clippy::large_enum_variant)]
33pub enum FFI_Distribution {
34    Uniform(FFI_UniformDistribution),
35    Exponential(FFI_ExponentialDistribution),
36    Gaussian(FFI_GaussianDistribution),
37    Bernoulli(FFI_BernoulliDistribution),
38    Generic(FFI_GenericDistribution),
39}
40
41impl TryFrom<&Distribution> for FFI_Distribution {
42    type Error = DataFusionError;
43    fn try_from(value: &Distribution) -> Result<Self, Self::Error> {
44        match value {
45            Distribution::Uniform(d) => Ok(FFI_Distribution::Uniform(d.try_into()?)),
46            Distribution::Exponential(d) => {
47                Ok(FFI_Distribution::Exponential(d.try_into()?))
48            }
49            Distribution::Gaussian(d) => Ok(FFI_Distribution::Gaussian(d.try_into()?)),
50            Distribution::Bernoulli(d) => Ok(FFI_Distribution::Bernoulli(d.try_into()?)),
51            Distribution::Generic(d) => Ok(FFI_Distribution::Generic(d.try_into()?)),
52        }
53    }
54}
55
56impl TryFrom<FFI_Distribution> for Distribution {
57    type Error = DataFusionError;
58    fn try_from(value: FFI_Distribution) -> Result<Self, Self::Error> {
59        match value {
60            FFI_Distribution::Uniform(d) => d.try_into(),
61            FFI_Distribution::Exponential(d) => d.try_into(),
62            FFI_Distribution::Gaussian(d) => d.try_into(),
63            FFI_Distribution::Bernoulli(d) => d.try_into(),
64            FFI_Distribution::Generic(d) => d.try_into(),
65        }
66    }
67}
68
69#[repr(C)]
70#[derive(Debug, StableAbi)]
71pub struct FFI_UniformDistribution {
72    interval: FFI_Interval,
73}
74
75#[repr(C)]
76#[derive(Debug, StableAbi)]
77pub struct FFI_ExponentialDistribution {
78    rate: WrappedArray,
79    offset: WrappedArray,
80    positive_tail: bool,
81}
82
83#[repr(C)]
84#[derive(Debug, StableAbi)]
85pub struct FFI_GaussianDistribution {
86    mean: WrappedArray,
87    variance: WrappedArray,
88}
89
90#[repr(C)]
91#[derive(Debug, StableAbi)]
92pub struct FFI_BernoulliDistribution {
93    p: WrappedArray,
94}
95
96#[repr(C)]
97#[derive(Debug, StableAbi)]
98pub struct FFI_GenericDistribution {
99    mean: WrappedArray,
100    median: WrappedArray,
101    variance: WrappedArray,
102    range: FFI_Interval,
103}
104
105impl TryFrom<&UniformDistribution> for FFI_UniformDistribution {
106    type Error = DataFusionError;
107    fn try_from(value: &UniformDistribution) -> Result<Self, Self::Error> {
108        Ok(Self {
109            interval: value.range().try_into()?,
110        })
111    }
112}
113
114impl TryFrom<&ExponentialDistribution> for FFI_ExponentialDistribution {
115    type Error = DataFusionError;
116    fn try_from(value: &ExponentialDistribution) -> Result<Self, Self::Error> {
117        let rate = value.rate().try_into()?;
118        let offset = value.offset().try_into()?;
119
120        Ok(Self {
121            rate,
122            offset,
123            positive_tail: value.positive_tail(),
124        })
125    }
126}
127
128impl TryFrom<&GaussianDistribution> for FFI_GaussianDistribution {
129    type Error = DataFusionError;
130    fn try_from(value: &GaussianDistribution) -> Result<Self, Self::Error> {
131        let mean = value.mean().try_into()?;
132        let variance = value.variance().try_into()?;
133
134        Ok(Self { mean, variance })
135    }
136}
137
138impl TryFrom<&BernoulliDistribution> for FFI_BernoulliDistribution {
139    type Error = DataFusionError;
140    fn try_from(value: &BernoulliDistribution) -> Result<Self, Self::Error> {
141        let p = value.p_value().try_into()?;
142
143        Ok(Self { p })
144    }
145}
146
147impl TryFrom<&GenericDistribution> for FFI_GenericDistribution {
148    type Error = DataFusionError;
149    fn try_from(value: &GenericDistribution) -> Result<Self, Self::Error> {
150        let mean = value.mean().try_into()?;
151        let median = value.median().try_into()?;
152        let variance = value.variance().try_into()?;
153
154        Ok(Self {
155            mean,
156            median,
157            variance,
158            range: value.range().try_into()?,
159        })
160    }
161}
162
163impl TryFrom<FFI_UniformDistribution> for Distribution {
164    type Error = DataFusionError;
165    fn try_from(value: FFI_UniformDistribution) -> Result<Self, Self::Error> {
166        let interval = value.interval.try_into()?;
167        Distribution::new_uniform(interval)
168    }
169}
170
171impl TryFrom<FFI_ExponentialDistribution> for Distribution {
172    type Error = DataFusionError;
173    fn try_from(value: FFI_ExponentialDistribution) -> Result<Self, Self::Error> {
174        let rate = value.rate.try_into()?;
175        let offset = value.offset.try_into()?;
176
177        Distribution::new_exponential(rate, offset, value.positive_tail)
178    }
179}
180
181impl TryFrom<FFI_GaussianDistribution> for Distribution {
182    type Error = DataFusionError;
183    fn try_from(value: FFI_GaussianDistribution) -> Result<Self, Self::Error> {
184        let mean = value.mean.try_into()?;
185        let variance = value.variance.try_into()?;
186
187        Distribution::new_gaussian(mean, variance)
188    }
189}
190
191impl TryFrom<FFI_BernoulliDistribution> for Distribution {
192    type Error = DataFusionError;
193    fn try_from(value: FFI_BernoulliDistribution) -> Result<Self, Self::Error> {
194        let p = value.p.try_into()?;
195
196        Distribution::new_bernoulli(p)
197    }
198}
199
200impl TryFrom<FFI_GenericDistribution> for Distribution {
201    type Error = DataFusionError;
202    fn try_from(value: FFI_GenericDistribution) -> Result<Self, Self::Error> {
203        let mean = value.mean.try_into()?;
204        let median = value.median.try_into()?;
205        let variance = value.variance.try_into()?;
206        let range = value.range.try_into()?;
207
208        Distribution::new_generic(mean, median, variance, range)
209    }
210}
datafusion_ffi/expr/distribution.rs

datafusion_ffi/expr/
distribution.rs
