polars_ops/series/ops/
clip.rs

Help
1use polars_core::prelude::arity::{binary_elementwise, ternary_elementwise, unary_elementwise};
2use polars_core::prelude::*;
3use polars_core::with_match_physical_numeric_polars_type;
4
5/// Set values outside the given boundaries to the boundary value.
6pub fn clip(s: &Series, min: &Series, max: &Series) -> PolarsResult<Series> {
7    polars_ensure!(
8        s.dtype().to_physical().is_primitive_numeric(),
9        InvalidOperation: "`clip` only supports physical numeric types"
10    );
11
12    let original_type = s.dtype();
13    let (min, max) = (min.strict_cast(s.dtype())?, max.strict_cast(s.dtype())?);
14
15    let (s, min, max) = (
16        s.to_physical_repr(),
17        min.to_physical_repr(),
18        max.to_physical_repr(),
19    );
20
21    with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
22        let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
23        let min: &ChunkedArray<$T> = min.as_ref().as_ref().as_ref();
24        let max: &ChunkedArray<$T> = max.as_ref().as_ref().as_ref();
25        let out = clip_helper_both_bounds(ca, min, max).into_series();
26        match original_type {
27            #[cfg(feature = "dtype-decimal")]
28            DataType::Decimal(precision, scale) => {
29                let phys = out.i128()?.as_ref().clone();
30                Ok(phys.into_decimal_unchecked(*precision, scale.unwrap()).into_series())
31            },
32            dt if dt.is_logical() => out.cast(original_type),
33            _ => Ok(out)
34        }
35    })
36}
37
38/// Set values above the given maximum to the maximum value.
39pub fn clip_max(s: &Series, max: &Series) -> PolarsResult<Series> {
40    polars_ensure!(
41        s.dtype().to_physical().is_primitive_numeric(),
42        InvalidOperation: "`clip` only supports physical numeric types"
43    );
44
45    let original_type = s.dtype();
46    let max = max.strict_cast(s.dtype())?;
47
48    let (s, max) = (s.to_physical_repr(), max.to_physical_repr());
49
50    with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
51        let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
52        let max: &ChunkedArray<$T> = max.as_ref().as_ref().as_ref();
53        let out = clip_helper_single_bound(ca, max, num_traits::clamp_max).into_series();
54        match original_type {
55            #[cfg(feature = "dtype-decimal")]
56            DataType::Decimal(precision, scale) => {
57                let phys = out.i128()?.as_ref().clone();
58                Ok(phys.into_decimal_unchecked(*precision, scale.unwrap()).into_series())
59            },
60            dt if dt.is_logical() => out.cast(original_type),
61            _ => Ok(out)
62        }
63    })
64}
65
66/// Set values below the given minimum to the minimum value.
67pub fn clip_min(s: &Series, min: &Series) -> PolarsResult<Series> {
68    polars_ensure!(
69        s.dtype().to_physical().is_primitive_numeric(),
70        InvalidOperation: "`clip` only supports physical numeric types"
71    );
72
73    let original_type = s.dtype();
74    let min = min.strict_cast(s.dtype())?;
75
76    let (s, min) = (s.to_physical_repr(), min.to_physical_repr());
77
78    with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
79        let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
80        let min: &ChunkedArray<$T> = min.as_ref().as_ref().as_ref();
81        let out = clip_helper_single_bound(ca, min, num_traits::clamp_min).into_series();
82        match original_type {
83            #[cfg(feature = "dtype-decimal")]
84            DataType::Decimal(precision, scale) => {
85                let phys = out.i128()?.as_ref().clone();
86                Ok(phys.into_decimal_unchecked(*precision, scale.unwrap()).into_series())
87            },
88            dt if dt.is_logical() => out.cast(original_type),
89            _ => Ok(out)
90        }
91    })
92}
93
94fn clip_helper_both_bounds<T>(
95    ca: &ChunkedArray<T>,
96    min: &ChunkedArray<T>,
97    max: &ChunkedArray<T>,
98) -> ChunkedArray<T>
99where
100    T: PolarsNumericType,
101    T::Native: PartialOrd,
102{
103    match (min.len(), max.len()) {
104        (1, 1) => match (min.get(0), max.get(0)) {
105            (Some(min), Some(max)) => clip_unary(ca, |v| num_traits::clamp(v, min, max)),
106            (Some(min), None) => clip_unary(ca, |v| num_traits::clamp_min(v, min)),
107            (None, Some(max)) => clip_unary(ca, |v| num_traits::clamp_max(v, max)),
108            (None, None) => ca.clone(),
109        },
110        (1, _) => match min.get(0) {
111            Some(min) => clip_binary(ca, max, |v, b| num_traits::clamp(v, min, b)),
112            None => clip_binary(ca, max, num_traits::clamp_max),
113        },
114        (_, 1) => match max.get(0) {
115            Some(max) => clip_binary(ca, min, |v, b| num_traits::clamp(v, b, max)),
116            None => clip_binary(ca, min, num_traits::clamp_min),
117        },
118        _ => clip_ternary(ca, min, max),
119    }
120}
121
122fn clip_helper_single_bound<T, F>(
123    ca: &ChunkedArray<T>,
124    bound: &ChunkedArray<T>,
125    op: F,
126) -> ChunkedArray<T>
127where
128    T: PolarsNumericType,
129    T::Native: PartialOrd,
130    F: Fn(T::Native, T::Native) -> T::Native,
131{
132    match bound.len() {
133        1 => match bound.get(0) {
134            Some(bound) => clip_unary(ca, |v| op(v, bound)),
135            None => ca.clone(),
136        },
137        _ => clip_binary(ca, bound, op),
138    }
139}
140
141fn clip_unary<T, F>(ca: &ChunkedArray<T>, op: F) -> ChunkedArray<T>
142where
143    T: PolarsNumericType,
144    F: Fn(T::Native) -> T::Native + Copy,
145{
146    unary_elementwise(ca, |v| v.map(op))
147}
148
149fn clip_binary<T, F>(ca: &ChunkedArray<T>, bound: &ChunkedArray<T>, op: F) -> ChunkedArray<T>
150where
151    T: PolarsNumericType,
152    T::Native: PartialOrd,
153    F: Fn(T::Native, T::Native) -> T::Native,
154{
155    binary_elementwise(ca, bound, |opt_s, opt_bound| match (opt_s, opt_bound) {
156        (Some(s), Some(bound)) => Some(op(s, bound)),
157        (Some(s), None) => Some(s),
158        (None, _) => None,
159    })
160}
161
162fn clip_ternary<T>(
163    ca: &ChunkedArray<T>,
164    min: &ChunkedArray<T>,
165    max: &ChunkedArray<T>,
166) -> ChunkedArray<T>
167where
168    T: PolarsNumericType,
169    T::Native: PartialOrd,
170{
171    ternary_elementwise(ca, min, max, |opt_v, opt_min, opt_max| {
172        match (opt_v, opt_min, opt_max) {
173            (Some(v), Some(min), Some(max)) => Some(num_traits::clamp(v, min, max)),
174            (Some(v), Some(min), None) => Some(num_traits::clamp_min(v, min)),
175            (Some(v), None, Some(max)) => Some(num_traits::clamp_max(v, max)),
176            (Some(v), None, None) => Some(v),
177            (None, _, _) => None,
178        }
179    })
180}
polars_ops/series/ops/clip.rs

polars_ops/series/ops/
clip.rs
