polars_arrow/types/
aligned_bytes.rs

1use bytemuck::{Pod, Zeroable};
2
3use super::{days_ms, f16, i256, months_days_ns};
4use crate::array::View;
5
6/// A representation of a type as raw bytes with the same alignment as the original type.
7pub trait AlignedBytes: Pod + Default + Eq {
8    const ALIGNMENT: usize;
9    const SIZE: usize;
10    const SIZE_ALIGNMENT_PAIR: PrimitiveSizeAlignmentPair;
11
12    type Unaligned: AsRef<[u8]>
13        + AsMut<[u8]>
14        + std::ops::Index<usize, Output = u8>
15        + std::ops::IndexMut<usize, Output = u8>
16        + for<'a> TryFrom<&'a [u8]>
17        + std::fmt::Debug
18        + Default
19        + IntoIterator<Item = u8>
20        + Pod;
21
22    fn to_unaligned(&self) -> Self::Unaligned;
23    fn from_unaligned(unaligned: Self::Unaligned) -> Self;
24
25    fn zeros() -> Self;
26    fn ones() -> Self;
27
28    fn unsigned_leq(self, other: Self) -> bool;
29}
30
31macro_rules! impl_aligned_bytes {
32    (
33        $(($name:ident, $size:literal, $alignment:literal, $sap:ident, [$($eq_type:ty),*]$(, $unsigned:ty)?),)+
34    ) => {
35        $(
36        /// Bytes with a size and alignment.
37        ///
38        /// This is used to reduce the monomorphizations for routines that solely rely on the size
39        /// and alignment of types.
40        #[derive(Debug, Copy, Clone, PartialEq, Eq, Default, Pod, Zeroable)]
41        #[repr(C, align($alignment))]
42        pub struct $name([u8; $size]);
43
44        impl AlignedBytes for $name {
45            const ALIGNMENT: usize = $alignment;
46            const SIZE: usize = $size;
47            const SIZE_ALIGNMENT_PAIR: PrimitiveSizeAlignmentPair = PrimitiveSizeAlignmentPair::$sap;
48
49            type Unaligned = [u8; $size];
50
51            #[inline(always)]
52            fn to_unaligned(&self) -> Self::Unaligned {
53                self.0
54            }
55            #[inline(always)]
56            fn from_unaligned(unaligned: Self::Unaligned) -> Self {
57                Self(unaligned)
58            }
59            fn zeros() -> Self {
60                Self([0u8; _])
61            }
62            fn ones() -> Self {
63                Self([0xFFu8; _])
64            }
65            #[inline(always)]
66            fn unsigned_leq(self, _other: Self) -> bool {
67                $(
68                    return <$unsigned>::from(self) <= <$unsigned>::from(_other);
69                )?
70
71                #[allow(unreachable_code)]
72                {
73                    unreachable!()
74                }
75            }
76        }
77
78        impl AsRef<[u8; $size]> for $name {
79            #[inline(always)]
80            fn as_ref(&self) -> &[u8; $size] {
81                &self.0
82            }
83        }
84
85        $(
86        impl From<$eq_type> for $name {
87            #[inline(always)]
88            fn from(value: $eq_type) -> Self {
89                bytemuck::must_cast(value)
90            }
91        }
92        impl From<$name> for $eq_type {
93            #[inline(always)]
94            fn from(value: $name) -> Self {
95                bytemuck::must_cast(value)
96            }
97        }
98        )*
99        )+
100    }
101}
102
103#[derive(Clone, Copy)]
104pub enum PrimitiveSizeAlignmentPair {
105    S1A1,
106    S2A2,
107    S4A4,
108    S8A4,
109    S8A8,
110    S12A4,
111    S16A4,
112    S16A8,
113    S16A16,
114    S32A16,
115}
116
117impl PrimitiveSizeAlignmentPair {
118    pub const fn size(self) -> usize {
119        match self {
120            Self::S1A1 => 1,
121            Self::S2A2 => 2,
122            Self::S4A4 => 4,
123            Self::S8A4 | Self::S8A8 => 8,
124            Self::S12A4 => 12,
125            Self::S16A4 | Self::S16A8 | Self::S16A16 => 16,
126            Self::S32A16 => 32,
127        }
128    }
129
130    pub const fn alignment(self) -> usize {
131        match self {
132            Self::S1A1 => 1,
133            Self::S2A2 => 2,
134            Self::S4A4 | Self::S8A4 | Self::S12A4 | Self::S16A4 => 4,
135            Self::S8A8 | Self::S16A8 => 8,
136            Self::S16A16 | Self::S32A16 => 16,
137        }
138    }
139}
140
141impl_aligned_bytes! {
142    (Bytes1Alignment1, 1, 1, S1A1, [u8, i8], u8),
143    (Bytes2Alignment2, 2, 2, S2A2, [u16, i16, f16], u16),
144    (Bytes4Alignment4, 4, 4, S4A4, [u32, i32, f32], u32),
145    (Bytes8Alignment8, 8, 8, S8A8, [u64, i64, f64], u64),
146    (Bytes8Alignment4, 8, 4, S8A4, [days_ms]),
147    (Bytes12Alignment4, 12, 4, S12A4, [[u32; 3]]),
148    (Bytes16Alignment4, 16, 4, S16A4, [View]),
149    (Bytes16Alignment8, 16, 8, S16A8, [months_days_ns]),
150    (Bytes16Alignment16, 16, 16, S16A16, [u128, i128]),
151    (Bytes32Alignment16, 32, 16, S32A16, [i256]),
152}
polars_arrow/types/aligned_bytes.rs

polars_arrow/types/
aligned_bytes.rs
