Skip to main content

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