arrow_buffer/
native.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::{i256, IntervalDayTime, IntervalMonthDayNano};
19use half::f16;
20
21mod private {
22    pub trait Sealed {}
23}
24
25/// Trait expressing a Rust type that has the same in-memory representation as
26/// Arrow.
27///
28/// This includes `i16`, `f32`, but excludes `bool` (which in arrow is
29/// represented in bits).
30///
31/// In little endian machines, types that implement [`ArrowNativeType`] can be
32/// memcopied to arrow buffers as is.
33///
34/// # Transmute Safety
35///
36/// A type T implementing this trait means that any arbitrary slice of bytes of length and
37/// alignment `size_of::<T>()` can be safely interpreted as a value of that type without
38/// being unsound, i.e. potentially resulting in undefined behaviour.
39///
40/// Note: in the case of floating point numbers this transmutation can result in a signalling
41/// NaN, which, whilst sound, can be unwieldy. In general, whilst it is perfectly sound to
42/// reinterpret bytes as different types using this trait, it is likely unwise. For more information
43/// see [f32::from_bits] and [f64::from_bits].
44///
45/// Note: `bool` is restricted to `0` or `1`, and so `bool: !ArrowNativeType`
46///
47/// # Sealed
48///
49/// Due to the above restrictions, this trait is sealed to prevent accidental misuse
50pub trait ArrowNativeType:
51    std::fmt::Debug + Send + Sync + Copy + PartialOrd + Default + private::Sealed + 'static
52{
53    /// Returns the byte width of this native type.
54    fn get_byte_width() -> usize {
55        std::mem::size_of::<Self>()
56    }
57
58    /// Convert native integer type from usize
59    ///
60    /// Returns `None` if [`Self`] is not an integer or conversion would result
61    /// in truncation/overflow
62    fn from_usize(_: usize) -> Option<Self>;
63
64    /// Convert to usize according to the [`as`] operator
65    ///
66    /// [`as`]: https://doc.rust-lang.org/reference/expressions/operator-expr.html#numeric-cast
67    fn as_usize(self) -> usize;
68
69    /// Convert from usize according to the [`as`] operator
70    ///
71    /// [`as`]: https://doc.rust-lang.org/reference/expressions/operator-expr.html#numeric-cast
72    fn usize_as(i: usize) -> Self;
73
74    /// Convert native type to usize.
75    ///
76    /// Returns `None` if [`Self`] is not an integer or conversion would result
77    /// in truncation/overflow
78    fn to_usize(self) -> Option<usize>;
79
80    /// Convert native type to isize.
81    ///
82    /// Returns `None` if [`Self`] is not an integer or conversion would result
83    /// in truncation/overflow
84    fn to_isize(self) -> Option<isize>;
85
86    /// Convert native type to i64.
87    ///
88    /// Returns `None` if [`Self`] is not an integer or conversion would result
89    /// in truncation/overflow
90    fn to_i64(self) -> Option<i64>;
91}
92
93macro_rules! native_integer {
94    ($t: ty $(, $from:ident)*) => {
95        impl private::Sealed for $t {}
96        impl ArrowNativeType for $t {
97            #[inline]
98            fn from_usize(v: usize) -> Option<Self> {
99                v.try_into().ok()
100            }
101
102            #[inline]
103            fn to_usize(self) -> Option<usize> {
104                self.try_into().ok()
105            }
106
107            #[inline]
108            fn to_isize(self) -> Option<isize> {
109                self.try_into().ok()
110            }
111
112            #[inline]
113            fn to_i64(self) -> Option<i64> {
114                self.try_into().ok()
115            }
116
117            #[inline]
118            fn as_usize(self) -> usize {
119                self as _
120            }
121
122            #[inline]
123            fn usize_as(i: usize) -> Self {
124                i as _
125            }
126        }
127    };
128}
129
130native_integer!(i8);
131native_integer!(i16);
132native_integer!(i32);
133native_integer!(i64);
134native_integer!(i128);
135native_integer!(u8);
136native_integer!(u16);
137native_integer!(u32);
138native_integer!(u64);
139native_integer!(u128);
140
141macro_rules! native_float {
142    ($t:ty, $s:ident, $as_usize: expr, $i:ident, $usize_as: expr) => {
143        impl private::Sealed for $t {}
144        impl ArrowNativeType for $t {
145            #[inline]
146            fn from_usize(_: usize) -> Option<Self> {
147                None
148            }
149
150            #[inline]
151            fn to_usize(self) -> Option<usize> {
152                None
153            }
154
155            #[inline]
156            fn to_isize(self) -> Option<isize> {
157                None
158            }
159
160            #[inline]
161            fn to_i64(self) -> Option<i64> {
162                None
163            }
164
165            #[inline]
166            fn as_usize($s) -> usize {
167                $as_usize
168            }
169
170            #[inline]
171            fn usize_as($i: usize) -> Self {
172                $usize_as
173            }
174        }
175    };
176}
177
178native_float!(f16, self, self.to_f32() as _, i, f16::from_f32(i as _));
179native_float!(f32, self, self as _, i, i as _);
180native_float!(f64, self, self as _, i, i as _);
181
182impl private::Sealed for i256 {}
183impl ArrowNativeType for i256 {
184    fn from_usize(u: usize) -> Option<Self> {
185        Some(Self::from_parts(u as u128, 0))
186    }
187
188    fn as_usize(self) -> usize {
189        self.to_parts().0 as usize
190    }
191
192    fn usize_as(i: usize) -> Self {
193        Self::from_parts(i as u128, 0)
194    }
195
196    fn to_usize(self) -> Option<usize> {
197        let (low, high) = self.to_parts();
198        if high != 0 {
199            return None;
200        }
201        low.try_into().ok()
202    }
203
204    fn to_isize(self) -> Option<isize> {
205        self.to_i128()?.try_into().ok()
206    }
207
208    fn to_i64(self) -> Option<i64> {
209        self.to_i128()?.try_into().ok()
210    }
211}
212
213impl private::Sealed for IntervalMonthDayNano {}
214impl ArrowNativeType for IntervalMonthDayNano {
215    fn from_usize(_: usize) -> Option<Self> {
216        None
217    }
218
219    fn as_usize(self) -> usize {
220        ((self.months as u64) | ((self.days as u64) << 32)) as usize
221    }
222
223    fn usize_as(i: usize) -> Self {
224        Self::new(i as _, ((i as u64) >> 32) as _, 0)
225    }
226
227    fn to_usize(self) -> Option<usize> {
228        None
229    }
230
231    fn to_isize(self) -> Option<isize> {
232        None
233    }
234
235    fn to_i64(self) -> Option<i64> {
236        None
237    }
238}
239
240impl private::Sealed for IntervalDayTime {}
241impl ArrowNativeType for IntervalDayTime {
242    fn from_usize(_: usize) -> Option<Self> {
243        None
244    }
245
246    fn as_usize(self) -> usize {
247        ((self.days as u64) | ((self.milliseconds as u64) << 32)) as usize
248    }
249
250    fn usize_as(i: usize) -> Self {
251        Self::new(i as _, ((i as u64) >> 32) as _)
252    }
253
254    fn to_usize(self) -> Option<usize> {
255        None
256    }
257
258    fn to_isize(self) -> Option<isize> {
259        None
260    }
261
262    fn to_i64(self) -> Option<i64> {
263        None
264    }
265}
266
267/// Allows conversion from supported Arrow types to a byte slice.
268pub trait ToByteSlice {
269    /// Converts this instance into a byte slice
270    fn to_byte_slice(&self) -> &[u8];
271}
272
273impl<T: ArrowNativeType> ToByteSlice for [T] {
274    #[inline]
275    fn to_byte_slice(&self) -> &[u8] {
276        let raw_ptr = self.as_ptr() as *const u8;
277        unsafe { std::slice::from_raw_parts(raw_ptr, std::mem::size_of_val(self)) }
278    }
279}
280
281impl<T: ArrowNativeType> ToByteSlice for T {
282    #[inline]
283    fn to_byte_slice(&self) -> &[u8] {
284        let raw_ptr = self as *const T as *const u8;
285        unsafe { std::slice::from_raw_parts(raw_ptr, std::mem::size_of::<T>()) }
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292
293    #[test]
294    fn test_i256() {
295        let a = i256::from_parts(0, 0);
296        assert_eq!(a.as_usize(), 0);
297        assert_eq!(a.to_usize().unwrap(), 0);
298        assert_eq!(a.to_isize().unwrap(), 0);
299
300        let a = i256::from_parts(0, -1);
301        assert_eq!(a.as_usize(), 0);
302        assert!(a.to_usize().is_none());
303        assert!(a.to_usize().is_none());
304
305        let a = i256::from_parts(u128::MAX, -1);
306        assert_eq!(a.as_usize(), usize::MAX);
307        assert!(a.to_usize().is_none());
308        assert_eq!(a.to_isize().unwrap(), -1);
309    }
310
311    #[test]
312    fn test_interval_usize() {
313        assert_eq!(IntervalDayTime::new(1, 0).as_usize(), 1);
314        assert_eq!(IntervalMonthDayNano::new(1, 0, 0).as_usize(), 1);
315
316        let a = IntervalDayTime::new(23, 53);
317        let b = IntervalDayTime::usize_as(a.as_usize());
318        assert_eq!(a, b);
319
320        let a = IntervalMonthDayNano::new(23, 53, 0);
321        let b = IntervalMonthDayNano::usize_as(a.as_usize());
322        assert_eq!(a, b);
323    }
324}