Skip to main content

polars_core/chunked_array/ops/
mod.rs

1//! Traits for miscellaneous operations on ChunkedArray
2use polars_arrow::offset::OffsetsBuffer;
3use polars_compute::rolling::QuantileMethod;
4
5use crate::prelude::*;
6
7pub(crate) mod aggregate;
8pub(crate) mod any_value;
9pub(crate) mod append;
10mod apply;
11#[cfg(feature = "approx_unique")]
12mod approx_n_unique;
13pub mod arity;
14pub mod binning;
15mod bit_repr;
16mod bits;
17#[cfg(feature = "bitwise")]
18mod bitwise_reduce;
19pub(crate) mod chunkops;
20pub(crate) mod compare_inner;
21#[cfg(feature = "dtype-decimal")]
22mod decimal;
23pub(crate) mod downcast;
24pub(crate) mod explode;
25mod explode_and_offsets;
26mod extend;
27pub mod fill_null;
28mod filter;
29pub mod float_sorted_arg_max;
30mod for_each;
31pub mod full;
32pub mod gather;
33mod nesting_utils;
34pub(crate) mod nulls;
35mod reverse;
36#[cfg(feature = "rolling_window")]
37pub(crate) mod rolling_window;
38pub mod row_encode;
39pub mod search_sorted;
40mod set;
41mod shift;
42pub mod sort;
43#[cfg(feature = "algorithm_group_by")]
44pub(crate) mod unique;
45#[cfg(feature = "zip_with")]
46pub mod zip;
47
48pub use bit_repr::reinterpret;
49pub use chunkops::_set_check_length;
50pub use explode::explode_with_row_offsets;
51pub use nesting_utils::ChunkNestingUtils;
52#[cfg(feature = "serde-lazy")]
53use serde::{Deserialize, Serialize};
54pub use sort::options::*;
55
56use crate::chunked_array::cast::CastOptions;
57use crate::series::{BitRepr, IsSorted};
58
59/// Transmute [`ChunkedArray`] to bit representation.
60/// This is useful in hashing context and reduces no.
61/// of compiled code paths.
62pub(crate) trait ToBitRepr {
63    fn to_bit_repr(&self) -> BitRepr;
64}
65
66pub trait ChunkAnyValue {
67    /// Get a single value. Beware this is slow.
68    /// If you need to use this slightly performant, cast Categorical to UInt32
69    ///
70    /// # Safety
71    /// Does not do any bounds checking.
72    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_>;
73
74    /// Get a single value. Beware this is slow.
75    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>>;
76}
77
78pub trait ChunkAnyValueBypassValidity {
79    /// Get a single value bypassing the validity map. Beware this is slow.
80    ///
81    /// # Safety
82    /// Does not do any bounds checking.
83    unsafe fn get_any_value_bypass_validity(&self, index: usize) -> AnyValue<'_>;
84}
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
87#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
88#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
89pub struct ExplodeOptions {
90    /// Explode an empty list into a `null`.
91    pub empty_as_null: bool,
92    /// Explode a `null` into a `null`.
93    pub keep_nulls: bool,
94}
95
96/// Explode/flatten a List or String Series
97pub trait ChunkExplode {
98    fn explode(&self, options: ExplodeOptions) -> PolarsResult<Series> {
99        self.explode_and_offsets(options).map(|t| t.0)
100    }
101    fn offsets(&self) -> PolarsResult<OffsetsBuffer<i64>>;
102    fn explode_and_offsets(
103        &self,
104        options: ExplodeOptions,
105    ) -> PolarsResult<(Series, OffsetsBuffer<i64>)>;
106}
107
108/// This differs from ChunkWindowCustom and ChunkWindow
109/// by not using a fold aggregator, but reusing a `Series` wrapper and calling `Series` aggregators.
110/// This likely is a bit slower than ChunkWindow
111#[cfg(feature = "rolling_window")]
112pub trait ChunkRollApply: AsRefDataType {
113    fn rolling_map(
114        &self,
115        f: &dyn Fn(&Series) -> PolarsResult<Series>,
116        options: RollingOptionsFixedWindow,
117    ) -> PolarsResult<Series>
118    where
119        Self: Sized;
120}
121
122pub trait ChunkTake<Idx: ?Sized>: ChunkTakeUnchecked<Idx> {
123    /// Gather values from ChunkedArray by index.
124    fn take(&self, indices: &Idx) -> PolarsResult<Self>
125    where
126        Self: Sized;
127}
128
129pub trait ChunkTakeUnchecked<Idx: ?Sized> {
130    /// Gather values from ChunkedArray by index.
131    ///
132    /// # Safety
133    /// The non-null indices must be valid.
134    unsafe fn take_unchecked(&self, indices: &Idx) -> Self;
135}
136
137/// Create a `ChunkedArray` with new values by index or by boolean mask.
138///
139/// Note that these operations clone data. This is however the only way we can modify at mask or
140/// index level as the underlying Arrow arrays are immutable.
141pub trait ChunkSet<'a, A, B> {
142    /// Set the values at indexes `idx` to some optional value `Option<T>`.
143    ///
144    /// # Example
145    ///
146    /// ```rust
147    /// # use polars_core::prelude::*;
148    /// let ca = UInt32Chunked::new("a".into(), &[1, 2, 3]);
149    /// let new = ca.scatter_single(vec![0, 1], Some(10)).unwrap();
150    ///
151    /// assert_eq!(Vec::from(&new), &[Some(10), Some(10), Some(3)]);
152    /// ```
153    fn scatter_single<I: IntoIterator<Item = IdxSize>>(
154        &'a self,
155        idx: I,
156        opt_value: Option<A>,
157    ) -> PolarsResult<Self>
158    where
159        Self: Sized;
160
161    /// Set the values at indexes `idx` by applying a closure to these values.
162    ///
163    /// # Example
164    ///
165    /// ```rust
166    /// # use polars_core::prelude::*;
167    /// let ca = Int32Chunked::new("a".into(), &[1, 2, 3]);
168    /// let new = ca.scatter_with(vec![0, 1], |opt_v| opt_v.map(|v| v - 5)).unwrap();
169    ///
170    /// assert_eq!(Vec::from(&new), &[Some(-4), Some(-3), Some(3)]);
171    /// ```
172    fn scatter_with<I: IntoIterator<Item = IdxSize>, F>(
173        &'a self,
174        idx: I,
175        f: F,
176    ) -> PolarsResult<Self>
177    where
178        Self: Sized,
179        F: Fn(Option<A>) -> Option<B>;
180    /// Set the values where the mask evaluates to `true` to some optional value `Option<T>`.
181    ///
182    /// # Example
183    ///
184    /// ```rust
185    /// # use polars_core::prelude::*;
186    /// let ca = Int32Chunked::new("a".into(), &[1, 2, 3]);
187    /// let mask = BooleanChunked::new("mask".into(), &[false, true, false]);
188    /// let new = ca.set(&mask, Some(5)).unwrap();
189    /// assert_eq!(Vec::from(&new), &[Some(1), Some(5), Some(3)]);
190    /// ```
191    fn set(&'a self, mask: &BooleanChunked, opt_value: Option<A>) -> PolarsResult<Self>
192    where
193        Self: Sized;
194}
195
196/// Cast `ChunkedArray<T>` to `ChunkedArray<N>`
197pub trait ChunkCast {
198    /// Cast a [`ChunkedArray`] to [`DataType`]
199    fn cast(&self, dtype: &DataType) -> PolarsResult<Series> {
200        self.cast_with_options(dtype, CastOptions::NonStrict)
201    }
202
203    /// Cast a [`ChunkedArray`] to [`DataType`]
204    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series>;
205
206    /// Does not check if the cast is a valid one and may over/underflow
207    ///
208    /// # Safety
209    /// - This doesn't do utf8 validation checking when casting from binary
210    /// - This doesn't do categorical bound checking when casting from UInt32
211    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series>;
212}
213
214/// Fastest way to do elementwise operations on a [`ChunkedArray<T>`] when the operation is cheaper than
215/// branching due to null checking.
216pub trait ChunkApply<'a, T> {
217    type FuncRet;
218
219    /// Apply a closure elementwise. This is fastest when the null check branching is more expensive
220    /// than the closure application. Often it is.
221    ///
222    /// Null values remain null.
223    ///
224    /// # Example
225    ///
226    /// ```
227    /// use polars_core::prelude::*;
228    /// fn double(ca: &UInt32Chunked) -> UInt32Chunked {
229    ///     ca.apply_values(|v| v * 2)
230    /// }
231    /// ```
232    #[must_use]
233    fn apply_values<F>(&'a self, f: F) -> Self
234    where
235        F: Fn(T) -> Self::FuncRet + Copy;
236
237    /// Apply a closure elementwise including null values.
238    #[must_use]
239    fn apply<F>(&'a self, f: F) -> Self
240    where
241        F: Fn(Option<T>) -> Option<Self::FuncRet> + Copy;
242
243    /// Apply a closure elementwise and write results to a mutable slice.
244    fn apply_to_slice<F, S>(&'a self, f: F, slice: &mut [S])
245    // (value of chunkedarray, value of slice) -> value of slice
246    where
247        F: Fn(Option<T>, &S) -> S;
248}
249
250/// Aggregation operations.
251pub trait ChunkAgg<T> {
252    /// Aggregate the sum of the ChunkedArray.
253    /// Returns `None` if not implemented for `T`.
254    /// If the array is empty, `0` is returned
255    fn sum(&self) -> Option<T> {
256        None
257    }
258
259    fn _sum_as_f64(&self) -> f64;
260
261    fn min(&self) -> Option<T> {
262        None
263    }
264
265    /// Returns the maximum value in the array, according to the natural order.
266    /// Returns `None` if the array is empty or only contains null values.
267    fn max(&self) -> Option<T> {
268        None
269    }
270
271    fn min_max(&self) -> Option<(T, T)> {
272        Some((self.min()?, self.max()?))
273    }
274
275    /// Returns the mean value in the array.
276    /// Returns `None` if the array is empty or only contains null values.
277    fn mean(&self) -> Option<f64> {
278        None
279    }
280}
281
282/// Quantile and median aggregation.
283pub trait ChunkQuantile<T> {
284    /// Returns the mean value in the array.
285    /// Returns `None` if the array is empty or only contains null values.
286    fn median(&self) -> Option<T> {
287        None
288    }
289    /// Aggregate a given quantile of the ChunkedArray.
290    /// Returns `None` if the array is empty or only contains null values.
291    fn quantile(&self, _quantile: f64, _method: QuantileMethod) -> PolarsResult<Option<T>> {
292        Ok(None)
293    }
294    /// Aggregate a given set of quantiles of the ChunkedArray.
295    /// Returns `None` if the array is empty or only contains null values.
296    fn quantiles(&self, quantiles: &[f64], _method: QuantileMethod) -> PolarsResult<Vec<Option<T>>>
297    where
298        T: Clone,
299    {
300        Ok(vec![None; quantiles.len()])
301    }
302}
303
304/// Variance and standard deviation aggregation.
305pub trait ChunkVar {
306    /// Compute the variance of this ChunkedArray/Series.
307    fn var(&self, _ddof: u8) -> Option<f64> {
308        None
309    }
310
311    /// Compute the standard deviation of this ChunkedArray/Series.
312    fn std(&self, _ddof: u8) -> Option<f64> {
313        None
314    }
315}
316
317/// Bitwise Reduction Operations.
318#[cfg(feature = "bitwise")]
319pub trait ChunkBitwiseReduce {
320    type Physical;
321
322    fn and_reduce(&self) -> Option<Self::Physical>;
323    fn or_reduce(&self) -> Option<Self::Physical>;
324    fn xor_reduce(&self) -> Option<Self::Physical>;
325}
326
327/// Compare [`Series`] and [`ChunkedArray`]'s and get a `boolean` mask that
328/// can be used to filter rows.
329///
330/// # Example
331///
332/// ```
333/// use polars_core::prelude::*;
334/// fn filter_all_ones(df: &DataFrame) -> PolarsResult<DataFrame> {
335///     let mask = df
336///     .column("column_a")?
337///     .as_materialized_series()
338///     .equal(1)?;
339///
340///     df.filter(&mask)
341/// }
342/// ```
343pub trait ChunkCompareEq<Rhs> {
344    type Item;
345
346    /// Check for equality.
347    fn equal(&self, rhs: Rhs) -> Self::Item;
348
349    /// Check for equality where `None == None`.
350    fn equal_missing(&self, rhs: Rhs) -> Self::Item;
351
352    /// Check for inequality.
353    fn not_equal(&self, rhs: Rhs) -> Self::Item;
354
355    /// Check for inequality where `None == None`.
356    fn not_equal_missing(&self, rhs: Rhs) -> Self::Item;
357}
358
359/// Compare [`Series`] and [`ChunkedArray`]'s using inequality operators (`<`, `>=`, etc.) and get
360/// a `boolean` mask that can be used to filter rows.
361pub trait ChunkCompareIneq<Rhs> {
362    type Item;
363
364    /// Greater than comparison.
365    fn gt(&self, rhs: Rhs) -> Self::Item;
366
367    /// Greater than or equal comparison.
368    fn gt_eq(&self, rhs: Rhs) -> Self::Item;
369
370    /// Less than comparison.
371    fn lt(&self, rhs: Rhs) -> Self::Item;
372
373    /// Less than or equal comparison
374    fn lt_eq(&self, rhs: Rhs) -> Self::Item;
375}
376
377/// Get unique values in a `ChunkedArray`
378pub trait ChunkUnique {
379    // We don't return Self to be able to use AutoRef specialization
380    /// Get unique values of a ChunkedArray
381    fn unique(&self) -> PolarsResult<Self>
382    where
383        Self: Sized;
384
385    /// Get first index of the unique values in a `ChunkedArray`.
386    /// This Vec is sorted.
387    fn arg_unique(&self) -> PolarsResult<IdxCa>;
388
389    /// Number of unique values in the `ChunkedArray`
390    fn n_unique(&self) -> PolarsResult<usize> {
391        self.arg_unique().map(|v| v.len())
392    }
393
394    /// Get dense ids for each unique value.
395    ///
396    /// Returns: (n_unique, unique_ids)
397    fn unique_id(&self) -> PolarsResult<(IdxSize, Vec<IdxSize>)>;
398}
399
400#[cfg(feature = "approx_unique")]
401pub trait ChunkApproxNUnique {
402    fn approx_n_unique(&self) -> IdxSize;
403}
404
405/// Sort operations on `ChunkedArray`.
406pub trait ChunkSort<T: PolarsDataType> {
407    #[allow(unused_variables)]
408    fn sort_with(&self, options: SortOptions) -> ChunkedArray<T>;
409
410    /// Returned a sorted `ChunkedArray`.
411    fn sort(&self, descending: bool) -> ChunkedArray<T>;
412
413    /// Retrieve the indexes needed to sort this array.
414    fn arg_sort(&self, options: SortOptions) -> IdxCa;
415
416    /// Retrieve the indexes need to sort this and the other arrays.
417    #[allow(unused_variables)]
418    fn arg_sort_multiple(
419        &self,
420        by: &[Column],
421        _options: &SortMultipleOptions,
422    ) -> PolarsResult<IdxCa> {
423        polars_bail!(opq = arg_sort_multiple, T::get_static_dtype());
424    }
425}
426
427pub type FillNullLimit = Option<IdxSize>;
428
429#[derive(Copy, Clone, Debug, PartialEq, Hash, Eq)]
430#[cfg_attr(feature = "serde-lazy", derive(Serialize, Deserialize))]
431#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
432pub enum FillNullStrategy {
433    /// previous value in array
434    Backward(FillNullLimit),
435    /// next value in array
436    Forward(FillNullLimit),
437    /// mean value of array
438    Mean,
439    /// minimal value in array
440    Min,
441    /// maximum value in array
442    Max,
443    /// replace with the value zero
444    Zero,
445    /// replace with the value one
446    One,
447}
448
449impl FillNullStrategy {
450    pub fn is_elementwise(&self) -> bool {
451        matches!(self, Self::One | Self::Zero)
452    }
453}
454
455/// Replace None values with a value
456pub trait ChunkFillNullValue<T> {
457    /// Replace None values with a give value `T`.
458    fn fill_null_with_values(&self, value: T) -> PolarsResult<Self>
459    where
460        Self: Sized;
461}
462
463/// Fill a ChunkedArray with one value.
464pub trait ChunkFull<T> {
465    /// Create a ChunkedArray with a single value.
466    fn full(name: PlSmallStr, value: T, length: usize) -> Self
467    where
468        Self: Sized;
469}
470
471pub trait ChunkFullNull {
472    fn full_null(_name: PlSmallStr, _length: usize) -> Self
473    where
474        Self: Sized;
475}
476
477/// Reverse a [`ChunkedArray<T>`]
478pub trait ChunkReverse {
479    /// Return a reversed version of this array.
480    fn reverse(&self) -> Self;
481}
482
483/// Filter values by a boolean mask.
484pub trait ChunkFilter<T: PolarsDataType> {
485    /// Filter values in the ChunkedArray with a boolean mask.
486    ///
487    /// ```rust
488    /// # use polars_core::prelude::*;
489    /// let array = Int32Chunked::new("array".into(), &[1, 2, 3]);
490    /// let mask = BooleanChunked::new("mask".into(), &[true, false, true]);
491    ///
492    /// let filtered = array.filter(&mask).unwrap();
493    /// assert_eq!(Vec::from(&filtered), [Some(1), Some(3)])
494    /// ```
495    fn filter(&self, filter: &BooleanChunked) -> PolarsResult<ChunkedArray<T>>
496    where
497        Self: Sized;
498}
499
500/// Create a new ChunkedArray filled with values at that index.
501pub trait ChunkExpandAtIndex<T: PolarsDataType> {
502    /// Create a new ChunkedArray filled with values at that index.
503    fn new_from_index(&self, index: usize, length: usize) -> ChunkedArray<T>;
504}
505
506macro_rules! impl_chunk_expand {
507    ($self:ident, $length:ident, $index:ident) => {{
508        if $self.is_empty() {
509            return $self.clone();
510        }
511        let opt_val = $self.get($index);
512        match opt_val {
513            Some(val) => ChunkedArray::full($self.name().clone(), val, $length),
514            None => ChunkedArray::full_null($self.name().clone(), $length),
515        }
516    }};
517}
518
519impl<T: PolarsNumericType> ChunkExpandAtIndex<T> for ChunkedArray<T>
520where
521    ChunkedArray<T>: ChunkFull<T::Native>,
522{
523    fn new_from_index(&self, index: usize, length: usize) -> ChunkedArray<T> {
524        let mut out = impl_chunk_expand!(self, length, index);
525        out.set_sorted_flag(IsSorted::Ascending);
526        out
527    }
528}
529
530impl ChunkExpandAtIndex<BooleanType> for BooleanChunked {
531    fn new_from_index(&self, index: usize, length: usize) -> BooleanChunked {
532        let mut out = impl_chunk_expand!(self, length, index);
533        out.set_sorted_flag(IsSorted::Ascending);
534        out
535    }
536}
537
538impl ChunkExpandAtIndex<StringType> for StringChunked {
539    fn new_from_index(&self, index: usize, length: usize) -> StringChunked {
540        let mut out = impl_chunk_expand!(self, length, index);
541        out.set_sorted_flag(IsSorted::Ascending);
542        out
543    }
544}
545
546impl ChunkExpandAtIndex<BinaryType> for BinaryChunked {
547    fn new_from_index(&self, index: usize, length: usize) -> BinaryChunked {
548        let mut out = impl_chunk_expand!(self, length, index);
549        out.set_sorted_flag(IsSorted::Ascending);
550        out
551    }
552}
553
554impl ChunkExpandAtIndex<BinaryOffsetType> for BinaryOffsetChunked {
555    fn new_from_index(&self, index: usize, length: usize) -> BinaryOffsetChunked {
556        let mut out = impl_chunk_expand!(self, length, index);
557        out.set_sorted_flag(IsSorted::Ascending);
558        out
559    }
560}
561
562impl ChunkExpandAtIndex<ListType> for ListChunked {
563    fn new_from_index(&self, index: usize, length: usize) -> ListChunked {
564        let opt_val = self.get_as_series(index);
565        match opt_val {
566            Some(val) => {
567                let mut ca = ListChunked::full(self.name().clone(), &val, length);
568                unsafe { ca.to_logical(self.inner_dtype().clone()) };
569                ca
570            },
571            None => {
572                ListChunked::full_null_with_dtype(self.name().clone(), length, self.inner_dtype())
573            },
574        }
575    }
576}
577
578#[cfg(feature = "dtype-struct")]
579impl ChunkExpandAtIndex<StructType> for StructChunked {
580    fn new_from_index(&self, index: usize, length: usize) -> ChunkedArray<StructType> {
581        let (chunk_idx, idx) = self.index_to_chunked_index(index);
582        let chunk = self.downcast_chunks().get(chunk_idx).unwrap();
583        let chunk = if chunk.is_null(idx) {
584            new_null_array(chunk.dtype().clone(), length)
585        } else {
586            let values = chunk
587                .values()
588                .iter()
589                .map(|arr| {
590                    let s = Series::try_from((PlSmallStr::EMPTY, arr.clone())).unwrap();
591                    let s = s.new_from_index(idx, length);
592                    s.chunks()[0].clone()
593                })
594                .collect::<Vec<_>>();
595
596            StructArray::new(chunk.dtype().clone(), length, values, None).boxed()
597        };
598
599        // SAFETY: chunks are from self.
600        unsafe { self.copy_with_chunks(vec![chunk]) }
601    }
602}
603
604#[cfg(feature = "dtype-array")]
605impl ChunkExpandAtIndex<FixedSizeListType> for ArrayChunked {
606    fn new_from_index(&self, index: usize, length: usize) -> ArrayChunked {
607        let opt_val = self.get_as_series(index);
608        match opt_val {
609            Some(val) => {
610                let mut ca = ArrayChunked::full(self.name().clone(), &val, length);
611                unsafe { ca.to_logical(self.inner_dtype().clone()) };
612                ca
613            },
614            None => ArrayChunked::full_null_with_dtype(
615                self.name().clone(),
616                length,
617                self.inner_dtype(),
618                self.width(),
619            ),
620        }
621    }
622}
623
624#[cfg(feature = "object")]
625impl<T: PolarsObject> ChunkExpandAtIndex<ObjectType<T>> for ObjectChunked<T> {
626    fn new_from_index(&self, index: usize, length: usize) -> ObjectChunked<T> {
627        let opt_val = self.get(index);
628        match opt_val {
629            Some(val) => ObjectChunked::<T>::full(self.name().clone(), val.clone(), length),
630            None => ObjectChunked::<T>::full_null(self.name().clone(), length),
631        }
632    }
633}
634
635/// Shift the values of a [`ChunkedArray`] by a number of periods.
636pub trait ChunkShiftFill<T: PolarsDataType, V> {
637    /// Shift the values by a given period and fill the parts that will be empty due to this operation
638    /// with `fill_value`.
639    fn shift_and_fill(&self, periods: i64, fill_value: V) -> ChunkedArray<T>;
640}
641
642pub trait ChunkShift<T: PolarsDataType> {
643    fn shift(&self, periods: i64) -> ChunkedArray<T>;
644}
645
646/// Combine two [`ChunkedArray`] based on some predicate.
647pub trait ChunkZip<T: PolarsDataType> {
648    /// Create a new ChunkedArray with values from self where the mask evaluates `true` and values
649    /// from `other` where the mask evaluates `false`
650    fn zip_with(
651        &self,
652        mask: &BooleanChunked,
653        other: &ChunkedArray<T>,
654    ) -> PolarsResult<ChunkedArray<T>>;
655}
656
657/// Apply kernels on the arrow array chunks in a ChunkedArray.
658pub trait ChunkApplyKernel<A: Array> {
659    /// Apply kernel and return result as a new ChunkedArray.
660    #[must_use]
661    fn apply_kernel(&self, f: &dyn Fn(&A) -> ArrayRef) -> Self;
662
663    /// Apply a kernel that outputs an array of different type.
664    fn apply_kernel_cast<S>(&self, f: &dyn Fn(&A) -> ArrayRef) -> ChunkedArray<S>
665    where
666        S: PolarsDataType;
667}