polars_core/chunked_array/ops/
fill_null.rs

1use arrow::bitmap::{Bitmap, BitmapBuilder};
2use arrow::legacy::kernels::set::set_at_nulls;
3use bytemuck::Zeroable;
4use num_traits::{NumCast, One, Zero};
5use polars_utils::itertools::Itertools;
6
7use crate::prelude::*;
8
9fn err_fill_null() -> PolarsError {
10    polars_err!(ComputeError: "could not determine the fill value")
11}
12
13impl Series {
14    /// Replace None values with one of the following strategies:
15    /// * Forward fill (replace None with the previous value)
16    /// * Backward fill (replace None with the next value)
17    /// * Mean fill (replace None with the mean of the whole array)
18    /// * Min fill (replace None with the minimum of the whole array)
19    /// * Max fill (replace None with the maximum of the whole array)
20    /// * Zero fill (replace None with the value zero)
21    /// * One fill (replace None with the value one)
22    ///
23    /// *NOTE: If you want to fill the Nones with a value use the
24    /// [`fill_null` operation on `ChunkedArray<T>`](crate::chunked_array::ops::ChunkFillNullValue)*.
25    ///
26    /// # Example
27    ///
28    /// ```rust
29    /// # use polars_core::prelude::*;
30    /// fn example() -> PolarsResult<()> {
31    ///     let s = Column::new("some_missing".into(), &[Some(1), None, Some(2)]);
32    ///
33    ///     let filled = s.fill_null(FillNullStrategy::Forward(None))?;
34    ///     assert_eq!(Vec::from(filled.i32()?), &[Some(1), Some(1), Some(2)]);
35    ///
36    ///     let filled = s.fill_null(FillNullStrategy::Backward(None))?;
37    ///     assert_eq!(Vec::from(filled.i32()?), &[Some(1), Some(2), Some(2)]);
38    ///
39    ///     let filled = s.fill_null(FillNullStrategy::Min)?;
40    ///     assert_eq!(Vec::from(filled.i32()?), &[Some(1), Some(1), Some(2)]);
41    ///
42    ///     let filled = s.fill_null(FillNullStrategy::Max)?;
43    ///     assert_eq!(Vec::from(filled.i32()?), &[Some(1), Some(2), Some(2)]);
44    ///
45    ///     let filled = s.fill_null(FillNullStrategy::Mean)?;
46    ///     assert_eq!(Vec::from(filled.i32()?), &[Some(1), Some(1), Some(2)]);
47    ///
48    ///     let filled = s.fill_null(FillNullStrategy::Zero)?;
49    ///     assert_eq!(Vec::from(filled.i32()?), &[Some(1), Some(0), Some(2)]);
50    ///
51    ///     let filled = s.fill_null(FillNullStrategy::One)?;
52    ///     assert_eq!(Vec::from(filled.i32()?), &[Some(1), Some(1), Some(2)]);
53    ///
54    ///     Ok(())
55    /// }
56    /// example();
57    /// ```
58    pub fn fill_null(&self, strategy: FillNullStrategy) -> PolarsResult<Series> {
59        // Nothing to fill.
60        let nc = self.null_count();
61        if nc == 0
62            || (nc == self.len()
63                && matches!(
64                    strategy,
65                    FillNullStrategy::Forward(_)
66                        | FillNullStrategy::Backward(_)
67                        | FillNullStrategy::Max
68                        | FillNullStrategy::Min
69                        | FillNullStrategy::Mean
70                ))
71        {
72            return Ok(self.clone());
73        }
74
75        let physical_type = self.dtype().to_physical();
76
77        match strategy {
78            FillNullStrategy::Forward(None) if !physical_type.is_primitive_numeric() => {
79                fill_forward_gather(self)
80            },
81
82            // Fast path to remove limit.
83            FillNullStrategy::Forward(Some(limit)) if limit >= nc as IdxSize => {
84                self.fill_null(FillNullStrategy::Forward(None))
85            },
86            FillNullStrategy::Backward(Some(limit)) if limit >= nc as IdxSize => {
87                self.fill_null(FillNullStrategy::Backward(None))
88            },
89
90            FillNullStrategy::Forward(Some(limit)) => fill_forward_gather_limit(self, limit),
91            FillNullStrategy::Backward(None) if !physical_type.is_primitive_numeric() => {
92                fill_backward_gather(self)
93            },
94            FillNullStrategy::Backward(Some(limit)) => fill_backward_gather_limit(self, limit),
95            #[cfg(feature = "dtype-decimal")]
96            FillNullStrategy::One if self.dtype().is_decimal() => {
97                use polars_compute::decimal::i128_to_dec128;
98
99                let ca = self.decimal().unwrap();
100                let precision = ca.precision();
101                let scale = ca.scale();
102                let fill_value = i128_to_dec128(1, precision, scale).ok_or_else(|| {
103                    polars_err!(ComputeError: "value '1' is out of range for Decimal({precision}, {scale})")
104                })?;
105                let phys = ca.physical().fill_null_with_values(fill_value)?;
106                Ok(phys.into_decimal_unchecked(precision, scale).into_series())
107            },
108            _ => {
109                let logical_type = self.dtype();
110
111                let s = self.to_physical_repr();
112                use DataType::*;
113                let out = match s.dtype() {
114                    #[cfg(feature = "dtype-categorical")]
115                    _ if (logical_type.is_categorical() || logical_type.is_enum())
116                        && (!matches!(
117                            strategy,
118                            FillNullStrategy::Forward(_) | FillNullStrategy::Backward(_)
119                        )) =>
120                    {
121                        with_match_categorical_physical_type!(logical_type.cat_physical().unwrap(), |$C| {
122                            let ca = self.cat::<$C>().unwrap();
123                            fill_null_cat(ca, strategy)
124                        })
125                    },
126
127                    Boolean => fill_null_bool(s.bool().unwrap(), strategy),
128                    String => {
129                        let s = unsafe { s.cast_unchecked(&Binary)? };
130                        let out = s.fill_null(strategy)?;
131                        return unsafe { out.cast_unchecked(&String) };
132                    },
133                    Binary => {
134                        let ca = s.binary().unwrap();
135                        fill_null_binary(ca, strategy).map(|ca| ca.into_series())
136                    },
137                    dt if dt.is_primitive_numeric() => {
138                        with_match_physical_numeric_polars_type!(dt, |$T| {
139                            let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
140                                fill_null_numeric(ca, strategy).map(|ca| ca.into_series())
141                        })
142                    },
143                    dt => {
144                        polars_bail!(InvalidOperation: "fill null strategy not yet supported for dtype: {}", dt)
145                    },
146                }?;
147                unsafe { out.from_physical_unchecked(logical_type) }
148            },
149        }
150    }
151}
152
153fn fill_forward_numeric<'a, T>(ca: &'a ChunkedArray<T>) -> ChunkedArray<T>
154where
155    T: PolarsDataType,
156    T::ZeroablePhysical<'a>: Copy,
157{
158    // Compute values.
159    let mut last = T::ZeroablePhysical::zeroed();
160    let values: Vec<T::ZeroablePhysical<'a>> = ca
161        .iter()
162        .map(|v| {
163            last = v.map(|v| v.into()).unwrap_or(last);
164            last
165        })
166        .collect_trusted();
167
168    // Compute bitmask.
169    let num_start_nulls = ca.first_non_null().unwrap_or(ca.len());
170    let mut bm = BitmapBuilder::with_capacity(ca.len());
171    bm.extend_constant(num_start_nulls, false);
172    bm.extend_constant(ca.len() - num_start_nulls, true);
173    ChunkedArray::from_chunk_iter_like(
174        ca,
175        [
176            T::Array::from_zeroable_vec(values, ca.dtype().to_arrow(CompatLevel::newest()))
177                .with_validity_typed(bm.into_opt_validity()),
178        ],
179    )
180}
181
182fn fill_backward_numeric<'a, T>(ca: &'a ChunkedArray<T>) -> ChunkedArray<T>
183where
184    T: PolarsDataType,
185    T::ZeroablePhysical<'a>: Copy,
186{
187    // Compute values.
188    let mut last = T::ZeroablePhysical::zeroed();
189    let values: Vec<T::ZeroablePhysical<'a>> = ca
190        .iter()
191        .rev()
192        .map(|v| {
193            last = v.map(|v| v.into()).unwrap_or(last);
194            last
195        })
196        .collect_reversed();
197
198    // Compute bitmask.
199    let num_end_nulls = ca
200        .last_non_null()
201        .map(|i| ca.len() - 1 - i)
202        .unwrap_or(ca.len());
203    let mut bm = BitmapBuilder::with_capacity(ca.len());
204    bm.extend_constant(ca.len() - num_end_nulls, true);
205    bm.extend_constant(num_end_nulls, false);
206    ChunkedArray::from_chunk_iter_like(
207        ca,
208        [
209            T::Array::from_zeroable_vec(values, ca.dtype().to_arrow(CompatLevel::newest()))
210                .with_validity_typed(bm.into_opt_validity()),
211        ],
212    )
213}
214
215fn fill_null_numeric<T>(
216    ca: &ChunkedArray<T>,
217    strategy: FillNullStrategy,
218) -> PolarsResult<ChunkedArray<T>>
219where
220    T: PolarsNumericType,
221    ChunkedArray<T>: ChunkAgg<T::Native>,
222{
223    // Nothing to fill.
224    let mut out = match strategy {
225        FillNullStrategy::Min => {
226            ca.fill_null_with_values(ChunkAgg::min(ca).ok_or_else(err_fill_null)?)?
227        },
228        FillNullStrategy::Max => {
229            ca.fill_null_with_values(ChunkAgg::max(ca).ok_or_else(err_fill_null)?)?
230        },
231        FillNullStrategy::Mean => ca.fill_null_with_values(
232            ca.mean()
233                .map(|v| NumCast::from(v).unwrap())
234                .ok_or_else(err_fill_null)?,
235        )?,
236        FillNullStrategy::One => return ca.fill_null_with_values(One::one()),
237        FillNullStrategy::Zero => return ca.fill_null_with_values(Zero::zero()),
238        FillNullStrategy::Forward(None) => fill_forward_numeric(ca),
239        FillNullStrategy::Backward(None) => fill_backward_numeric(ca),
240        // Handled earlier
241        FillNullStrategy::Forward(_) => unreachable!(),
242        FillNullStrategy::Backward(_) => unreachable!(),
243    };
244    out.rename(ca.name().clone());
245    Ok(out)
246}
247
248fn fill_with_gather<F: Fn(&Bitmap) -> Vec<IdxSize>>(
249    s: &Series,
250    bits_to_idx: F,
251) -> PolarsResult<Series> {
252    let s = s.rechunk();
253    let arr = s.chunks()[0].clone();
254    let validity = arr.validity().expect("nulls");
255
256    let idx = bits_to_idx(validity);
257
258    Ok(unsafe { s.take_slice_unchecked(&idx) })
259}
260
261fn fill_forward_gather(s: &Series) -> PolarsResult<Series> {
262    fill_with_gather(s, |validity| {
263        let mut last_valid = 0;
264        validity
265            .iter()
266            .enumerate_idx()
267            .map(|(i, v)| {
268                if v {
269                    last_valid = i;
270                    i
271                } else {
272                    last_valid
273                }
274            })
275            .collect::<Vec<_>>()
276    })
277}
278
279fn fill_forward_gather_limit(s: &Series, limit: IdxSize) -> PolarsResult<Series> {
280    fill_with_gather(s, |validity| {
281        let mut last_valid = 0;
282        let mut conseq_invalid_count = 0;
283        validity
284            .iter()
285            .enumerate_idx()
286            .map(|(i, v)| {
287                if v {
288                    last_valid = i;
289                    conseq_invalid_count = 0;
290                    i
291                } else if conseq_invalid_count < limit {
292                    conseq_invalid_count += 1;
293                    last_valid
294                } else {
295                    i
296                }
297            })
298            .collect::<Vec<_>>()
299    })
300}
301
302fn fill_backward_gather(s: &Series) -> PolarsResult<Series> {
303    fill_with_gather(s, |validity| {
304        let last = validity.len() as IdxSize - 1;
305        let mut last_valid = last;
306        unsafe {
307            validity
308                .iter()
309                .rev()
310                .enumerate_idx()
311                .map(|(i, v)| {
312                    if v {
313                        last_valid = last - i;
314                        last - i
315                    } else {
316                        last_valid
317                    }
318                })
319                .trust_my_length((last + 1) as usize)
320                .collect_reversed::<Vec<_>>()
321        }
322    })
323}
324
325fn fill_backward_gather_limit(s: &Series, limit: IdxSize) -> PolarsResult<Series> {
326    fill_with_gather(s, |validity| {
327        let last = validity.len() as IdxSize - 1;
328        let mut last_valid = last;
329        let mut conseq_invalid_count = 0;
330        unsafe {
331            validity
332                .iter()
333                .rev()
334                .enumerate_idx()
335                .map(|(i, v)| {
336                    if v {
337                        last_valid = last - i;
338                        conseq_invalid_count = 0;
339                        last - i
340                    } else if conseq_invalid_count < limit {
341                        conseq_invalid_count += 1;
342                        last_valid
343                    } else {
344                        last - i
345                    }
346                })
347                .trust_my_length((last + 1) as usize)
348                .collect_reversed()
349        }
350    })
351}
352
353#[cfg(feature = "dtype-categorical")]
354fn fill_null_cat<T: PolarsCategoricalType>(
355    ca: &CategoricalChunked<T>,
356    strategy: FillNullStrategy,
357) -> PolarsResult<Series>
358where
359    ChunkedArray<T::PolarsPhysical>: ChunkAgg<T::Native>,
360{
361    let cat = match strategy {
362        FillNullStrategy::Max => ca.max_categorical(),
363        FillNullStrategy::Min => ca.min_categorical(),
364        FillNullStrategy::Forward(_) => unreachable!(),
365        FillNullStrategy::Backward(_) => unreachable!(),
366        strat => {
367            polars_bail!(InvalidOperation: "fill-null strategy {:?} not supported for datatype {}", strat, ca.dtype)
368        },
369    }.ok_or_else(err_fill_null)?;
370    Ok(ca
371        .physical()
372        .fill_null_with_values(T::Native::from_cat(cat))?
373        .into_series())
374}
375
376fn fill_null_bool(ca: &BooleanChunked, strategy: FillNullStrategy) -> PolarsResult<Series> {
377    match strategy {
378        FillNullStrategy::Min => ca
379            .fill_null_with_values(ca.min().ok_or_else(err_fill_null)?)
380            .map(|ca| ca.into_series()),
381        FillNullStrategy::Max => ca
382            .fill_null_with_values(ca.max().ok_or_else(err_fill_null)?)
383            .map(|ca| ca.into_series()),
384        FillNullStrategy::Mean => polars_bail!(opq = mean, "Boolean"),
385        FillNullStrategy::One => ca.fill_null_with_values(true).map(|ca| ca.into_series()),
386        FillNullStrategy::Zero => ca.fill_null_with_values(false).map(|ca| ca.into_series()),
387        FillNullStrategy::Forward(_) => unreachable!(),
388        FillNullStrategy::Backward(_) => unreachable!(),
389    }
390}
391
392fn fill_null_binary(ca: &BinaryChunked, strategy: FillNullStrategy) -> PolarsResult<BinaryChunked> {
393    match strategy {
394        FillNullStrategy::Min => {
395            ca.fill_null_with_values(ca.min_binary().ok_or_else(err_fill_null)?)
396        },
397        FillNullStrategy::Max => {
398            ca.fill_null_with_values(ca.max_binary().ok_or_else(err_fill_null)?)
399        },
400        FillNullStrategy::Zero => ca.fill_null_with_values(&[]),
401        FillNullStrategy::Forward(_) => unreachable!(),
402        FillNullStrategy::Backward(_) => unreachable!(),
403        strat => polars_bail!(InvalidOperation: "fill-null strategy {:?} is not supported", strat),
404    }
405}
406
407impl<T> ChunkFillNullValue<T::Native> for ChunkedArray<T>
408where
409    T: PolarsNumericType,
410{
411    fn fill_null_with_values(&self, value: T::Native) -> PolarsResult<Self> {
412        Ok(self.apply_kernel(&|arr| Box::new(set_at_nulls(arr, value))))
413    }
414}
415
416impl ChunkFillNullValue<bool> for BooleanChunked {
417    fn fill_null_with_values(&self, value: bool) -> PolarsResult<Self> {
418        self.set(&self.is_null(), Some(value))
419    }
420}
421
422impl ChunkFillNullValue<&[u8]> for BinaryChunked {
423    fn fill_null_with_values(&self, value: &[u8]) -> PolarsResult<Self> {
424        self.set(&self.is_null(), Some(value))
425    }
426}
polars_core/chunked_array/ops/fill_null.rs

polars_core/chunked_array/ops/
fill_null.rs
