polars_core/chunked_array/
float.rs

1use arrow::bitmap::Bitmap;
2use arrow::legacy::kernels::set::set_at_nulls;
3use num_traits::Float;
4use polars_compute::nan::{is_nan, is_not_nan};
5use polars_utils::float16::pf16;
6use polars_utils::total_ord::{canonical_f16, canonical_f32, canonical_f64};
7
8use crate::prelude::arity::{unary_elementwise_values, unary_kernel};
9use crate::prelude::*;
10
11impl<T> ChunkedArray<T>
12where
13    T: PolarsFloatType,
14    T::Native: Float,
15{
16    pub fn is_nan(&self) -> BooleanChunked {
17        unary_kernel(self, |arr| {
18            let out = is_nan(arr.values()).unwrap_or_else(|| Bitmap::new_zeroed(arr.len()));
19            BooleanArray::from(out).with_validity(arr.validity().cloned())
20        })
21    }
22    pub fn is_not_nan(&self) -> BooleanChunked {
23        unary_kernel(self, |arr| {
24            let out =
25                is_not_nan(arr.values()).unwrap_or_else(|| Bitmap::new_with_value(true, arr.len()));
26            BooleanArray::from(out).with_validity(arr.validity().cloned())
27        })
28    }
29    pub fn is_finite(&self) -> BooleanChunked {
30        unary_elementwise_values(self, |x| x.is_finite())
31    }
32    pub fn is_infinite(&self) -> BooleanChunked {
33        unary_elementwise_values(self, |x| x.is_infinite())
34    }
35
36    #[must_use]
37    /// Convert missing values to `NaN` values.
38    pub fn none_to_nan(&self) -> Self {
39        let chunks = self
40            .downcast_iter()
41            .map(|arr| set_at_nulls(arr, T::Native::nan()));
42        ChunkedArray::from_chunk_iter(self.name().clone(), chunks)
43    }
44}
45
46pub trait Canonical {
47    fn canonical(self) -> Self;
48}
49
50impl Canonical for pf16 {
51    #[inline]
52    fn canonical(self) -> Self {
53        canonical_f16(self)
54    }
55}
56
57impl Canonical for f32 {
58    #[inline]
59    fn canonical(self) -> Self {
60        canonical_f32(self)
61    }
62}
63
64impl Canonical for f64 {
65    #[inline]
66    fn canonical(self) -> Self {
67        canonical_f64(self)
68    }
69}
70
71impl<T> ChunkedArray<T>
72where
73    T: PolarsFloatType,
74    T::Native: Float + Canonical,
75{
76    pub fn to_canonical(&self) -> Self {
77        unary_elementwise_values(self, |v| v.canonical())
78    }
79}