polars_ops/series/ops/
clip.rs

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#[inline]
6fn clamp<T: PartialOrd>(input: T, min: T, max: T) -> T {
7    if input < min {
8        min
9    } else if input > max {
10        max
11    } else {
12        input
13    }
14}
15
16#[inline]
17fn clamp_min<T: PartialOrd>(input: T, min: T) -> T {
18    if input < min { min } else { input }
19}
20
21#[inline]
22fn clamp_max<T: PartialOrd>(input: T, max: T) -> T {
23    if input > max { max } else { input }
24}
25
26/// Set values outside the given boundaries to the boundary value.
27pub fn clip(s: &Series, min: &Series, max: &Series) -> PolarsResult<Series> {
28    polars_ensure!(
29        s.dtype().to_physical().is_primitive_numeric(),
30        InvalidOperation: "`clip` only supports physical numeric types"
31    );
32    let n = [s.len(), min.len(), max.len()]
33        .into_iter()
34        .find(|l| *l != 1)
35        .unwrap_or(1);
36
37    for (i, (name, length)) in [("self", s.len()), ("min", min.len()), ("max", max.len())]
38        .into_iter()
39        .enumerate()
40    {
41        polars_ensure!(
42            length == n || length == 1,
43            length_mismatch = "clip",
44            length,
45            n,
46            argument = name,
47            argument_idx = i
48        );
49    }
50
51    let original_type = s.dtype();
52    let (min, max) = (min.strict_cast(s.dtype())?, max.strict_cast(s.dtype())?);
53
54    let (s, min, max) = (
55        s.to_physical_repr(),
56        min.to_physical_repr(),
57        max.to_physical_repr(),
58    );
59
60    with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
61        let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
62        let min: &ChunkedArray<$T> = min.as_ref().as_ref().as_ref();
63        let max: &ChunkedArray<$T> = max.as_ref().as_ref().as_ref();
64        let out = clip_helper_both_bounds(ca, min, max).into_series();
65        match original_type {
66            #[cfg(feature = "dtype-decimal")]
67            DataType::Decimal(precision, scale) => {
68                let phys = out.i128()?.as_ref().clone();
69                Ok(phys.into_decimal_unchecked(*precision, *scale).into_series())
70            },
71            dt if dt.is_logical() => out.cast(original_type),
72            _ => Ok(out)
73        }
74    })
75}
76
77/// Set values above the given maximum to the maximum value.
78pub fn clip_max(s: &Series, max: &Series) -> PolarsResult<Series> {
79    polars_ensure!(
80        s.dtype().to_physical().is_primitive_numeric(),
81        InvalidOperation: "`clip` only supports physical numeric types"
82    );
83    polars_ensure!(
84        s.len() == max.len() || s.len() == 1 || max.len() == 1,
85        length_mismatch = "clip(max)",
86        s.len(),
87        max.len()
88    );
89
90    let original_type = s.dtype();
91    let max = max.strict_cast(s.dtype())?;
92
93    let (s, max) = (s.to_physical_repr(), max.to_physical_repr());
94
95    with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
96        let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
97        let max: &ChunkedArray<$T> = max.as_ref().as_ref().as_ref();
98        let out = clip_helper_single_bound(ca, max, clamp_max).into_series();
99        match original_type {
100            #[cfg(feature = "dtype-decimal")]
101            DataType::Decimal(precision, scale) => {
102                let phys = out.i128()?.as_ref().clone();
103                Ok(phys.into_decimal_unchecked(*precision, *scale).into_series())
104            },
105            dt if dt.is_logical() => out.cast(original_type),
106            _ => Ok(out)
107        }
108    })
109}
110
111/// Set values below the given minimum to the minimum value.
112pub fn clip_min(s: &Series, min: &Series) -> PolarsResult<Series> {
113    polars_ensure!(
114        s.dtype().to_physical().is_primitive_numeric(),
115        InvalidOperation: "`clip` only supports physical numeric types"
116    );
117    polars_ensure!(
118        s.len() == min.len() || s.len() == 1 || min.len() == 1,
119        length_mismatch = "clip(min)",
120        s.len(),
121        min.len()
122    );
123
124    let original_type = s.dtype();
125    let min = min.strict_cast(s.dtype())?;
126
127    let (s, min) = (s.to_physical_repr(), min.to_physical_repr());
128
129    with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
130        let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
131        let min: &ChunkedArray<$T> = min.as_ref().as_ref().as_ref();
132        let out = clip_helper_single_bound(ca, min, clamp_min).into_series();
133        match original_type {
134            #[cfg(feature = "dtype-decimal")]
135            DataType::Decimal(precision, scale) => {
136                let phys = out.i128()?.as_ref().clone();
137                Ok(phys.into_decimal_unchecked(*precision, *scale).into_series())
138            },
139            dt if dt.is_logical() => out.cast(original_type),
140            _ => Ok(out)
141        }
142    })
143}
144
145fn clip_helper_both_bounds<T>(
146    ca: &ChunkedArray<T>,
147    min: &ChunkedArray<T>,
148    max: &ChunkedArray<T>,
149) -> ChunkedArray<T>
150where
151    T: PolarsNumericType,
152    T::Native: PartialOrd,
153{
154    match (min.len(), max.len()) {
155        (1, 1) => match (min.get(0), max.get(0)) {
156            (Some(min), Some(max)) => clip_unary(ca, |v| clamp(v, min, max)),
157            (Some(min), None) => clip_unary(ca, |v| clamp_min(v, min)),
158            (None, Some(max)) => clip_unary(ca, |v| clamp_max(v, max)),
159            (None, None) => ca.clone(),
160        },
161        (1, _) => match min.get(0) {
162            Some(min) => binary_elementwise(ca, max, |opt_s, opt_max| match (opt_s, opt_max) {
163                (Some(s), Some(max)) => Some(clamp(s, min, max)),
164                (Some(s), None) => Some(clamp_min(s, min)),
165                (None, _) => None,
166            }),
167            None => binary_elementwise(ca, max, |opt_s, opt_max| match (opt_s, opt_max) {
168                (Some(s), Some(max)) => Some(clamp_max(s, max)),
169                (Some(s), None) => Some(s),
170                (None, _) => None,
171            }),
172        },
173        (_, 1) => match max.get(0) {
174            Some(max) => binary_elementwise(ca, min, |opt_s, opt_min| match (opt_s, opt_min) {
175                (Some(s), Some(min)) => Some(clamp(s, min, max)),
176                (Some(s), None) => Some(clamp_max(s, max)),
177                (None, _) => None,
178            }),
179            None => binary_elementwise(ca, min, |opt_s, opt_min| match (opt_s, opt_min) {
180                (Some(s), Some(min)) => Some(clamp_min(s, min)),
181                (Some(s), None) => Some(s),
182                (None, _) => None,
183            }),
184        },
185        _ => clip_ternary(ca, min, max),
186    }
187}
188
189fn clip_helper_single_bound<T, F>(
190    ca: &ChunkedArray<T>,
191    bound: &ChunkedArray<T>,
192    op: F,
193) -> ChunkedArray<T>
194where
195    T: PolarsNumericType,
196    T::Native: PartialOrd,
197    F: Fn(T::Native, T::Native) -> T::Native,
198{
199    match bound.len() {
200        1 => match bound.get(0) {
201            Some(bound) => clip_unary(ca, |v| op(v, bound)),
202            None => ca.clone(),
203        },
204        _ => binary_elementwise(ca, bound, |opt_s, opt_bound| match (opt_s, opt_bound) {
205            (Some(s), Some(bound)) => Some(op(s, bound)),
206            (Some(s), None) => Some(s),
207            (None, _) => None,
208        }),
209    }
210}
211
212fn clip_unary<T, F>(ca: &ChunkedArray<T>, op: F) -> ChunkedArray<T>
213where
214    T: PolarsNumericType,
215    F: Fn(T::Native) -> T::Native + Copy,
216{
217    unary_elementwise(ca, |v| v.map(op))
218}
219
220fn clip_ternary<T>(
221    ca: &ChunkedArray<T>,
222    min: &ChunkedArray<T>,
223    max: &ChunkedArray<T>,
224) -> ChunkedArray<T>
225where
226    T: PolarsNumericType,
227    T::Native: PartialOrd,
228{
229    ternary_elementwise(ca, min, max, |opt_v, opt_min, opt_max| {
230        match (opt_v, opt_min, opt_max) {
231            (Some(v), Some(min), Some(max)) => Some(clamp(v, min, max)),
232            (Some(v), Some(min), None) => Some(clamp_min(v, min)),
233            (Some(v), None, Some(max)) => Some(clamp_max(v, max)),
234            (Some(v), None, None) => Some(v),
235            (None, _, _) => None,
236        }
237    })
238}
polars_ops/series/ops/clip.rs

polars_ops/series/ops/
clip.rs
