polars_arrow/types/
aligned_bytes.rs

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

polars_arrow/types/
aligned_bytes.rs
