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

polars_core/chunked_array/ops/
mod.rs
