polars_core/chunked_array/ops/
binning.rs

1use std::hash::{Hash, Hasher};
2use std::num::NonZeroUsize;
3use std::ops::Deref;
4
5#[cfg(feature = "serde")]
6use serde::{Deserialize, Serialize};
7
8use crate::CHEAP_SERIES_HASH_LIMIT;
9use crate::prelude::*;
10use crate::utils::Wrap;
11
12/// Breakpoints delimiting bins by value.
13///
14/// Always free of nulls and non-decreasing.
15#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
16#[derive(Clone, PartialEq, Debug)]
17pub struct Breaks(Series);
18
19impl Breaks {
20    pub fn new(breaks: Series) -> PolarsResult<Self> {
21        polars_ensure!(
22            breaks.null_count() == 0,
23            ComputeError: "breakpoints cannot contain nulls"
24        );
25        let n = breaks.len();
26        if n >= 2 {
27            let non_decreasing = breaks.slice(0, n - 1).lt_eq(&breaks.slice(1, n - 1))?;
28            polars_ensure!(
29                non_decreasing.all(),
30                ComputeError: "breakpoints must be non-decreasing"
31            );
32        }
33        Ok(Self(breaks))
34    }
35
36    pub fn into_series(self) -> Series {
37        self.0
38    }
39}
40
41/// Cumulative fractions delimiting bins by proportion.
42///
43/// Always non-decreasing and within `[0, 1]`.
44#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
45#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
46#[derive(Clone, PartialEq, Debug)]
47pub struct Fractions(Vec<f64>);
48
49impl Fractions {
50    pub fn new(fractions: Vec<f64>) -> PolarsResult<Self> {
51        for x in &fractions {
52            polars_ensure!(
53                (0.0..=1.0).contains(x),
54                ComputeError: "fractions must be between 0.0 and 1.0, got {}", x
55            );
56        }
57        polars_ensure!(
58            fractions.is_sorted(),
59            ComputeError: "fractions must be non-decreasing"
60        );
61        Ok(Self(fractions))
62    }
63}
64
65/// How interval binning delimits its bins, with its breakpoints resolved.
66#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
67#[derive(Clone, PartialEq, Debug, Hash)]
68pub enum IntervalSpec {
69    /// Explicit breakpoints.
70    Breaks(Breaks),
71    /// `n` equal-width bins spanning `[min, max]`.
72    Count(NonZeroUsize),
73}
74
75/// How quantile and rank binning delimit their bins.
76#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
77#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
78#[derive(Clone, PartialEq, Debug, Hash)]
79pub enum FractionSpec {
80    /// Explicit fractions.
81    Explicit(Fractions),
82    /// `n` bins of equal probability, or of equal size for rank binning.
83    Count(NonZeroUsize),
84}
85
86fn n_bins(n_bins: usize) -> PolarsResult<NonZeroUsize> {
87    NonZeroUsize::new(n_bins)
88        .ok_or_else(|| polars_err!(ComputeError: "binning requires at least one bin"))
89}
90
91impl IntervalSpec {
92    pub fn from_breaks(breaks: Series) -> PolarsResult<Self> {
93        Ok(Self::Breaks(Breaks::new(breaks)?))
94    }
95
96    pub fn from_count(count: usize) -> PolarsResult<Self> {
97        Ok(Self::Count(n_bins(count)?))
98    }
99
100    pub fn n_bins(&self) -> usize {
101        match self {
102            Self::Breaks(breaks) => breaks.len() + 1,
103            Self::Count(n_bins) => n_bins.get(),
104        }
105    }
106
107    pub fn breaks(&self) -> Option<&Series> {
108        match self {
109            Self::Breaks(breaks) => Some(&breaks.0),
110            Self::Count(_) => None,
111        }
112    }
113}
114
115impl FractionSpec {
116    pub fn from_fractions(fractions: Vec<f64>) -> PolarsResult<Self> {
117        Ok(Self::Explicit(Fractions::new(fractions)?))
118    }
119
120    pub fn from_count(count: usize) -> PolarsResult<Self> {
121        Ok(Self::Count(n_bins(count)?))
122    }
123
124    pub fn n_bins(&self) -> usize {
125        match self {
126            Self::Explicit(fractions) => fractions.len() + 1,
127            Self::Count(n_bins) => n_bins.get(),
128        }
129    }
130}
131
132impl Deref for Breaks {
133    type Target = Series;
134
135    fn deref(&self) -> &Self::Target {
136        &self.0
137    }
138}
139
140impl Deref for Fractions {
141    type Target = [f64];
142
143    fn deref(&self) -> &Self::Target {
144        &self.0
145    }
146}
147
148impl Hash for Breaks {
149    fn hash<H: Hasher>(&self, state: &mut H) {
150        Wrap(self.0.slice(0, CHEAP_SERIES_HASH_LIMIT)).hash(state)
151    }
152}
153
154impl Hash for Fractions {
155    fn hash<H: Hasher>(&self, state: &mut H) {
156        bytemuck::cast_slice::<_, u64>(&self.0).hash(state)
157    }
158}
polars_core/chunked_array/ops/binning.rs

polars_core/chunked_array/ops/
binning.rs
