Skip to main content

polars_arrow/types/
native.rs

1use std::hash::Hash;
2use std::ops::Neg;
3use std::panic::RefUnwindSafe;
4
5use bytemuck::{Pod, Zeroable};
6use half;
7use polars_utils::float16::pf16;
8use polars_utils::min_max::MinMax;
9use polars_utils::nulls::IsNull;
10use polars_utils::total_ord::{TotalEq, TotalOrd};
11
12use super::PrimitiveType;
13use super::aligned_bytes::*;
14
15/// Sealed trait implemented by all physical types that can be allocated,
16/// serialized and deserialized by this crate.
17/// All O(N) allocations in this crate are done for this trait alone.
18pub trait NativeType:
19    super::private::Sealed
20    + Pod
21    + Send
22    + Sync
23    + Sized
24    + RefUnwindSafe
25    + std::fmt::Debug
26    + std::fmt::Display
27    + PartialEq
28    + Default
29    + Copy
30    + TotalOrd
31    + IsNull
32    + MinMax
33{
34    /// The corresponding variant of [`PrimitiveType`].
35    const PRIMITIVE: PrimitiveType;
36
37    /// Type denoting its representation as bytes.
38    /// This is `[u8; N]` where `N = size_of::<T>`.
39    type Bytes: AsRef<[u8]>
40        + AsMut<[u8]>
41        + std::ops::Index<usize, Output = u8>
42        + std::ops::IndexMut<usize, Output = u8>
43        + for<'a> TryFrom<&'a [u8]>
44        + std::fmt::Debug
45        + Default
46        + IntoIterator<Item = u8>;
47
48    /// Type denoting its representation as aligned bytes.
49    ///
50    /// This is `[u8; N]` where `N = size_of::<Self>` and has alignment `align_of::<Self>`.
51    type AlignedBytes: AlignedBytes<Unaligned = Self::Bytes> + From<Self> + Into<Self>;
52
53    /// To bytes in little endian
54    fn to_le_bytes(&self) -> Self::Bytes;
55
56    /// To bytes in big endian
57    fn to_be_bytes(&self) -> Self::Bytes;
58
59    /// From bytes in little endian
60    fn from_le_bytes(bytes: Self::Bytes) -> Self;
61
62    /// From bytes in big endian
63    fn from_be_bytes(bytes: Self::Bytes) -> Self;
64}
65
66macro_rules! native_type {
67    ($type:ty, $aligned:ty, $primitive_type:expr) => {
68        impl NativeType for $type {
69            const PRIMITIVE: PrimitiveType = $primitive_type;
70
71            type Bytes = [u8; std::mem::size_of::<Self>()];
72            type AlignedBytes = $aligned;
73
74            #[inline]
75            fn to_le_bytes(&self) -> Self::Bytes {
76                Self::to_le_bytes(*self)
77            }
78
79            #[inline]
80            fn to_be_bytes(&self) -> Self::Bytes {
81                Self::to_be_bytes(*self)
82            }
83
84            #[inline]
85            fn from_le_bytes(bytes: Self::Bytes) -> Self {
86                Self::from_le_bytes(bytes)
87            }
88
89            #[inline]
90            fn from_be_bytes(bytes: Self::Bytes) -> Self {
91                Self::from_be_bytes(bytes)
92            }
93        }
94    };
95}
96
97native_type!(u8, Bytes1Alignment1, PrimitiveType::UInt8);
98native_type!(u16, Bytes2Alignment2, PrimitiveType::UInt16);
99native_type!(u32, Bytes4Alignment4, PrimitiveType::UInt32);
100native_type!(u64, Bytes8Alignment8, PrimitiveType::UInt64);
101native_type!(u128, Bytes16Alignment16, PrimitiveType::UInt128);
102native_type!(i8, Bytes1Alignment1, PrimitiveType::Int8);
103native_type!(i16, Bytes2Alignment2, PrimitiveType::Int16);
104native_type!(i32, Bytes4Alignment4, PrimitiveType::Int32);
105native_type!(i64, Bytes8Alignment8, PrimitiveType::Int64);
106native_type!(i128, Bytes16Alignment16, PrimitiveType::Int128);
107native_type!(f32, Bytes4Alignment4, PrimitiveType::Float32);
108native_type!(f64, Bytes8Alignment8, PrimitiveType::Float64);
109
110/// The in-memory representation of the DayMillisecond variant of arrow's "Interval" logical type.
111#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Zeroable, Pod)]
112#[allow(non_camel_case_types)]
113#[repr(C)]
114pub struct days_ms(pub i32, pub i32);
115
116impl days_ms {
117    /// A new [`days_ms`].
118    #[inline]
119    pub fn new(days: i32, milliseconds: i32) -> Self {
120        Self(days, milliseconds)
121    }
122
123    /// The number of days
124    #[inline]
125    pub fn days(&self) -> i32 {
126        self.0
127    }
128
129    /// The number of milliseconds
130    #[inline]
131    pub fn milliseconds(&self) -> i32 {
132        self.1
133    }
134}
135
136impl TotalEq for days_ms {
137    #[inline]
138    fn tot_eq(&self, other: &Self) -> bool {
139        self == other
140    }
141}
142
143impl TotalOrd for days_ms {
144    #[inline]
145    fn tot_cmp(&self, other: &Self) -> std::cmp::Ordering {
146        self.days()
147            .cmp(&other.days())
148            .then(self.milliseconds().cmp(&other.milliseconds()))
149    }
150}
151
152impl MinMax for days_ms {
153    fn nan_min_lt(&self, other: &Self) -> bool {
154        self < other
155    }
156
157    fn nan_max_lt(&self, other: &Self) -> bool {
158        self < other
159    }
160}
161
162impl NativeType for days_ms {
163    const PRIMITIVE: PrimitiveType = PrimitiveType::DaysMs;
164
165    type Bytes = [u8; 8];
166    type AlignedBytes = Bytes8Alignment4;
167
168    #[inline]
169    fn to_le_bytes(&self) -> Self::Bytes {
170        let days = self.0.to_le_bytes();
171        let ms = self.1.to_le_bytes();
172        let mut result = [0; 8];
173        result[0] = days[0];
174        result[1] = days[1];
175        result[2] = days[2];
176        result[3] = days[3];
177        result[4] = ms[0];
178        result[5] = ms[1];
179        result[6] = ms[2];
180        result[7] = ms[3];
181        result
182    }
183
184    #[inline]
185    fn to_be_bytes(&self) -> Self::Bytes {
186        let days = self.0.to_be_bytes();
187        let ms = self.1.to_be_bytes();
188        let mut result = [0; 8];
189        result[0] = days[0];
190        result[1] = days[1];
191        result[2] = days[2];
192        result[3] = days[3];
193        result[4] = ms[0];
194        result[5] = ms[1];
195        result[6] = ms[2];
196        result[7] = ms[3];
197        result
198    }
199
200    #[inline]
201    fn from_le_bytes(bytes: Self::Bytes) -> Self {
202        let mut days = [0; 4];
203        days[0] = bytes[0];
204        days[1] = bytes[1];
205        days[2] = bytes[2];
206        days[3] = bytes[3];
207        let mut ms = [0; 4];
208        ms[0] = bytes[4];
209        ms[1] = bytes[5];
210        ms[2] = bytes[6];
211        ms[3] = bytes[7];
212        Self(i32::from_le_bytes(days), i32::from_le_bytes(ms))
213    }
214
215    #[inline]
216    fn from_be_bytes(bytes: Self::Bytes) -> Self {
217        let mut days = [0; 4];
218        days[0] = bytes[0];
219        days[1] = bytes[1];
220        days[2] = bytes[2];
221        days[3] = bytes[3];
222        let mut ms = [0; 4];
223        ms[0] = bytes[4];
224        ms[1] = bytes[5];
225        ms[2] = bytes[6];
226        ms[3] = bytes[7];
227        Self(i32::from_be_bytes(days), i32::from_be_bytes(ms))
228    }
229}
230
231/// The in-memory representation of the MonthDayNano variant of the "Interval" logical type.
232#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Zeroable, Pod)]
233#[allow(non_camel_case_types)]
234#[repr(C)]
235pub struct months_days_ns(pub i32, pub i32, pub i64);
236
237impl IsNull for months_days_ns {
238    const HAS_NULLS: bool = false;
239    type Inner = months_days_ns;
240
241    fn is_null(&self) -> bool {
242        false
243    }
244
245    fn unwrap_inner(self) -> Self::Inner {
246        self
247    }
248}
249
250impl months_days_ns {
251    /// A new [`months_days_ns`].
252    #[inline]
253    pub fn new(months: i32, days: i32, nanoseconds: i64) -> Self {
254        Self(months, days, nanoseconds)
255    }
256
257    /// The number of months
258    #[inline]
259    pub fn months(&self) -> i32 {
260        self.0
261    }
262
263    /// The number of days
264    #[inline]
265    pub fn days(&self) -> i32 {
266        self.1
267    }
268
269    /// The number of nanoseconds
270    #[inline]
271    pub fn ns(&self) -> i64 {
272        self.2
273    }
274}
275
276impl TotalEq for months_days_ns {
277    #[inline]
278    fn tot_eq(&self, other: &Self) -> bool {
279        self == other
280    }
281}
282
283impl TotalOrd for months_days_ns {
284    #[inline]
285    fn tot_cmp(&self, other: &Self) -> std::cmp::Ordering {
286        self.months()
287            .cmp(&other.months())
288            .then(self.days().cmp(&other.days()))
289            .then(self.ns().cmp(&other.ns()))
290    }
291}
292
293impl MinMax for months_days_ns {
294    fn nan_min_lt(&self, other: &Self) -> bool {
295        self < other
296    }
297
298    fn nan_max_lt(&self, other: &Self) -> bool {
299        self < other
300    }
301}
302
303impl NativeType for months_days_ns {
304    const PRIMITIVE: PrimitiveType = PrimitiveType::MonthDayNano;
305
306    type Bytes = [u8; 16];
307    type AlignedBytes = Bytes16Alignment8;
308
309    #[inline]
310    fn to_le_bytes(&self) -> Self::Bytes {
311        let months = self.months().to_le_bytes();
312        let days = self.days().to_le_bytes();
313        let ns = self.ns().to_le_bytes();
314        let mut result = [0; 16];
315        result[0] = months[0];
316        result[1] = months[1];
317        result[2] = months[2];
318        result[3] = months[3];
319        result[4] = days[0];
320        result[5] = days[1];
321        result[6] = days[2];
322        result[7] = days[3];
323        (0..8).for_each(|i| {
324            result[8 + i] = ns[i];
325        });
326        result
327    }
328
329    #[inline]
330    fn to_be_bytes(&self) -> Self::Bytes {
331        let months = self.months().to_be_bytes();
332        let days = self.days().to_be_bytes();
333        let ns = self.ns().to_be_bytes();
334        let mut result = [0; 16];
335        result[0] = months[0];
336        result[1] = months[1];
337        result[2] = months[2];
338        result[3] = months[3];
339        result[4] = days[0];
340        result[5] = days[1];
341        result[6] = days[2];
342        result[7] = days[3];
343        (0..8).for_each(|i| {
344            result[8 + i] = ns[i];
345        });
346        result
347    }
348
349    #[inline]
350    fn from_le_bytes(bytes: Self::Bytes) -> Self {
351        let mut months = [0; 4];
352        months[0] = bytes[0];
353        months[1] = bytes[1];
354        months[2] = bytes[2];
355        months[3] = bytes[3];
356        let mut days = [0; 4];
357        days[0] = bytes[4];
358        days[1] = bytes[5];
359        days[2] = bytes[6];
360        days[3] = bytes[7];
361        let mut ns = [0; 8];
362        (0..8).for_each(|i| {
363            ns[i] = bytes[8 + i];
364        });
365        Self(
366            i32::from_le_bytes(months),
367            i32::from_le_bytes(days),
368            i64::from_le_bytes(ns),
369        )
370    }
371
372    #[inline]
373    fn from_be_bytes(bytes: Self::Bytes) -> Self {
374        let mut months = [0; 4];
375        months[0] = bytes[0];
376        months[1] = bytes[1];
377        months[2] = bytes[2];
378        months[3] = bytes[3];
379        let mut days = [0; 4];
380        days[0] = bytes[4];
381        days[1] = bytes[5];
382        days[2] = bytes[6];
383        days[3] = bytes[7];
384        let mut ns = [0; 8];
385        (0..8).for_each(|i| {
386            ns[i] = bytes[8 + i];
387        });
388        Self(
389            i32::from_be_bytes(months),
390            i32::from_be_bytes(days),
391            i64::from_be_bytes(ns),
392        )
393    }
394}
395
396impl IsNull for days_ms {
397    const HAS_NULLS: bool = false;
398    type Inner = days_ms;
399    fn is_null(&self) -> bool {
400        false
401    }
402    fn unwrap_inner(self) -> Self::Inner {
403        self
404    }
405}
406
407impl std::fmt::Display for days_ms {
408    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
409        write!(f, "{}d {}ms", self.days(), self.milliseconds())
410    }
411}
412
413impl std::fmt::Display for months_days_ns {
414    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
415        write!(f, "{}m {}d {}ns", self.months(), self.days(), self.ns())
416    }
417}
418
419impl Neg for days_ms {
420    type Output = Self;
421
422    #[inline(always)]
423    fn neg(self) -> Self::Output {
424        Self::new(-self.days(), -self.milliseconds())
425    }
426}
427
428impl Neg for months_days_ns {
429    type Output = Self;
430
431    #[inline(always)]
432    fn neg(self) -> Self::Output {
433        Self::new(-self.months(), -self.days(), -self.ns())
434    }
435}
436
437/// Physical representation of a decimal
438#[derive(Clone, Copy, Default, Eq, Hash, PartialEq, PartialOrd, Ord)]
439#[allow(non_camel_case_types)]
440#[repr(C)]
441pub struct i256(pub ethnum::I256);
442
443impl i256 {
444    /// Returns a new [`i256`] from two `i128`.
445    pub fn from_words(hi: i128, lo: i128) -> Self {
446        Self(ethnum::I256::from_words(hi, lo))
447    }
448}
449
450impl TryFrom<i256> for i128 {
451    type Error = core::num::TryFromIntError;
452
453    fn try_from(value: i256) -> Result<Self, Self::Error> {
454        value.0.try_into()
455    }
456}
457
458impl IsNull for i256 {
459    const HAS_NULLS: bool = false;
460    type Inner = i256;
461    #[inline(always)]
462    fn is_null(&self) -> bool {
463        false
464    }
465    fn unwrap_inner(self) -> Self::Inner {
466        self
467    }
468}
469
470impl Neg for i256 {
471    type Output = Self;
472
473    #[inline]
474    fn neg(self) -> Self::Output {
475        let (a, b) = self.0.into_words();
476        Self(ethnum::I256::from_words(-a, b))
477    }
478}
479
480impl std::fmt::Debug for i256 {
481    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
482        write!(f, "{:?}", self.0)
483    }
484}
485
486impl std::fmt::Display for i256 {
487    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
488        write!(f, "{}", self.0)
489    }
490}
491
492unsafe impl Pod for i256 {}
493unsafe impl Zeroable for i256 {}
494
495impl TotalEq for i256 {
496    #[inline]
497    fn tot_eq(&self, other: &Self) -> bool {
498        self == other
499    }
500}
501
502impl TotalOrd for i256 {
503    #[inline]
504    fn tot_cmp(&self, other: &Self) -> std::cmp::Ordering {
505        self.cmp(other)
506    }
507}
508
509impl MinMax for i256 {
510    fn nan_min_lt(&self, other: &Self) -> bool {
511        self < other
512    }
513
514    fn nan_max_lt(&self, other: &Self) -> bool {
515        self < other
516    }
517}
518
519impl NativeType for i256 {
520    const PRIMITIVE: PrimitiveType = PrimitiveType::Int256;
521
522    type Bytes = [u8; 32];
523    type AlignedBytes = Bytes32Alignment16;
524
525    #[inline]
526    fn to_le_bytes(&self) -> Self::Bytes {
527        let mut bytes = [0u8; 32];
528        let (a, b) = self.0.into_words();
529        let a = a.to_le_bytes();
530        (0..16).for_each(|i| {
531            bytes[i] = a[i];
532        });
533
534        let b = b.to_le_bytes();
535        (0..16).for_each(|i| {
536            bytes[i + 16] = b[i];
537        });
538
539        bytes
540    }
541
542    #[inline]
543    fn to_be_bytes(&self) -> Self::Bytes {
544        let mut bytes = [0u8; 32];
545        let (a, b) = self.0.into_words();
546
547        let a = a.to_be_bytes();
548        (0..16).for_each(|i| {
549            bytes[i] = a[i];
550        });
551
552        let b = b.to_be_bytes();
553        (0..16).for_each(|i| {
554            bytes[i + 16] = b[i];
555        });
556
557        bytes
558    }
559
560    #[inline]
561    fn from_be_bytes(bytes: Self::Bytes) -> Self {
562        let (a, b) = bytes.split_at(16);
563        let a: [u8; 16] = a.try_into().unwrap();
564        let b: [u8; 16] = b.try_into().unwrap();
565        let a = i128::from_be_bytes(a);
566        let b = i128::from_be_bytes(b);
567        Self(ethnum::I256::from_words(a, b))
568    }
569
570    #[inline]
571    fn from_le_bytes(bytes: Self::Bytes) -> Self {
572        let (b, a) = bytes.split_at(16);
573        let a: [u8; 16] = a.try_into().unwrap();
574        let b: [u8; 16] = b.try_into().unwrap();
575        let a = i128::from_le_bytes(a);
576        let b = i128::from_le_bytes(b);
577        Self(ethnum::I256::from_words(a, b))
578    }
579}
580
581impl NativeType for pf16 {
582    const PRIMITIVE: PrimitiveType = PrimitiveType::Float16;
583    type Bytes = [u8; 2];
584    type AlignedBytes = Bytes2Alignment2;
585
586    #[inline]
587    fn to_le_bytes(&self) -> Self::Bytes {
588        self.0.to_le_bytes()
589    }
590
591    #[inline]
592    fn to_be_bytes(&self) -> Self::Bytes {
593        self.0.to_be_bytes()
594    }
595
596    #[inline]
597    fn from_le_bytes(bytes: Self::Bytes) -> Self {
598        pf16(half::f16::from_le_bytes(bytes))
599    }
600
601    #[inline]
602    fn from_be_bytes(bytes: Self::Bytes) -> Self {
603        pf16(half::f16::from_be_bytes(bytes))
604    }
605}