1use 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
59pub(crate) trait ToBitRepr {
63 fn to_bit_repr(&self) -> BitRepr;
64}
65
66pub trait ChunkAnyValue {
67 unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_>;
73
74 fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>>;
76}
77
78pub trait ChunkAnyValueBypassValidity {
79 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 pub empty_as_null: bool,
92 pub keep_nulls: bool,
94}
95
96pub 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#[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 fn take(&self, indices: &Idx) -> PolarsResult<Self>
125 where
126 Self: Sized;
127}
128
129pub trait ChunkTakeUnchecked<Idx: ?Sized> {
130 unsafe fn take_unchecked(&self, indices: &Idx) -> Self;
135}
136
137pub trait ChunkSet<'a, A, B> {
142 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 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 fn set(&'a self, mask: &BooleanChunked, opt_value: Option<A>) -> PolarsResult<Self>
192 where
193 Self: Sized;
194}
195
196pub trait ChunkCast {
198 fn cast(&self, dtype: &DataType) -> PolarsResult<Series> {
200 self.cast_with_options(dtype, CastOptions::NonStrict)
201 }
202
203 fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series>;
205
206 unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series>;
212}
213
214pub trait ChunkApply<'a, T> {
217 type FuncRet;
218
219 #[must_use]
233 fn apply_values<F>(&'a self, f: F) -> Self
234 where
235 F: Fn(T) -> Self::FuncRet + Copy;
236
237 #[must_use]
239 fn apply<F>(&'a self, f: F) -> Self
240 where
241 F: Fn(Option<T>) -> Option<Self::FuncRet> + Copy;
242
243 fn apply_to_slice<F, S>(&'a self, f: F, slice: &mut [S])
245 where
247 F: Fn(Option<T>, &S) -> S;
248}
249
250pub trait ChunkAgg<T> {
252 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 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 fn mean(&self) -> Option<f64> {
278 None
279 }
280}
281
282pub trait ChunkQuantile<T> {
284 fn median(&self) -> Option<T> {
287 None
288 }
289 fn quantile(&self, _quantile: f64, _method: QuantileMethod) -> PolarsResult<Option<T>> {
292 Ok(None)
293 }
294 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
304pub trait ChunkVar {
306 fn var(&self, _ddof: u8) -> Option<f64> {
308 None
309 }
310
311 fn std(&self, _ddof: u8) -> Option<f64> {
313 None
314 }
315}
316
317#[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
327pub trait ChunkCompareEq<Rhs> {
344 type Item;
345
346 fn equal(&self, rhs: Rhs) -> Self::Item;
348
349 fn equal_missing(&self, rhs: Rhs) -> Self::Item;
351
352 fn not_equal(&self, rhs: Rhs) -> Self::Item;
354
355 fn not_equal_missing(&self, rhs: Rhs) -> Self::Item;
357}
358
359pub trait ChunkCompareIneq<Rhs> {
362 type Item;
363
364 fn gt(&self, rhs: Rhs) -> Self::Item;
366
367 fn gt_eq(&self, rhs: Rhs) -> Self::Item;
369
370 fn lt(&self, rhs: Rhs) -> Self::Item;
372
373 fn lt_eq(&self, rhs: Rhs) -> Self::Item;
375}
376
377pub trait ChunkUnique {
379 fn unique(&self) -> PolarsResult<Self>
382 where
383 Self: Sized;
384
385 fn arg_unique(&self) -> PolarsResult<IdxCa>;
388
389 fn n_unique(&self) -> PolarsResult<usize> {
391 self.arg_unique().map(|v| v.len())
392 }
393
394 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
405pub trait ChunkSort<T: PolarsDataType> {
407 #[allow(unused_variables)]
408 fn sort_with(&self, options: SortOptions) -> ChunkedArray<T>;
409
410 fn sort(&self, descending: bool) -> ChunkedArray<T>;
412
413 fn arg_sort(&self, options: SortOptions) -> IdxCa;
415
416 #[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 Backward(FillNullLimit),
435 Forward(FillNullLimit),
437 Mean,
439 Min,
441 Max,
443 Zero,
445 One,
447}
448
449impl FillNullStrategy {
450 pub fn is_elementwise(&self) -> bool {
451 matches!(self, Self::One | Self::Zero)
452 }
453}
454
455pub trait ChunkFillNullValue<T> {
457 fn fill_null_with_values(&self, value: T) -> PolarsResult<Self>
459 where
460 Self: Sized;
461}
462
463pub trait ChunkFull<T> {
465 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
477pub trait ChunkReverse {
479 fn reverse(&self) -> Self;
481}
482
483pub trait ChunkFilter<T: PolarsDataType> {
485 fn filter(&self, filter: &BooleanChunked) -> PolarsResult<ChunkedArray<T>>
496 where
497 Self: Sized;
498}
499
500pub trait ChunkExpandAtIndex<T: PolarsDataType> {
502 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 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
635pub trait ChunkShiftFill<T: PolarsDataType, V> {
637 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
646pub trait ChunkZip<T: PolarsDataType> {
648 fn zip_with(
651 &self,
652 mask: &BooleanChunked,
653 other: &ChunkedArray<T>,
654 ) -> PolarsResult<ChunkedArray<T>>;
655}
656
657pub trait ChunkApplyKernel<A: Array> {
659 #[must_use]
661 fn apply_kernel(&self, f: &dyn Fn(&A) -> ArrayRef) -> Self;
662
663 fn apply_kernel_cast<S>(&self, f: &dyn Fn(&A) -> ArrayRef) -> ChunkedArray<S>
665 where
666 S: PolarsDataType;
667}