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

polars_ops/series/ops/
clip.rs
