polars_core/chunked_array/ops/
mod.rs1use 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;
14mod bit_repr;
15mod bits;
16#[cfg(feature = "bitwise")]
17mod bitwise_reduce;
18pub(crate) mod chunkops;
19pub(crate) mod compare_inner;
20#[cfg(feature = "dtype-decimal")]
21mod decimal;
22pub(crate) mod downcast;
23pub(crate) mod explode;
24mod explode_and_offsets;
25mod extend;
26pub mod fill_null;
27mod filter;
28pub mod float_sorted_arg_max;
29mod for_each;
30pub mod full;
31pub mod gather;
32mod nesting_utils;
33pub(crate) mod nulls;
34mod reverse;
35#[cfg(feature = "rolling_window")]
36pub(crate) mod rolling_window;
37pub mod row_encode;
38pub mod search_sorted;
39mod set;
40mod shift;
41pub mod sort;
42#[cfg(feature = "algorithm_group_by")]
43pub(crate) mod unique;
44#[cfg(feature = "zip_with")]
45pub mod zip;
46
47pub use chunkops::_set_check_length;
48pub use nesting_utils::ChunkNestingUtils;
49#[cfg(feature = "serde-lazy")]
50use serde::{Deserialize, Serialize};
51pub use sort::options::*;
52
53use crate::chunked_array::cast::CastOptions;
54use crate::series::{BitRepr, IsSorted};
55pub trait Reinterpret {
56 fn reinterpret_signed(&self) -> Series {
57 unimplemented!()
58 }
59
60 fn reinterpret_unsigned(&self) -> Series {
61 unimplemented!()
62 }
63}
64
65pub(crate) trait ToBitRepr {
69 fn to_bit_repr(&self) -> BitRepr;
70}
71
72pub trait ChunkAnyValue {
73 unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_>;
79
80 fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>>;
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
85#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
86#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
87pub struct ExplodeOptions {
88 pub empty_as_null: bool,
90 pub keep_nulls: bool,
92}
93
94pub trait ChunkExplode {
96 fn explode(&self, options: ExplodeOptions) -> PolarsResult<Series> {
97 self.explode_and_offsets(options).map(|t| t.0)
98 }
99 fn offsets(&self) -> PolarsResult<OffsetsBuffer<i64>>;
100 fn explode_and_offsets(
101 &self,
102 options: ExplodeOptions,
103 ) -> PolarsResult<(Series, OffsetsBuffer<i64>)>;
104}
105
106pub trait ChunkBytes {
107 fn to_byte_slices(&self) -> Vec<&[u8]>;
108}
109
110#[cfg(feature = "rolling_window")]
114pub trait ChunkRollApply: AsRefDataType {
115 fn rolling_map(
116 &self,
117 f: &dyn Fn(&Series) -> PolarsResult<Series>,
118 options: RollingOptionsFixedWindow,
119 ) -> PolarsResult<Series>
120 where
121 Self: Sized;
122}
123
124pub trait ChunkTake<Idx: ?Sized>: ChunkTakeUnchecked<Idx> {
125 fn take(&self, indices: &Idx) -> PolarsResult<Self>
127 where
128 Self: Sized;
129}
130
131pub trait ChunkTakeUnchecked<Idx: ?Sized> {
132 unsafe fn take_unchecked(&self, indices: &Idx) -> Self;
137}
138
139pub trait ChunkSet<'a, A, B> {
144 fn scatter_single<I: IntoIterator<Item = IdxSize>>(
156 &'a self,
157 idx: I,
158 opt_value: Option<A>,
159 ) -> PolarsResult<Self>
160 where
161 Self: Sized;
162
163 fn scatter_with<I: IntoIterator<Item = IdxSize>, F>(
175 &'a self,
176 idx: I,
177 f: F,
178 ) -> PolarsResult<Self>
179 where
180 Self: Sized,
181 F: Fn(Option<A>) -> Option<B>;
182 fn set(&'a self, mask: &BooleanChunked, opt_value: Option<A>) -> PolarsResult<Self>
194 where
195 Self: Sized;
196}
197
198pub trait ChunkCast {
200 fn cast(&self, dtype: &DataType) -> PolarsResult<Series> {
202 self.cast_with_options(dtype, CastOptions::NonStrict)
203 }
204
205 fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series>;
207
208 unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series>;
214}
215
216pub trait ChunkApply<'a, T> {
219 type FuncRet;
220
221 #[must_use]
235 fn apply_values<F>(&'a self, f: F) -> Self
236 where
237 F: Fn(T) -> Self::FuncRet + Copy;
238
239 #[must_use]
241 fn apply<F>(&'a self, f: F) -> Self
242 where
243 F: Fn(Option<T>) -> Option<Self::FuncRet> + Copy;
244
245 fn apply_to_slice<F, S>(&'a self, f: F, slice: &mut [S])
247 where
249 F: Fn(Option<T>, &S) -> S;
250}
251
252pub trait ChunkAgg<T> {
254 fn sum(&self) -> Option<T> {
258 None
259 }
260
261 fn _sum_as_f64(&self) -> f64;
262
263 fn min(&self) -> Option<T> {
264 None
265 }
266
267 fn max(&self) -> Option<T> {
270 None
271 }
272
273 fn min_max(&self) -> Option<(T, T)> {
274 Some((self.min()?, self.max()?))
275 }
276
277 fn mean(&self) -> Option<f64> {
280 None
281 }
282}
283
284pub trait ChunkQuantile<T> {
286 fn median(&self) -> Option<T> {
289 None
290 }
291 fn quantile(&self, _quantile: f64, _method: QuantileMethod) -> PolarsResult<Option<T>> {
294 Ok(None)
295 }
296}
297
298pub trait ChunkVar {
300 fn var(&self, _ddof: u8) -> Option<f64> {
302 None
303 }
304
305 fn std(&self, _ddof: u8) -> Option<f64> {
307 None
308 }
309}
310
311#[cfg(feature = "bitwise")]
313pub trait ChunkBitwiseReduce {
314 type Physical;
315
316 fn and_reduce(&self) -> Option<Self::Physical>;
317 fn or_reduce(&self) -> Option<Self::Physical>;
318 fn xor_reduce(&self) -> Option<Self::Physical>;
319}
320
321pub trait ChunkCompareEq<Rhs> {
338 type Item;
339
340 fn equal(&self, rhs: Rhs) -> Self::Item;
342
343 fn equal_missing(&self, rhs: Rhs) -> Self::Item;
345
346 fn not_equal(&self, rhs: Rhs) -> Self::Item;
348
349 fn not_equal_missing(&self, rhs: Rhs) -> Self::Item;
351}
352
353pub trait ChunkCompareIneq<Rhs> {
356 type Item;
357
358 fn gt(&self, rhs: Rhs) -> Self::Item;
360
361 fn gt_eq(&self, rhs: Rhs) -> Self::Item;
363
364 fn lt(&self, rhs: Rhs) -> Self::Item;
366
367 fn lt_eq(&self, rhs: Rhs) -> Self::Item;
369}
370
371pub trait ChunkUnique {
373 fn unique(&self) -> PolarsResult<Self>
376 where
377 Self: Sized;
378
379 fn arg_unique(&self) -> PolarsResult<IdxCa>;
382
383 fn n_unique(&self) -> PolarsResult<usize> {
385 self.arg_unique().map(|v| v.len())
386 }
387
388 fn unique_id(&self) -> PolarsResult<(IdxSize, Vec<IdxSize>)>;
392}
393
394#[cfg(feature = "approx_unique")]
395pub trait ChunkApproxNUnique {
396 fn approx_n_unique(&self) -> IdxSize;
397}
398
399pub trait ChunkSort<T: PolarsDataType> {
401 #[allow(unused_variables)]
402 fn sort_with(&self, options: SortOptions) -> ChunkedArray<T>;
403
404 fn sort(&self, descending: bool) -> ChunkedArray<T>;
406
407 fn arg_sort(&self, options: SortOptions) -> IdxCa;
409
410 #[allow(unused_variables)]
412 fn arg_sort_multiple(
413 &self,
414 by: &[Column],
415 _options: &SortMultipleOptions,
416 ) -> PolarsResult<IdxCa> {
417 polars_bail!(opq = arg_sort_multiple, T::get_static_dtype());
418 }
419}
420
421pub type FillNullLimit = Option<IdxSize>;
422
423#[derive(Copy, Clone, Debug, PartialEq, Hash)]
424#[cfg_attr(feature = "serde-lazy", derive(Serialize, Deserialize))]
425#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
426pub enum FillNullStrategy {
427 Backward(FillNullLimit),
429 Forward(FillNullLimit),
431 Mean,
433 Min,
435 Max,
437 Zero,
439 One,
441}
442
443impl FillNullStrategy {
444 pub fn is_elementwise(&self) -> bool {
445 matches!(self, Self::One | Self::Zero)
446 }
447}
448
449pub trait ChunkFillNullValue<T> {
451 fn fill_null_with_values(&self, value: T) -> PolarsResult<Self>
453 where
454 Self: Sized;
455}
456
457pub trait ChunkFull<T> {
459 fn full(name: PlSmallStr, value: T, length: usize) -> Self
461 where
462 Self: Sized;
463}
464
465pub trait ChunkFullNull {
466 fn full_null(_name: PlSmallStr, _length: usize) -> Self
467 where
468 Self: Sized;
469}
470
471pub trait ChunkReverse {
473 fn reverse(&self) -> Self;
475}
476
477pub trait ChunkFilter<T: PolarsDataType> {
479 fn filter(&self, filter: &BooleanChunked) -> PolarsResult<ChunkedArray<T>>
490 where
491 Self: Sized;
492}
493
494pub trait ChunkExpandAtIndex<T: PolarsDataType> {
496 fn new_from_index(&self, index: usize, length: usize) -> ChunkedArray<T>;
498}
499
500macro_rules! impl_chunk_expand {
501 ($self:ident, $length:ident, $index:ident) => {{
502 if $self.is_empty() {
503 return $self.clone();
504 }
505 let opt_val = $self.get($index);
506 match opt_val {
507 Some(val) => ChunkedArray::full($self.name().clone(), val, $length),
508 None => ChunkedArray::full_null($self.name().clone(), $length),
509 }
510 }};
511}
512
513impl<T: PolarsNumericType> ChunkExpandAtIndex<T> for ChunkedArray<T>
514where
515 ChunkedArray<T>: ChunkFull<T::Native>,
516{
517 fn new_from_index(&self, index: usize, length: usize) -> ChunkedArray<T> {
518 let mut out = impl_chunk_expand!(self, length, index);
519 out.set_sorted_flag(IsSorted::Ascending);
520 out
521 }
522}
523
524impl ChunkExpandAtIndex<BooleanType> for BooleanChunked {
525 fn new_from_index(&self, index: usize, length: usize) -> BooleanChunked {
526 let mut out = impl_chunk_expand!(self, length, index);
527 out.set_sorted_flag(IsSorted::Ascending);
528 out
529 }
530}
531
532impl ChunkExpandAtIndex<StringType> for StringChunked {
533 fn new_from_index(&self, index: usize, length: usize) -> StringChunked {
534 let mut out = impl_chunk_expand!(self, length, index);
535 out.set_sorted_flag(IsSorted::Ascending);
536 out
537 }
538}
539
540impl ChunkExpandAtIndex<BinaryType> for BinaryChunked {
541 fn new_from_index(&self, index: usize, length: usize) -> BinaryChunked {
542 let mut out = impl_chunk_expand!(self, length, index);
543 out.set_sorted_flag(IsSorted::Ascending);
544 out
545 }
546}
547
548impl ChunkExpandAtIndex<BinaryOffsetType> for BinaryOffsetChunked {
549 fn new_from_index(&self, index: usize, length: usize) -> BinaryOffsetChunked {
550 let mut out = impl_chunk_expand!(self, length, index);
551 out.set_sorted_flag(IsSorted::Ascending);
552 out
553 }
554}
555
556impl ChunkExpandAtIndex<ListType> for ListChunked {
557 fn new_from_index(&self, index: usize, length: usize) -> ListChunked {
558 let opt_val = self.get_as_series(index);
559 match opt_val {
560 Some(val) => {
561 let mut ca = ListChunked::full(self.name().clone(), &val, length);
562 unsafe { ca.to_logical(self.inner_dtype().clone()) };
563 ca
564 },
565 None => {
566 ListChunked::full_null_with_dtype(self.name().clone(), length, self.inner_dtype())
567 },
568 }
569 }
570}
571
572#[cfg(feature = "dtype-struct")]
573impl ChunkExpandAtIndex<StructType> for StructChunked {
574 fn new_from_index(&self, index: usize, length: usize) -> ChunkedArray<StructType> {
575 let (chunk_idx, idx) = self.index_to_chunked_index(index);
576 let chunk = self.downcast_chunks().get(chunk_idx).unwrap();
577 let chunk = if chunk.is_null(idx) {
578 new_null_array(chunk.dtype().clone(), length)
579 } else {
580 let values = chunk
581 .values()
582 .iter()
583 .map(|arr| {
584 let s = Series::try_from((PlSmallStr::EMPTY, arr.clone())).unwrap();
585 let s = s.new_from_index(idx, length);
586 s.chunks()[0].clone()
587 })
588 .collect::<Vec<_>>();
589
590 StructArray::new(chunk.dtype().clone(), length, values, None).boxed()
591 };
592
593 unsafe { self.copy_with_chunks(vec![chunk]) }
595 }
596}
597
598#[cfg(feature = "dtype-array")]
599impl ChunkExpandAtIndex<FixedSizeListType> for ArrayChunked {
600 fn new_from_index(&self, index: usize, length: usize) -> ArrayChunked {
601 let opt_val = self.get_as_series(index);
602 match opt_val {
603 Some(val) => {
604 let mut ca = ArrayChunked::full(self.name().clone(), &val, length);
605 unsafe { ca.to_logical(self.inner_dtype().clone()) };
606 ca
607 },
608 None => ArrayChunked::full_null_with_dtype(
609 self.name().clone(),
610 length,
611 self.inner_dtype(),
612 self.width(),
613 ),
614 }
615 }
616}
617
618#[cfg(feature = "object")]
619impl<T: PolarsObject> ChunkExpandAtIndex<ObjectType<T>> for ObjectChunked<T> {
620 fn new_from_index(&self, index: usize, length: usize) -> ObjectChunked<T> {
621 let opt_val = self.get(index);
622 match opt_val {
623 Some(val) => ObjectChunked::<T>::full(self.name().clone(), val.clone(), length),
624 None => ObjectChunked::<T>::full_null(self.name().clone(), length),
625 }
626 }
627}
628
629pub trait ChunkShiftFill<T: PolarsDataType, V> {
631 fn shift_and_fill(&self, periods: i64, fill_value: V) -> ChunkedArray<T>;
634}
635
636pub trait ChunkShift<T: PolarsDataType> {
637 fn shift(&self, periods: i64) -> ChunkedArray<T>;
638}
639
640pub trait ChunkZip<T: PolarsDataType> {
642 fn zip_with(
645 &self,
646 mask: &BooleanChunked,
647 other: &ChunkedArray<T>,
648 ) -> PolarsResult<ChunkedArray<T>>;
649}
650
651pub trait ChunkApplyKernel<A: Array> {
653 #[must_use]
655 fn apply_kernel(&self, f: &dyn Fn(&A) -> ArrayRef) -> Self;
656
657 fn apply_kernel_cast<S>(&self, f: &dyn Fn(&A) -> ArrayRef) -> ChunkedArray<S>
659 where
660 S: PolarsDataType;
661}
662
663#[cfg(feature = "is_first_distinct")]
664pub trait IsFirstDistinct<T: PolarsDataType> {
666 fn is_first_distinct(&self) -> PolarsResult<BooleanChunked> {
667 polars_bail!(opq = is_first_distinct, T::get_static_dtype());
668 }
669}
670
671#[cfg(feature = "is_last_distinct")]
672pub trait IsLastDistinct<T: PolarsDataType> {
674 fn is_last_distinct(&self) -> PolarsResult<BooleanChunked> {
675 polars_bail!(opq = is_last_distinct, T::get_static_dtype());
676 }
677}