polars_utils/
float16.rs

1#[cfg(feature = "python")]
2use std::convert::Infallible;
3use std::fmt::{Display, LowerExp};
4use std::iter::Sum;
5use std::ops::*;
6
7use bytemuck::{Pod, Zeroable};
8use half;
9use num_derive::*;
10use num_traits::real::Real;
11use num_traits::{AsPrimitive, Bounded, FromBytes, One, Pow, ToBytes, Zero};
12#[cfg(feature = "python")]
13use numpy::Element;
14#[cfg(feature = "python")]
15use pyo3::types::{PyAnyMethods, PyFloat};
16#[cfg(feature = "python")]
17use pyo3::{FromPyObject, IntoPyObject, Python};
18#[cfg(feature = "serde")]
19use serde::{Deserialize, Serialize};
20
21use crate::nulls::IsNull;
22
23/// A portable float16 type.
24///
25/// This type is a newtype wrapper around `half::f16`.
26/// We intend to replace it by Rust's builtin `f16` type once it is stabilized.
27#[derive(
28    Debug,
29    Copy,
30    Clone,
31    Default,
32    PartialEq,
33    PartialOrd,
34    Zeroable,
35    Pod,
36    Float,
37    FromPrimitive,
38    Num,
39    NumCast,
40    NumOps,
41    One,
42    ToPrimitive,
43    Zero,
44)]
45#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
46#[allow(non_camel_case_types)]
47#[repr(transparent)]
48pub struct pf16(pub half::f16);
49
50#[cfg(feature = "dsl-schema")]
51impl schemars::JsonSchema for pf16 {
52    fn schema_name() -> std::borrow::Cow<'static, str> {
53        "f16".into()
54    }
55
56    fn schema_id() -> std::borrow::Cow<'static, str> {
57        std::borrow::Cow::Borrowed(concat!(module_path!(), "::", "pf16"))
58    }
59
60    fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
61        f32::json_schema(generator)
62    }
63}
64
65impl pf16 {
66    pub const NAN: Self = pf16(half::f16::NAN);
67    pub const INFINITY: Self = pf16(half::f16::INFINITY);
68    pub const NEG_INFINITY: Self = pf16(half::f16::NEG_INFINITY);
69
70    #[inline]
71    pub fn to_le_bytes(&self) -> [u8; 2] {
72        self.0.to_le_bytes()
73    }
74
75    #[inline]
76    pub fn is_nan(self) -> bool {
77        self.0.is_nan()
78    }
79
80    #[inline]
81    pub fn is_finite(self) -> bool {
82        self.0.is_finite()
83    }
84
85    #[inline]
86    pub fn abs(self) -> Self {
87        pf16(self.0.abs())
88    }
89
90    #[inline]
91    pub fn trunc(self) -> Self {
92        pf16(self.0.trunc())
93    }
94
95    #[inline]
96    pub fn round(self) -> Self {
97        pf16(self.0.round())
98    }
99
100    #[inline]
101    pub fn floor(self) -> Self {
102        pf16(self.0.floor())
103    }
104
105    #[inline]
106    pub fn ceil(self) -> Self {
107        pf16(self.0.ceil())
108    }
109
110    #[inline]
111    pub fn log(self, base: Self) -> Self {
112        pf16(self.0.log(base.0))
113    }
114
115    #[inline]
116    pub fn to_bits(self) -> u16 {
117        self.0.to_bits()
118    }
119
120    #[inline]
121    pub fn from_bits(b: u16) -> Self {
122        pf16(half::f16::from_bits(b))
123    }
124
125    #[inline]
126    pub fn min(self, other: Self) -> Self {
127        pf16(self.0.min(other.0))
128    }
129
130    #[inline]
131    pub fn max(self, other: Self) -> Self {
132        pf16(self.0.max(other.0))
133    }
134}
135
136impl Display for pf16 {
137    #[inline]
138    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
139        Display::fmt(&self.0.to_f32(), f)
140    }
141}
142
143impl Neg for pf16 {
144    type Output = Self;
145
146    #[inline]
147    fn neg(self) -> Self::Output {
148        pf16(-self.0)
149    }
150}
151
152impl AddAssign for pf16 {
153    #[inline]
154    fn add_assign(&mut self, other: Self) {
155        self.0 = self.0 + other.0;
156    }
157}
158
159impl SubAssign for pf16 {
160    #[inline]
161    fn sub_assign(&mut self, other: Self) {
162        self.0 = self.0 - other.0;
163    }
164}
165
166impl MulAssign for pf16 {
167    #[inline]
168    fn mul_assign(&mut self, other: Self) {
169        self.0 = self.0 * other.0;
170    }
171}
172
173impl DivAssign for pf16 {
174    #[inline]
175    fn div_assign(&mut self, other: Self) {
176        self.0 = self.0 / other.0;
177    }
178}
179
180impl RemAssign for pf16 {
181    #[inline]
182    fn rem_assign(&mut self, other: Self) {
183        self.0 = self.0 % other.0;
184    }
185}
186
187impl Sum for pf16 {
188    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
189        let mut acc = Zero::zero();
190        for v in iter {
191            acc += v.0;
192        }
193        pf16(acc)
194    }
195}
196
197impl Pow<pf16> for pf16 {
198    type Output = Self;
199
200    #[inline]
201    fn pow(self, rhs: pf16) -> Self::Output {
202        pf16(self.0.powf(rhs.0))
203    }
204}
205
206impl From<bool> for pf16 {
207    #[inline]
208    fn from(value: bool) -> Self {
209        if value { One::one() } else { Zero::zero() }
210    }
211}
212
213impl From<pf16> for f32 {
214    #[inline]
215    fn from(value: pf16) -> Self {
216        value.0.to_f32()
217    }
218}
219
220impl From<pf16> for f64 {
221    #[inline]
222    fn from(value: pf16) -> Self {
223        value.0.to_f64()
224    }
225}
226
227impl From<f32> for pf16 {
228    #[inline]
229    fn from(value: f32) -> Self {
230        pf16(half::f16::from_f32(value))
231    }
232}
233
234impl From<f64> for pf16 {
235    #[inline]
236    fn from(value: f64) -> Self {
237        pf16(half::f16::from_f64(value))
238    }
239}
240
241impl Bounded for pf16 {
242    fn min_value() -> Self {
243        pf16(half::f16::MIN)
244    }
245
246    fn max_value() -> Self {
247        pf16(half::f16::MAX)
248    }
249}
250
251macro_rules! impl_as_primitive {
252    ($ty:ty) => {
253        impl AsPrimitive<$ty> for pf16 {
254            #[inline]
255            fn as_(self) -> $ty {
256                self.0.as_()
257            }
258        }
259    };
260}
261
262impl_as_primitive!(u8);
263impl_as_primitive!(u16);
264impl_as_primitive!(u32);
265impl_as_primitive!(u64);
266impl_as_primitive!(i8);
267impl_as_primitive!(i16);
268impl_as_primitive!(i32);
269impl_as_primitive!(i64);
270impl_as_primitive!(f32);
271impl_as_primitive!(f64);
272
273impl AsPrimitive<pf16> for pf16 {
274    #[inline]
275    fn as_(self) -> pf16 {
276        self
277    }
278}
279
280impl AsPrimitive<u128> for pf16 {
281    #[inline]
282    fn as_(self) -> u128 {
283        self.0.to_f64() as u128
284    }
285}
286
287impl AsPrimitive<i128> for pf16 {
288    #[inline]
289    fn as_(self) -> i128 {
290        self.0.to_f64() as i128
291    }
292}
293
294macro_rules! impl_as_primitive_pf16_from {
295    ($ty:ty) => {
296        impl AsPrimitive<pf16> for $ty {
297            #[inline]
298            fn as_(self) -> pf16 {
299                pf16(<$ty>::as_(self))
300            }
301        }
302    };
303}
304impl_as_primitive_pf16_from!(usize);
305impl_as_primitive_pf16_from!(u8);
306impl_as_primitive_pf16_from!(u16);
307impl_as_primitive_pf16_from!(u32);
308impl_as_primitive_pf16_from!(u64);
309impl_as_primitive_pf16_from!(isize);
310impl_as_primitive_pf16_from!(i8);
311impl_as_primitive_pf16_from!(i16);
312impl_as_primitive_pf16_from!(i32);
313impl_as_primitive_pf16_from!(i64);
314impl_as_primitive_pf16_from!(f32);
315impl_as_primitive_pf16_from!(f64);
316
317impl AsPrimitive<pf16> for i128 {
318    #[inline]
319    fn as_(self) -> pf16 {
320        pf16(half::f16::from_f64(self as f64))
321    }
322}
323
324impl AsPrimitive<pf16> for u128 {
325    #[inline]
326    fn as_(self) -> pf16 {
327        pf16(half::f16::from_f64(self as f64))
328    }
329}
330
331impl IsNull for pf16 {
332    const HAS_NULLS: bool = false;
333    type Inner = pf16;
334
335    #[inline(always)]
336    fn is_null(&self) -> bool {
337        false
338    }
339    #[inline]
340    fn unwrap_inner(self) -> Self::Inner {
341        self
342    }
343}
344
345impl ToBytes for pf16 {
346    type Bytes = [u8; 2];
347
348    #[inline]
349    fn to_be_bytes(&self) -> Self::Bytes {
350        self.0.to_be_bytes()
351    }
352
353    #[inline]
354    fn to_le_bytes(&self) -> Self::Bytes {
355        self.0.to_le_bytes()
356    }
357
358    #[inline]
359    fn to_ne_bytes(&self) -> Self::Bytes {
360        self.0.to_ne_bytes()
361    }
362}
363
364impl FromBytes for pf16 {
365    type Bytes = [u8; 2];
366
367    #[inline]
368    fn from_be_bytes(bytes: &Self::Bytes) -> Self {
369        pf16(half::f16::from_be_bytes(*bytes))
370    }
371
372    #[inline]
373    fn from_le_bytes(bytes: &Self::Bytes) -> Self {
374        pf16(half::f16::from_le_bytes(*bytes))
375    }
376
377    #[inline]
378    fn from_ne_bytes(bytes: &Self::Bytes) -> Self {
379        pf16(half::f16::from_ne_bytes(*bytes))
380    }
381}
382
383impl LowerExp for pf16 {
384    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
385        LowerExp::fmt(&self.0.to_f32(), f)
386    }
387}
388
389#[cfg(feature = "python")]
390impl<'py> IntoPyObject<'py> for pf16 {
391    type Target = PyFloat;
392    type Output = pyo3::Bound<'py, Self::Target>;
393    type Error = Infallible;
394
395    #[inline]
396    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
397        f32::into_pyobject(self.into(), py)
398    }
399}
400
401#[cfg(feature = "python")]
402impl<'py> FromPyObject<'py> for pf16 {
403    fn extract_bound(ob: &pyo3::Bound<'py, pyo3::PyAny>) -> pyo3::PyResult<Self> {
404        let v: f32 = ob.extract()?;
405        Ok(v.as_())
406    }
407}
408
409#[cfg(feature = "python")]
410unsafe impl Element for pf16 {
411    const IS_COPY: bool = half::f16::IS_COPY;
412
413    fn get_dtype(py: Python<'_>) -> pyo3::Bound<'_, numpy::PyArrayDescr> {
414        half::f16::get_dtype(py)
415    }
416    fn clone_ref(&self, py: Python<'_>) -> Self {
417        pf16(half::f16::clone_ref(&self.0, py))
418    }
419}
polars_utils/float16.rs

polars_utils/
float16.rs
