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

polars_ops/series/ops/
clip.rs
