polars_core/utils/
mod.rs

1mod any_value;
2pub mod cut;
3use polars_arrow::compute::concatenate::concatenate_validities;
4use polars_arrow::compute::utils::combine_validities_and;
5pub mod flatten;
6pub(crate) mod series;
7mod supertype;
8use std::borrow::Cow;
9use std::ops::{Deref, DerefMut};
10mod schema;
11
12pub use any_value::*;
13use flatten::*;
14use num_traits::{One, Zero};
15pub use polars_arrow;
16use polars_arrow::bitmap::Bitmap;
17pub use polars_arrow::legacy::utils::*;
18pub use polars_arrow::trusted_len::TrustMyLength;
19pub use rayon;
20use rayon::prelude::*;
21pub use schema::*;
22pub use series::*;
23pub use supertype::*;
24
25use crate::prelude::*;
26use crate::runtime::RAYON;
27
28#[repr(transparent)]
29pub struct Wrap<T>(pub T);
30
31impl<T> Deref for Wrap<T> {
32    type Target = T;
33    fn deref(&self) -> &Self::Target {
34        &self.0
35    }
36}
37
38#[inline(always)]
39pub fn _set_partition_size() -> usize {
40    RAYON.current_num_threads()
41}
42
43/// Iterate `items` in parallel with the number of rayon tasks bounded by the thread count.
44///
45/// Use this when the length of `items` grows with the data. A worker adds a stack frame per
46/// stolen job, so an unbounded number of tasks can overflow a worker stack.
47pub fn par_iter_bounded<T: Sync>(items: &[T]) -> rayon::iter::MinLen<rayon::slice::Iter<'_, T>> {
48    const TASKS_PER_THREAD: usize = 8;
49    let min_len = items
50        .len()
51        .div_ceil(_set_partition_size() * TASKS_PER_THREAD);
52    items.par_iter().with_min_len(min_len.max(1))
53}
54
55/// Just a wrapper structure which is useful for certain impl specializations.
56///
57/// This is for instance use to implement
58/// `impl<T> FromIterator<T::Native> for NoNull<ChunkedArray<T>>`
59/// as `Option<T::Native>` was already implemented:
60/// `impl<T> FromIterator<Option<T::Native>> for ChunkedArray<T>`
61pub struct NoNull<T> {
62    inner: T,
63}
64
65impl<T> NoNull<T> {
66    pub fn new(inner: T) -> Self {
67        NoNull { inner }
68    }
69
70    pub fn into_inner(self) -> T {
71        self.inner
72    }
73}
74
75impl<T> Deref for NoNull<T> {
76    type Target = T;
77
78    fn deref(&self) -> &Self::Target {
79        &self.inner
80    }
81}
82
83impl<T> DerefMut for NoNull<T> {
84    fn deref_mut(&mut self) -> &mut Self::Target {
85        &mut self.inner
86    }
87}
88
89pub(crate) fn get_iter_capacity<T, I: Iterator<Item = T>>(iter: &I) -> usize {
90    match iter.size_hint() {
91        (_lower, Some(upper)) => upper,
92        (0, None) => 1024,
93        (lower, None) => lower,
94    }
95}
96
97// prefer this one over split_ca, as this can push the null_count into the thread pool
98// returns an `(offset, length)` tuple
99#[doc(hidden)]
100pub fn _split_offsets(len: usize, n: usize) -> Vec<(usize, usize)> {
101    if n == 1 {
102        vec![(0, len)]
103    } else {
104        let chunk_size = len / n;
105
106        (0..n)
107            .map(|partition| {
108                let offset = partition * chunk_size;
109                let len = if partition == (n - 1) {
110                    len - offset
111                } else {
112                    chunk_size
113                };
114                (partition * chunk_size, len)
115            })
116            .collect_trusted()
117    }
118}
119
120#[allow(clippy::len_without_is_empty)]
121pub trait Container: Clone {
122    fn slice(&self, offset: i64, len: usize) -> Self;
123
124    fn split_at(&self, offset: i64) -> (Self, Self);
125
126    fn len(&self) -> usize;
127
128    fn iter_chunks(&self) -> impl Iterator<Item = Self>;
129
130    fn should_rechunk(&self) -> bool;
131
132    fn n_chunks(&self) -> usize;
133
134    fn chunk_lengths(&self) -> impl Iterator<Item = usize>;
135}
136
137impl Container for DataFrame {
138    fn slice(&self, offset: i64, len: usize) -> Self {
139        DataFrame::slice(self, offset, len)
140    }
141
142    fn split_at(&self, offset: i64) -> (Self, Self) {
143        DataFrame::split_at(self, offset)
144    }
145
146    fn len(&self) -> usize {
147        self.height()
148    }
149
150    fn iter_chunks(&self) -> impl Iterator<Item = Self> {
151        flatten_df_iter(self)
152    }
153
154    fn should_rechunk(&self) -> bool {
155        self.should_rechunk()
156    }
157
158    fn n_chunks(&self) -> usize {
159        DataFrame::first_col_n_chunks(self)
160    }
161
162    fn chunk_lengths(&self) -> impl Iterator<Item = usize> {
163        // @scalar-correctness?
164        self.columns()[0].as_materialized_series().chunk_lengths()
165    }
166}
167
168impl<T: PolarsDataType> Container for ChunkedArray<T> {
169    fn slice(&self, offset: i64, len: usize) -> Self {
170        ChunkedArray::slice(self, offset, len)
171    }
172
173    fn split_at(&self, offset: i64) -> (Self, Self) {
174        ChunkedArray::split_at(self, offset)
175    }
176
177    fn len(&self) -> usize {
178        ChunkedArray::len(self)
179    }
180
181    fn iter_chunks(&self) -> impl Iterator<Item = Self> {
182        self.downcast_iter()
183            .map(|arr| Self::with_chunk(self.name().clone(), arr.clone()))
184    }
185
186    fn should_rechunk(&self) -> bool {
187        false
188    }
189
190    fn n_chunks(&self) -> usize {
191        self.chunks().len()
192    }
193
194    fn chunk_lengths(&self) -> impl Iterator<Item = usize> {
195        ChunkedArray::chunk_lengths(self)
196    }
197}
198
199impl Container for Series {
200    fn slice(&self, offset: i64, len: usize) -> Self {
201        self.0.slice(offset, len)
202    }
203
204    fn split_at(&self, offset: i64) -> (Self, Self) {
205        self.0.split_at(offset)
206    }
207
208    fn len(&self) -> usize {
209        self.0.len()
210    }
211
212    fn iter_chunks(&self) -> impl Iterator<Item = Self> {
213        (0..self.0.n_chunks()).map(|i| self.select_chunk(i))
214    }
215
216    fn should_rechunk(&self) -> bool {
217        false
218    }
219
220    fn n_chunks(&self) -> usize {
221        self.chunks().len()
222    }
223
224    fn chunk_lengths(&self) -> impl Iterator<Item = usize> {
225        self.0.chunk_lengths()
226    }
227}
228
229fn split_impl<C: Container>(container: &C, target: usize, chunk_size: usize) -> Vec<C> {
230    if target == 1 {
231        return vec![container.clone()];
232    }
233    let mut out = Vec::with_capacity(target);
234    let chunk_size = chunk_size as i64;
235
236    // First split
237    let (chunk, mut remainder) = container.split_at(chunk_size);
238    out.push(chunk);
239
240    // Take the rest of the splits of exactly chunk size, but skip the last remainder as we won't split that.
241    for _ in 1..target - 1 {
242        let (a, b) = remainder.split_at(chunk_size);
243        out.push(a);
244        remainder = b
245    }
246    // This can be slightly larger than `chunk_size`, but is smaller than `2 * chunk_size`.
247    out.push(remainder);
248    out
249}
250
251/// Splits, but doesn't flatten chunks. E.g. a container can still have multiple chunks.
252pub fn split<C: Container>(container: &C, target: usize) -> Vec<C> {
253    let total_len = container.len();
254    if total_len == 0 {
255        return vec![container.clone()];
256    }
257
258    let chunk_size = std::cmp::max(total_len / target, 1);
259
260    if container.n_chunks() == target
261        && container
262            .chunk_lengths()
263            .all(|len| len.abs_diff(chunk_size) < 100)
264        // We cannot get chunks if they are misaligned
265        && !container.should_rechunk()
266    {
267        return container.iter_chunks().collect();
268    }
269    split_impl(container, target, chunk_size)
270}
271
272/// Split a [`Container`] in `target` elements. The target doesn't have to be respected if not
273/// Deviation of the target might be done to create more equal size chunks.
274pub fn split_and_flatten<C: Container>(container: &C, target: usize) -> Vec<C> {
275    let total_len = container.len();
276    if total_len == 0 {
277        return vec![container.clone()];
278    }
279
280    let chunk_size = std::cmp::max(total_len / target, 1);
281
282    if container.n_chunks() == target
283        && container
284            .chunk_lengths()
285            .all(|len| len.abs_diff(chunk_size) < 100)
286        // We cannot get chunks if they are misaligned
287        && !container.should_rechunk()
288    {
289        return container.iter_chunks().collect();
290    }
291
292    if container.n_chunks() == 1 {
293        split_impl(container, target, chunk_size)
294    } else {
295        let mut out = Vec::with_capacity(target);
296        let chunks = container.iter_chunks();
297
298        'new_chunk: for mut chunk in chunks {
299            loop {
300                let h = chunk.len();
301                if h < chunk_size {
302                    // TODO if the chunk is much smaller than chunk size, we should try to merge it with the next one.
303                    out.push(chunk);
304                    continue 'new_chunk;
305                }
306
307                // If a split leads to the next chunk being smaller than 30% take the whole chunk
308                if ((h - chunk_size) as f64 / chunk_size as f64) < 0.3 {
309                    out.push(chunk);
310                    continue 'new_chunk;
311                }
312
313                let (a, b) = chunk.split_at(chunk_size as i64);
314                out.push(a);
315                chunk = b;
316            }
317        }
318        out
319    }
320}
321
322/// Split a [`DataFrame`] in `target` elements. The target doesn't have to be respected if not
323/// strict. Deviation of the target might be done to create more equal size chunks.
324///
325/// # Panics
326/// if chunks are not aligned
327pub fn split_df_as_ref(df: &DataFrame, target: usize, strict: bool) -> Vec<DataFrame> {
328    if strict {
329        split(df, target)
330    } else {
331        split_and_flatten(df, target)
332    }
333}
334
335#[doc(hidden)]
336/// Split a [`DataFrame`] into `n` parts. We take a `&mut` to be able to repartition/align chunks.
337/// `strict` in that it respects `n` even if the chunks are suboptimal.
338pub fn split_df(df: &mut DataFrame, target: usize, strict: bool) -> Vec<DataFrame> {
339    if target == 0 || df.height() == 0 {
340        return vec![df.clone()];
341    }
342    // make sure that chunks are aligned.
343    df.align_chunks_par();
344    split_df_as_ref(df, target, strict)
345}
346
347pub fn slice_slice<T>(vals: &[T], offset: i64, len: usize) -> &[T] {
348    let (raw_offset, slice_len) = slice_offsets(offset, len, vals.len());
349    &vals[raw_offset..raw_offset + slice_len]
350}
351
352#[inline]
353pub fn slice_offsets(offset: i64, length: usize, array_len: usize) -> (usize, usize) {
354    let signed_start_offset = if offset < 0 {
355        offset.saturating_add_unsigned(array_len as u64)
356    } else {
357        offset
358    };
359    let signed_stop_offset = signed_start_offset.saturating_add_unsigned(length as u64);
360
361    let signed_array_len: i64 = array_len
362        .try_into()
363        .expect("array length larger than i64::MAX");
364    let clamped_start_offset = signed_start_offset.clamp(0, signed_array_len);
365    let clamped_stop_offset = signed_stop_offset.clamp(0, signed_array_len);
366
367    let slice_start_idx = clamped_start_offset as usize;
368    let slice_len = (clamped_stop_offset - clamped_start_offset) as usize;
369    (slice_start_idx, slice_len)
370}
371
372/// Apply a macro on the Series
373#[macro_export]
374macro_rules! match_dtype_to_physical_apply_macro {
375    ($obj:expr, $macro:ident, $macro_string:ident, $macro_bool:ident $(, $opt_args:expr)*) => {{
376        match $obj {
377            DataType::String => $macro_string!($($opt_args)*),
378            DataType::Boolean => $macro_bool!($($opt_args)*),
379            #[cfg(feature = "dtype-u8")]
380            DataType::UInt8 => $macro!(u8 $(, $opt_args)*),
381            #[cfg(feature = "dtype-u16")]
382            DataType::UInt16 => $macro!(u16 $(, $opt_args)*),
383            DataType::UInt32 => $macro!(u32 $(, $opt_args)*),
384            DataType::UInt64 => $macro!(u64 $(, $opt_args)*),
385            #[cfg(feature = "dtype-i8")]
386            DataType::Int8 => $macro!(i8 $(, $opt_args)*),
387            #[cfg(feature = "dtype-i16")]
388            DataType::Int16 => $macro!(i16 $(, $opt_args)*),
389            DataType::Int32 => $macro!(i32 $(, $opt_args)*),
390            DataType::Int64 => $macro!(i64 $(, $opt_args)*),
391            #[cfg(feature = "dtype-i128")]
392            DataType::Int128 => $macro!(i128 $(, $opt_args)*),
393            #[cfg(feature = "dtype-f16")]
394            DataType::Float16 => $macro!(pf16 $(, $opt_args)*),
395            DataType::Float32 => $macro!(f32 $(, $opt_args)*),
396            DataType::Float64 => $macro!(f64 $(, $opt_args)*),
397            dt => panic!("not implemented for dtype {:?}", dt),
398        }
399    }};
400}
401
402/// Apply a macro on the Series
403#[macro_export]
404macro_rules! match_dtype_to_logical_apply_macro {
405    ($obj:expr, $macro:ident, $macro_string:ident, $macro_binary:ident, $macro_bool:ident $(, $opt_args:expr)*) => {{
406        match $obj {
407            DataType::String => $macro_string!($($opt_args)*),
408            DataType::Binary => $macro_binary!($($opt_args)*),
409            DataType::Boolean => $macro_bool!($($opt_args)*),
410            #[cfg(feature = "dtype-u8")]
411            DataType::UInt8 => $macro!(UInt8Type $(, $opt_args)*),
412            #[cfg(feature = "dtype-u16")]
413            DataType::UInt16 => $macro!(UInt16Type $(, $opt_args)*),
414            DataType::UInt32 => $macro!(UInt32Type $(, $opt_args)*),
415            DataType::UInt64 => $macro!(UInt64Type $(, $opt_args)*),
416            #[cfg(feature = "dtype-u128")]
417            DataType::UInt128 => $macro!(UInt128Type $(, $opt_args)*),
418            #[cfg(feature = "dtype-i8")]
419            DataType::Int8 => $macro!(Int8Type $(, $opt_args)*),
420            #[cfg(feature = "dtype-i16")]
421            DataType::Int16 => $macro!(Int16Type $(, $opt_args)*),
422            DataType::Int32 => $macro!(Int32Type $(, $opt_args)*),
423            DataType::Int64 => $macro!(Int64Type $(, $opt_args)*),
424            #[cfg(feature = "dtype-i128")]
425            DataType::Int128 => $macro!(Int128Type $(, $opt_args)*),
426            #[cfg(feature = "dtype-f16")]
427            DataType::Float16 => $macro!(Float16Type $(, $opt_args)*),
428            DataType::Float32 => $macro!(Float32Type $(, $opt_args)*),
429            DataType::Float64 => $macro!(Float64Type $(, $opt_args)*),
430            dt => panic!("not implemented for dtype {:?}", dt),
431        }
432    }};
433}
434
435/// Apply a macro on the Downcasted ChunkedArrays
436#[macro_export]
437macro_rules! match_arrow_dtype_apply_macro_ca {
438    ($self:expr, $macro:ident, $macro_string:ident, $macro_bool:ident $(, $opt_args:expr)*) => {{
439        match $self.dtype() {
440            DataType::String => $macro_string!($self.str().unwrap() $(, $opt_args)*),
441            DataType::Boolean => $macro_bool!($self.bool().unwrap() $(, $opt_args)*),
442            #[cfg(feature = "dtype-u8")]
443            DataType::UInt8 => $macro!($self.u8().unwrap() $(, $opt_args)*),
444            #[cfg(feature = "dtype-u16")]
445            DataType::UInt16 => $macro!($self.u16().unwrap() $(, $opt_args)*),
446            DataType::UInt32 => $macro!($self.u32().unwrap() $(, $opt_args)*),
447            DataType::UInt64 => $macro!($self.u64().unwrap() $(, $opt_args)*),
448            #[cfg(feature = "dtype-u128")]
449            DataType::UInt128 => $macro!($self.u128().unwrap() $(, $opt_args)*),
450            #[cfg(feature = "dtype-i8")]
451            DataType::Int8 => $macro!($self.i8().unwrap() $(, $opt_args)*),
452            #[cfg(feature = "dtype-i16")]
453            DataType::Int16 => $macro!($self.i16().unwrap() $(, $opt_args)*),
454            DataType::Int32 => $macro!($self.i32().unwrap() $(, $opt_args)*),
455            DataType::Int64 => $macro!($self.i64().unwrap() $(, $opt_args)*),
456            #[cfg(feature = "dtype-i128")]
457            DataType::Int128 => $macro!($self.i128().unwrap() $(, $opt_args)*),
458            #[cfg(feature = "dtype-f16")]
459            DataType::Float16 => $macro!($self.f16().unwrap() $(, $opt_args)*),
460            DataType::Float32 => $macro!($self.f32().unwrap() $(, $opt_args)*),
461            DataType::Float64 => $macro!($self.f64().unwrap() $(, $opt_args)*),
462            dt => panic!("not implemented for dtype {:?}", dt),
463        }
464    }};
465}
466
467#[macro_export]
468macro_rules! with_match_physical_numeric_type {(
469    $dtype:expr, | $_:tt $T:ident | $($body:tt)*
470) => ({
471    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
472    #[cfg(feature = "dtype-f16")]
473    use polars_utils::float16::pf16;
474    use $crate::datatypes::DataType::*;
475    match $dtype {
476        #[cfg(feature = "dtype-i8")]
477        Int8 => __with_ty__! { i8 },
478        #[cfg(feature = "dtype-i16")]
479        Int16 => __with_ty__! { i16 },
480        Int32 => __with_ty__! { i32 },
481        Int64 => __with_ty__! { i64 },
482        #[cfg(feature = "dtype-i128")]
483        Int128 => __with_ty__! { i128 },
484        #[cfg(feature = "dtype-u8")]
485        UInt8 => __with_ty__! { u8 },
486        #[cfg(feature = "dtype-u16")]
487        UInt16 => __with_ty__! { u16 },
488        UInt32 => __with_ty__! { u32 },
489        UInt64 => __with_ty__! { u64 },
490        #[cfg(feature = "dtype-u128")]
491        UInt128 => __with_ty__! { u128 },
492        #[cfg(feature = "dtype-f16")]
493        Float16 => __with_ty__! { pf16 },
494        Float32 => __with_ty__! { f32 },
495        Float64 => __with_ty__! { f64 },
496        dt => panic!("not implemented for dtype {:?}", dt),
497    }
498})}
499
500#[macro_export]
501macro_rules! with_match_physical_integer_type {(
502    $dtype:expr, | $_:tt $T:ident | $($body:tt)*
503) => ({
504    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
505    #[cfg(feature = "dtype-f16")]
506    use polars_utils::float16::pf16;
507    use $crate::datatypes::DataType::*;
508    match $dtype {
509        #[cfg(feature = "dtype-i8")]
510        Int8 => __with_ty__! { i8 },
511        #[cfg(feature = "dtype-i16")]
512        Int16 => __with_ty__! { i16 },
513        Int32 => __with_ty__! { i32 },
514        Int64 => __with_ty__! { i64 },
515        #[cfg(feature = "dtype-i128")]
516        Int128 => __with_ty__! { i128 },
517        #[cfg(feature = "dtype-u8")]
518        UInt8 => __with_ty__! { u8 },
519        #[cfg(feature = "dtype-u16")]
520        UInt16 => __with_ty__! { u16 },
521        UInt32 => __with_ty__! { u32 },
522        UInt64 => __with_ty__! { u64 },
523        #[cfg(feature = "dtype-u128")]
524        UInt128 => __with_ty__! { u128 },
525        dt => panic!("not implemented for dtype {:?}", dt),
526    }
527})}
528
529#[macro_export]
530macro_rules! with_match_physical_float_type {(
531    $dtype:expr, | $_:tt $T:ident | $($body:tt)*
532) => ({
533    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
534    use polars_utils::float16::pf16;
535    use $crate::datatypes::DataType::*;
536    match $dtype {
537        #[cfg(feature = "dtype-f16")]
538        Float16 => __with_ty__! { pf16 },
539        Float32 => __with_ty__! { f32 },
540        Float64 => __with_ty__! { f64 },
541        dt => panic!("not implemented for dtype {:?}", dt),
542    }
543})}
544
545#[macro_export]
546macro_rules! with_match_physical_float_polars_type {(
547    $key_type:expr, | $_:tt $T:ident | $($body:tt)*
548) => ({
549    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
550    use $crate::datatypes::DataType::*;
551    match $key_type {
552        #[cfg(feature = "dtype-f16")]
553        Float16 => __with_ty__! { Float16Type },
554        Float32 => __with_ty__! { Float32Type },
555        Float64 => __with_ty__! { Float64Type },
556        dt => panic!("not implemented for dtype {:?}", dt),
557    }
558})}
559
560#[macro_export]
561macro_rules! with_match_physical_numeric_polars_type {(
562    $key_type:expr, | $_:tt $T:ident | $($body:tt)*
563) => ({
564    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
565    use $crate::datatypes::DataType::*;
566    match $key_type {
567            #[cfg(feature = "dtype-i8")]
568        Int8 => __with_ty__! { Int8Type },
569            #[cfg(feature = "dtype-i16")]
570        Int16 => __with_ty__! { Int16Type },
571        Int32 => __with_ty__! { Int32Type },
572        Int64 => __with_ty__! { Int64Type },
573            #[cfg(feature = "dtype-i128")]
574        Int128 => __with_ty__! { Int128Type },
575            #[cfg(feature = "dtype-u8")]
576        UInt8 => __with_ty__! { UInt8Type },
577            #[cfg(feature = "dtype-u16")]
578        UInt16 => __with_ty__! { UInt16Type },
579        UInt32 => __with_ty__! { UInt32Type },
580        UInt64 => __with_ty__! { UInt64Type },
581            #[cfg(feature = "dtype-u128")]
582        UInt128 => __with_ty__! { UInt128Type },
583            #[cfg(feature = "dtype-f16")]
584        Float16 => __with_ty__! { Float16Type },
585        Float32 => __with_ty__! { Float32Type },
586        Float64 => __with_ty__! { Float64Type },
587        dt => panic!("not implemented for dtype {:?}", dt),
588    }
589})}
590
591#[macro_export]
592macro_rules! with_match_physical_integer_polars_type {(
593    $key_type:expr, | $_:tt $T:ident | $($body:tt)*
594) => ({
595    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
596    use $crate::datatypes::DataType::*;
597    use $crate::datatypes::*;
598    match $key_type {
599        #[cfg(feature = "dtype-i8")]
600        Int8 => __with_ty__! { Int8Type },
601        #[cfg(feature = "dtype-i16")]
602        Int16 => __with_ty__! { Int16Type },
603        Int32 => __with_ty__! { Int32Type },
604        Int64 => __with_ty__! { Int64Type },
605        #[cfg(feature = "dtype-i128")]
606        Int128 => __with_ty__! { Int128Type },
607        #[cfg(feature = "dtype-u8")]
608        UInt8 => __with_ty__! { UInt8Type },
609        #[cfg(feature = "dtype-u16")]
610        UInt16 => __with_ty__! { UInt16Type },
611        UInt32 => __with_ty__! { UInt32Type },
612        UInt64 => __with_ty__! { UInt64Type },
613        #[cfg(feature = "dtype-u128")]
614        UInt128 => __with_ty__! { UInt128Type },
615        dt => panic!("not implemented for dtype {:?}", dt),
616    }
617})}
618
619#[macro_export]
620macro_rules! with_match_categorical_physical_type {(
621    $dtype:expr, | $_:tt $T:ident | $($body:tt)*
622) => ({
623    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
624    match $dtype {
625        CategoricalPhysical::U8 => __with_ty__! { Categorical8Type },
626        CategoricalPhysical::U16 => __with_ty__! { Categorical16Type },
627        CategoricalPhysical::U32 => __with_ty__! { Categorical32Type },
628    }
629})}
630
631/// Apply a macro on the Downcasted ChunkedArrays of DataTypes that are logical numerics.
632/// So no logical.
633#[macro_export]
634macro_rules! downcast_as_macro_arg_physical {
635    ($self:expr, $macro:ident $(, $opt_args:expr)*) => {{
636        match $self.dtype() {
637            #[cfg(feature = "dtype-u8")]
638            DataType::UInt8 => $macro!($self.u8().unwrap() $(, $opt_args)*),
639            #[cfg(feature = "dtype-u16")]
640            DataType::UInt16 => $macro!($self.u16().unwrap() $(, $opt_args)*),
641            DataType::UInt32 => $macro!($self.u32().unwrap() $(, $opt_args)*),
642            DataType::UInt64 => $macro!($self.u64().unwrap() $(, $opt_args)*),
643            #[cfg(feature = "dtype-u128")]
644            DataType::UInt128 => $macro!($self.u128().unwrap() $(, $opt_args)*),
645            #[cfg(feature = "dtype-i8")]
646            DataType::Int8 => $macro!($self.i8().unwrap() $(, $opt_args)*),
647            #[cfg(feature = "dtype-i16")]
648            DataType::Int16 => $macro!($self.i16().unwrap() $(, $opt_args)*),
649            DataType::Int32 => $macro!($self.i32().unwrap() $(, $opt_args)*),
650            DataType::Int64 => $macro!($self.i64().unwrap() $(, $opt_args)*),
651            #[cfg(feature = "dtype-i128")]
652            DataType::Int128 => $macro!($self.i128().unwrap() $(, $opt_args)*),
653            #[cfg(feature = "dtype-f16")]
654            DataType::Float16 => $macro!($self.f16().unwrap() $(, $opt_args)*),
655            DataType::Float32 => $macro!($self.f32().unwrap() $(, $opt_args)*),
656            DataType::Float64 => $macro!($self.f64().unwrap() $(, $opt_args)*),
657            dt => panic!("not implemented for {:?}", dt),
658        }
659    }};
660}
661
662/// Apply a macro on the Downcasted ChunkedArrays of DataTypes that are logical numerics.
663/// So no logical.
664#[macro_export]
665macro_rules! downcast_as_macro_arg_physical_mut {
666    ($self:expr, $macro:ident $(, $opt_args:expr)*) => {{
667        // clone so that we do not borrow
668        match $self.dtype().clone() {
669            #[cfg(feature = "dtype-u8")]
670            DataType::UInt8 => {
671                let ca: &mut UInt8Chunked = $self.as_mut();
672                $macro!(UInt8Type, ca $(, $opt_args)*)
673            },
674            #[cfg(feature = "dtype-u16")]
675            DataType::UInt16 => {
676                let ca: &mut UInt16Chunked = $self.as_mut();
677                $macro!(UInt16Type, ca $(, $opt_args)*)
678            },
679            DataType::UInt32 => {
680                let ca: &mut UInt32Chunked = $self.as_mut();
681                $macro!(UInt32Type, ca $(, $opt_args)*)
682            },
683            DataType::UInt64 => {
684                let ca: &mut UInt64Chunked = $self.as_mut();
685                $macro!(UInt64Type, ca $(, $opt_args)*)
686            },
687            #[cfg(feature = "dtype-u128")]
688            DataType::UInt128 => {
689                let ca: &mut UInt128Chunked = $self.as_mut();
690                $macro!(UInt128Type, ca $(, $opt_args)*)
691            },
692            #[cfg(feature = "dtype-i8")]
693            DataType::Int8 => {
694                let ca: &mut Int8Chunked = $self.as_mut();
695                $macro!(Int8Type, ca $(, $opt_args)*)
696            },
697            #[cfg(feature = "dtype-i16")]
698            DataType::Int16 => {
699                let ca: &mut Int16Chunked = $self.as_mut();
700                $macro!(Int16Type, ca $(, $opt_args)*)
701            },
702            DataType::Int32 => {
703                let ca: &mut Int32Chunked = $self.as_mut();
704                $macro!(Int32Type, ca $(, $opt_args)*)
705            },
706            DataType::Int64 => {
707                let ca: &mut Int64Chunked = $self.as_mut();
708                $macro!(Int64Type, ca $(, $opt_args)*)
709            },
710            #[cfg(feature = "dtype-i128")]
711            DataType::Int128 => {
712                let ca: &mut Int128Chunked = $self.as_mut();
713                $macro!(Int128Type, ca $(, $opt_args)*)
714            },
715            #[cfg(feature = "dtype-f16")]
716            DataType::Float16 => {
717                let ca: &mut Float16Chunked = $self.as_mut();
718                $macro!(Float16Type, ca $(, $opt_args)*)
719            },
720            DataType::Float32 => {
721                let ca: &mut Float32Chunked = $self.as_mut();
722                $macro!(Float32Type, ca $(, $opt_args)*)
723            },
724            DataType::Float64 => {
725                let ca: &mut Float64Chunked = $self.as_mut();
726                $macro!(Float64Type, ca $(, $opt_args)*)
727            },
728            dt => panic!("not implemented for {:?}", dt),
729        }
730    }};
731}
732
733#[macro_export]
734macro_rules! apply_method_all_arrow_series {
735    ($self:expr, $method:ident, $($args:expr),*) => {
736        match $self.dtype() {
737            DataType::Boolean => $self.bool().unwrap().$method($($args),*),
738            DataType::String => $self.str().unwrap().$method($($args),*),
739            #[cfg(feature = "dtype-u8")]
740            DataType::UInt8 => $self.u8().unwrap().$method($($args),*),
741            #[cfg(feature = "dtype-u16")]
742            DataType::UInt16 => $self.u16().unwrap().$method($($args),*),
743            DataType::UInt32 => $self.u32().unwrap().$method($($args),*),
744            DataType::UInt64 => $self.u64().unwrap().$method($($args),*),
745            #[cfg(feature = "dtype-u128")]
746            DataType::UInt128 => $self.u128().unwrap().$medthod($($args),*),
747            #[cfg(feature = "dtype-i8")]
748            DataType::Int8 => $self.i8().unwrap().$method($($args),*),
749            #[cfg(feature = "dtype-i16")]
750            DataType::Int16 => $self.i16().unwrap().$method($($args),*),
751            DataType::Int32 => $self.i32().unwrap().$method($($args),*),
752            DataType::Int64 => $self.i64().unwrap().$method($($args),*),
753            #[cfg(feature = "dtype-i128")]
754            DataType::Int128 => $self.i128().unwrap().$method($($args),*),
755            #[cfg(feature = "dtype-f16")]
756            DataType::Float16 => $self.f16().unwrap().$method($($args),*),
757            DataType::Float32 => $self.f32().unwrap().$method($($args),*),
758            DataType::Float64 => $self.f64().unwrap().$method($($args),*),
759            DataType::Time => $self.time().unwrap().$method($($args),*),
760            DataType::Date => $self.date().unwrap().$method($($args),*),
761            DataType::Datetime(_, _) => $self.datetime().unwrap().$method($($args),*),
762            DataType::List(_) => $self.list().unwrap().$method($($args),*),
763            DataType::Struct(_) => $self.struct_().unwrap().$method($($args),*),
764            dt => panic!("dtype {:?} not supported", dt)
765        }
766    }
767}
768
769#[macro_export]
770macro_rules! apply_method_physical_integer {
771    ($self:expr, $method:ident, $($args:expr),*) => {
772        match $self.dtype() {
773            #[cfg(feature = "dtype-u8")]
774            DataType::UInt8 => $self.u8().unwrap().$method($($args),*),
775            #[cfg(feature = "dtype-u16")]
776            DataType::UInt16 => $self.u16().unwrap().$method($($args),*),
777            DataType::UInt32 => $self.u32().unwrap().$method($($args),*),
778            DataType::UInt64 => $self.u64().unwrap().$method($($args),*),
779            #[cfg(feature = "dtype-u128")]
780            DataType::UInt128 => $self.u128().unwrap().$method($($args),*),
781            #[cfg(feature = "dtype-i8")]
782            DataType::Int8 => $self.i8().unwrap().$method($($args),*),
783            #[cfg(feature = "dtype-i16")]
784            DataType::Int16 => $self.i16().unwrap().$method($($args),*),
785            DataType::Int32 => $self.i32().unwrap().$method($($args),*),
786            DataType::Int64 => $self.i64().unwrap().$method($($args),*),
787            #[cfg(feature = "dtype-i128")]
788            DataType::Int128 => $self.i128().unwrap().$method($($args),*),
789            dt => panic!("not implemented for dtype {:?}", dt),
790        }
791    }
792}
793
794// doesn't include Bool and String
795#[macro_export]
796macro_rules! apply_method_physical_numeric {
797    ($self:expr, $method:ident, $($args:expr),*) => {
798        match $self.dtype() {
799            #[cfg(feature = "dtype-f16")]
800            DataType::Float16 => $self.f16().unwrap().$method($($args),*),
801            DataType::Float32 => $self.f32().unwrap().$method($($args),*),
802            DataType::Float64 => $self.f64().unwrap().$method($($args),*),
803            _ => apply_method_physical_integer!($self, $method, $($args),*),
804        }
805    }
806}
807
808#[macro_export]
809macro_rules! df {
810    ($($col_name:expr => $slice:expr), + $(,)?) => {
811        $crate::prelude::DataFrame::new_infer_height(vec![
812            $($crate::prelude::Column::from(<$crate::prelude::Series as $crate::prelude::NamedFrom::<_, _>>::new($col_name.into(), $slice)),)+
813        ])
814    }
815}
816
817pub fn get_time_units(tu_l: &TimeUnit, tu_r: &TimeUnit) -> TimeUnit {
818    use crate::datatypes::time_unit::TimeUnit::*;
819    match (tu_l, tu_r) {
820        (Nanoseconds, Microseconds) => Microseconds,
821        (_, Milliseconds) => Milliseconds,
822        _ => *tu_l,
823    }
824}
825
826#[cold]
827#[inline(never)]
828fn width_mismatch(df1: &DataFrame, df2: &DataFrame) -> PolarsError {
829    let mut df1_extra = Vec::new();
830    let mut df2_extra = Vec::new();
831
832    let s1 = df1.schema();
833    let s2 = df2.schema();
834
835    s1.field_compare(s2, &mut df1_extra, &mut df2_extra);
836
837    let df1_extra = df1_extra
838        .into_iter()
839        .map(|(_, (n, _))| n.as_str())
840        .collect::<Vec<_>>()
841        .join(", ");
842    let df2_extra = df2_extra
843        .into_iter()
844        .map(|(_, (n, _))| n.as_str())
845        .collect::<Vec<_>>()
846        .join(", ");
847
848    polars_err!(
849        SchemaMismatch: r#"unable to vstack, dataframes have different widths ({} != {}).
850One dataframe has additional columns: [{df1_extra}].
851Other dataframe has additional columns: [{df2_extra}]."#,
852        df1.width(),
853        df2.width(),
854    )
855}
856
857/// This takes ownership of the DataFrame so that drop is called earlier.
858/// Does not check if schema is correct
859pub fn accumulate_dataframes_vertical_unchecked<I>(dfs: I) -> DataFrame
860where
861    I: IntoIterator<Item = DataFrame>,
862{
863    let mut iter = dfs.into_iter();
864    let additional = iter.size_hint().0;
865    let mut acc_df = iter.next().unwrap();
866    acc_df.reserve_chunks(additional);
867
868    for df in iter {
869        if acc_df.width() != df.width() {
870            panic!("{}", width_mismatch(&acc_df, &df));
871        }
872
873        acc_df.vstack_mut_owned_unchecked(df);
874    }
875    acc_df
876}
877
878/// This takes ownership of the DataFrame so that drop is called earlier.
879/// # Panics
880/// Panics if `dfs` is empty.
881pub fn accumulate_dataframes_vertical<I>(dfs: I) -> PolarsResult<DataFrame>
882where
883    I: IntoIterator<Item = DataFrame>,
884{
885    let mut iter = dfs.into_iter();
886    let additional = iter.size_hint().0;
887    let mut acc_df = iter.next().unwrap();
888    acc_df.reserve_chunks(additional);
889    for df in iter {
890        if acc_df.width() != df.width() {
891            return Err(width_mismatch(&acc_df, &df));
892        }
893
894        acc_df.vstack_mut_owned(df)?;
895    }
896
897    Ok(acc_df)
898}
899
900/// Concat the DataFrames to a single DataFrame.
901pub fn concat_df<'a, I>(dfs: I) -> PolarsResult<DataFrame>
902where
903    I: IntoIterator<Item = &'a DataFrame>,
904{
905    let mut iter = dfs.into_iter();
906    let additional = iter.size_hint().0;
907    let mut acc_df = iter.next().unwrap().clone();
908    acc_df.reserve_chunks(additional);
909    for df in iter {
910        acc_df.vstack_mut(df)?;
911    }
912    Ok(acc_df)
913}
914
915/// Concat the DataFrames to a single DataFrame.
916pub fn concat_df_unchecked<'a, I>(dfs: I) -> DataFrame
917where
918    I: IntoIterator<Item = &'a DataFrame>,
919{
920    let mut iter = dfs.into_iter();
921    let additional = iter.size_hint().0;
922    let mut acc_df = iter.next().unwrap().clone();
923    acc_df.reserve_chunks(additional);
924    for df in iter {
925        acc_df.vstack_mut_unchecked(df);
926    }
927    acc_df
928}
929
930pub fn accumulate_dataframes_horizontal(dfs: Vec<DataFrame>) -> PolarsResult<DataFrame> {
931    let mut iter = dfs.into_iter();
932    let mut acc_df = iter.next().unwrap();
933    for df in iter {
934        acc_df.hstack_mut(df.columns())?;
935    }
936    Ok(acc_df)
937}
938
939/// Ensure the chunks in both ChunkedArrays have the same length.
940/// # Panics
941/// This will panic if `left.len() != right.len()` and array is chunked.
942pub fn align_chunks_binary<'a, T, B>(
943    left: &'a ChunkedArray<T>,
944    right: &'a ChunkedArray<B>,
945) -> (Cow<'a, ChunkedArray<T>>, Cow<'a, ChunkedArray<B>>)
946where
947    B: PolarsDataType,
948    T: PolarsDataType,
949{
950    let assert = || {
951        assert_eq!(
952            left.len(),
953            right.len(),
954            "expected arrays of the same length"
955        )
956    };
957    match (left.chunks.len(), right.chunks.len()) {
958        // All chunks are equal length
959        (1, 1) => (Cow::Borrowed(left), Cow::Borrowed(right)),
960        // All chunks are equal length
961        (a, b)
962            if a == b
963                && left
964                    .chunk_lengths()
965                    .zip(right.chunk_lengths())
966                    .all(|(l, r)| l == r) =>
967        {
968            (Cow::Borrowed(left), Cow::Borrowed(right))
969        },
970        (_, 1) => {
971            assert();
972            (
973                Cow::Borrowed(left),
974                Cow::Owned(right.match_chunks(left.chunk_lengths())),
975            )
976        },
977        (1, _) => {
978            assert();
979            (
980                Cow::Owned(left.match_chunks(right.chunk_lengths())),
981                Cow::Borrowed(right),
982            )
983        },
984        (_, _) => {
985            assert();
986            // could optimize to choose to rechunk a primitive and not a string or list type
987            let left = left.rechunk();
988            (
989                Cow::Owned(left.match_chunks(right.chunk_lengths())),
990                Cow::Borrowed(right),
991            )
992        },
993    }
994}
995
996/// Ensure the chunks in ChunkedArray and Series have the same length.
997/// # Panics
998/// This will panic if `left.len() != right.len()` and array is chunked.
999pub fn align_chunks_binary_ca_series<'a, T>(
1000    left: &'a ChunkedArray<T>,
1001    right: &'a Series,
1002) -> (Cow<'a, ChunkedArray<T>>, Cow<'a, Series>)
1003where
1004    T: PolarsDataType,
1005{
1006    let assert = || {
1007        assert_eq!(
1008            left.len(),
1009            right.len(),
1010            "expected arrays of the same length"
1011        )
1012    };
1013    match (left.chunks.len(), right.chunks().len()) {
1014        // All chunks are equal length
1015        (1, 1) => (Cow::Borrowed(left), Cow::Borrowed(right)),
1016        // All chunks are equal length
1017        (a, b)
1018            if a == b
1019                && left
1020                    .chunk_lengths()
1021                    .zip(right.chunk_lengths())
1022                    .all(|(l, r)| l == r) =>
1023        {
1024            assert();
1025            (Cow::Borrowed(left), Cow::Borrowed(right))
1026        },
1027        (_, 1) => (left.rechunk(), Cow::Borrowed(right)),
1028        (1, _) => (Cow::Borrowed(left), Cow::Owned(right.rechunk())),
1029        (_, _) => {
1030            assert();
1031            (left.rechunk(), Cow::Owned(right.rechunk()))
1032        },
1033    }
1034}
1035
1036#[cfg(feature = "performant")]
1037pub(crate) fn align_chunks_binary_owned_series(left: Series, right: Series) -> (Series, Series) {
1038    match (left.chunks().len(), right.chunks().len()) {
1039        (1, 1) => (left, right),
1040        // All chunks are equal length
1041        (a, b)
1042            if a == b
1043                && left
1044                    .chunk_lengths()
1045                    .zip(right.chunk_lengths())
1046                    .all(|(l, r)| l == r) =>
1047        {
1048            (left, right)
1049        },
1050        (_, 1) => (left.rechunk(), right),
1051        (1, _) => (left, right.rechunk()),
1052        (_, _) => (left.rechunk(), right.rechunk()),
1053    }
1054}
1055
1056pub(crate) fn align_chunks_binary_owned<T, B>(
1057    left: ChunkedArray<T>,
1058    right: ChunkedArray<B>,
1059) -> (ChunkedArray<T>, ChunkedArray<B>)
1060where
1061    B: PolarsDataType,
1062    T: PolarsDataType,
1063{
1064    match (left.chunks.len(), right.chunks.len()) {
1065        (1, 1) => (left, right),
1066        // All chunks are equal length
1067        (a, b)
1068            if a == b
1069                && left
1070                    .chunk_lengths()
1071                    .zip(right.chunk_lengths())
1072                    .all(|(l, r)| l == r) =>
1073        {
1074            (left, right)
1075        },
1076        (_, 1) => (left.rechunk().into_owned(), right),
1077        (1, _) => (left, right.rechunk().into_owned()),
1078        (_, _) => (left.rechunk().into_owned(), right.rechunk().into_owned()),
1079    }
1080}
1081
1082/// # Panics
1083/// This will panic if `a.len() != b.len() || b.len() != c.len()` and array is chunked.
1084#[allow(clippy::type_complexity)]
1085pub fn align_chunks_ternary<'a, A, B, C>(
1086    a: &'a ChunkedArray<A>,
1087    b: &'a ChunkedArray<B>,
1088    c: &'a ChunkedArray<C>,
1089) -> (
1090    Cow<'a, ChunkedArray<A>>,
1091    Cow<'a, ChunkedArray<B>>,
1092    Cow<'a, ChunkedArray<C>>,
1093)
1094where
1095    A: PolarsDataType,
1096    B: PolarsDataType,
1097    C: PolarsDataType,
1098{
1099    if a.chunks.len() == 1 && b.chunks.len() == 1 && c.chunks.len() == 1 {
1100        return (Cow::Borrowed(a), Cow::Borrowed(b), Cow::Borrowed(c));
1101    }
1102
1103    assert!(
1104        a.len() == b.len() && b.len() == c.len(),
1105        "expected arrays of the same length"
1106    );
1107
1108    match (a.chunks.len(), b.chunks.len(), c.chunks.len()) {
1109        (_, 1, 1) => (
1110            Cow::Borrowed(a),
1111            Cow::Owned(b.match_chunks(a.chunk_lengths())),
1112            Cow::Owned(c.match_chunks(a.chunk_lengths())),
1113        ),
1114        (1, 1, _) => (
1115            Cow::Owned(a.match_chunks(c.chunk_lengths())),
1116            Cow::Owned(b.match_chunks(c.chunk_lengths())),
1117            Cow::Borrowed(c),
1118        ),
1119        (1, _, 1) => (
1120            Cow::Owned(a.match_chunks(b.chunk_lengths())),
1121            Cow::Borrowed(b),
1122            Cow::Owned(c.match_chunks(b.chunk_lengths())),
1123        ),
1124        (1, _, _) => {
1125            let b = b.rechunk();
1126            (
1127                Cow::Owned(a.match_chunks(c.chunk_lengths())),
1128                Cow::Owned(b.match_chunks(c.chunk_lengths())),
1129                Cow::Borrowed(c),
1130            )
1131        },
1132        (_, 1, _) => {
1133            let a = a.rechunk();
1134            (
1135                Cow::Owned(a.match_chunks(c.chunk_lengths())),
1136                Cow::Owned(b.match_chunks(c.chunk_lengths())),
1137                Cow::Borrowed(c),
1138            )
1139        },
1140        (_, _, 1) => {
1141            let b = b.rechunk();
1142            (
1143                Cow::Borrowed(a),
1144                Cow::Owned(b.match_chunks(a.chunk_lengths())),
1145                Cow::Owned(c.match_chunks(a.chunk_lengths())),
1146            )
1147        },
1148        (len_a, len_b, len_c)
1149            if len_a == len_b
1150                && len_b == len_c
1151                && a.chunk_lengths()
1152                    .zip(b.chunk_lengths())
1153                    .zip(c.chunk_lengths())
1154                    .all(|((a, b), c)| a == b && b == c) =>
1155        {
1156            (Cow::Borrowed(a), Cow::Borrowed(b), Cow::Borrowed(c))
1157        },
1158        _ => {
1159            // could optimize to choose to rechunk a primitive and not a string or list type
1160            let a = a.rechunk();
1161            let b = b.rechunk();
1162            (
1163                Cow::Owned(a.match_chunks(c.chunk_lengths())),
1164                Cow::Owned(b.match_chunks(c.chunk_lengths())),
1165                Cow::Borrowed(c),
1166            )
1167        },
1168    }
1169}
1170
1171pub fn binary_concatenate_validities<'a, T, B>(
1172    left: &'a ChunkedArray<T>,
1173    right: &'a ChunkedArray<B>,
1174) -> Option<Bitmap>
1175where
1176    B: PolarsDataType,
1177    T: PolarsDataType,
1178{
1179    let (left, right) = align_chunks_binary(left, right);
1180    let left_validity = concatenate_validities(left.chunks());
1181    let right_validity = concatenate_validities(right.chunks());
1182    combine_validities_and(left_validity.as_ref(), right_validity.as_ref())
1183}
1184
1185/// Convenience for `x.into_iter().map(Into::into).collect()` using an `into_vec()` function.
1186pub trait IntoVec<T> {
1187    fn into_vec(self) -> Vec<T>;
1188}
1189
1190impl<I, S> IntoVec<PlSmallStr> for I
1191where
1192    I: IntoIterator<Item = S>,
1193    S: Into<PlSmallStr>,
1194{
1195    fn into_vec(self) -> Vec<PlSmallStr> {
1196        self.into_iter().map(|s| s.into()).collect()
1197    }
1198}
1199
1200/// This logic is same as the impl on ChunkedArray
1201/// The difference is that there is less indirection because the caller should preallocate
1202/// `chunk_lens` once. On the `ChunkedArray` we indirect through an `ArrayRef` which is an indirection
1203/// and a vtable.
1204#[inline]
1205pub(crate) fn index_to_chunked_index<
1206    I: Iterator<Item = Idx>,
1207    Idx: PartialOrd + std::ops::AddAssign + std::ops::SubAssign + Zero + One,
1208>(
1209    chunk_lens: I,
1210    index: Idx,
1211) -> (Idx, Idx) {
1212    let mut index_remainder = index;
1213    let mut current_chunk_idx = Zero::zero();
1214
1215    for chunk_len in chunk_lens {
1216        if chunk_len > index_remainder {
1217            break;
1218        } else {
1219            index_remainder -= chunk_len;
1220            current_chunk_idx += One::one();
1221        }
1222    }
1223    (current_chunk_idx, index_remainder)
1224}
1225
1226pub(crate) fn index_to_chunked_index_rev<
1227    I: Iterator<Item = Idx>,
1228    Idx: PartialOrd
1229        + std::ops::AddAssign
1230        + std::ops::SubAssign
1231        + std::ops::Sub<Output = Idx>
1232        + Zero
1233        + One
1234        + Copy
1235        + std::fmt::Debug,
1236>(
1237    chunk_lens_rev: I,
1238    index_from_back: Idx,
1239    total_chunks: Idx,
1240) -> (Idx, Idx) {
1241    debug_assert!(index_from_back > Zero::zero(), "at least -1");
1242    let mut index_remainder = index_from_back;
1243    let mut current_chunk_idx = One::one();
1244    let mut current_chunk_len = Zero::zero();
1245
1246    for chunk_len in chunk_lens_rev {
1247        current_chunk_len = chunk_len;
1248        if chunk_len >= index_remainder {
1249            break;
1250        } else {
1251            index_remainder -= chunk_len;
1252            current_chunk_idx += One::one();
1253        }
1254    }
1255    (
1256        total_chunks - current_chunk_idx,
1257        current_chunk_len - index_remainder,
1258    )
1259}
1260
1261pub fn first_null<'a, I>(iter: I) -> Option<usize>
1262where
1263    I: Iterator<Item = &'a dyn Array>,
1264{
1265    let mut offset = 0;
1266    for arr in iter {
1267        if let Some(mask) = arr.validity() {
1268            let len_mask = mask.len();
1269            let n = mask.leading_ones();
1270            if n < len_mask {
1271                return Some(offset + n);
1272            }
1273            offset += len_mask
1274        } else {
1275            offset += arr.len();
1276        }
1277    }
1278    None
1279}
1280
1281pub fn first_non_null<'a, I>(iter: I) -> Option<usize>
1282where
1283    I: Iterator<Item = &'a dyn Array>,
1284{
1285    let mut offset = 0;
1286    for arr in iter {
1287        if let Some(mask) = arr.validity() {
1288            let len_mask = mask.len();
1289            let n = mask.leading_zeros();
1290            if n < len_mask {
1291                return Some(offset + n);
1292            }
1293            offset += len_mask
1294        } else if !arr.is_empty() {
1295            return Some(offset);
1296        }
1297    }
1298    None
1299}
1300
1301pub fn last_non_null<'a, I>(iter: I, len: usize) -> Option<usize>
1302where
1303    I: DoubleEndedIterator<Item = &'a dyn Array>,
1304{
1305    if len == 0 {
1306        return None;
1307    }
1308    let mut offset = 0;
1309    for arr in iter.rev() {
1310        if let Some(mask) = arr.validity() {
1311            let len_mask = mask.len();
1312            let n = mask.trailing_zeros();
1313            if n < len_mask {
1314                return Some(len - offset - n - 1);
1315            }
1316            offset += len_mask;
1317        } else if !arr.is_empty() {
1318            return Some(len - offset - 1);
1319        }
1320    }
1321    None
1322}
1323
1324pub fn coalesce_nulls_columns(a: &Column, b: &Column) -> (Column, Column) {
1325    if a.null_count() > 0 || b.null_count() > 0 {
1326        let mut a = a.as_materialized_series().rechunk();
1327        let mut b = b.as_materialized_series().rechunk();
1328        for (arr_a, arr_b) in unsafe { a.chunks_mut().iter_mut().zip(b.chunks_mut()) } {
1329            let validity = match (arr_a.validity(), arr_b.validity()) {
1330                (None, Some(b)) => Some(b.clone()),
1331                (Some(a), Some(b)) => Some(a & b),
1332                (Some(a), None) => Some(a.clone()),
1333                (None, None) => None,
1334            };
1335            *arr_a = arr_a.with_validity(validity.clone());
1336            *arr_b = arr_b.with_validity(validity);
1337        }
1338        a.compute_len();
1339        b.compute_len();
1340        (a.into(), b.into())
1341    } else {
1342        (a.clone(), b.clone())
1343    }
1344}
1345
1346#[cfg(test)]
1347mod test {
1348    use super::*;
1349
1350    #[test]
1351    fn test_split() {
1352        let ca: Int32Chunked = (0..10).collect_ca("a".into());
1353
1354        let out = split(&ca, 3);
1355        assert_eq!(out[0].len(), 3);
1356        assert_eq!(out[1].len(), 3);
1357        assert_eq!(out[2].len(), 4);
1358    }
1359
1360    #[test]
1361    fn test_align_chunks() -> PolarsResult<()> {
1362        let a = Int32Chunked::new(PlSmallStr::EMPTY, &[1, 2, 3, 4]);
1363        let mut b = Int32Chunked::new(PlSmallStr::EMPTY, &[1]);
1364        let b2 = Int32Chunked::new(PlSmallStr::EMPTY, &[2, 3, 4]);
1365
1366        b.append(&b2)?;
1367        let (a, b) = align_chunks_binary(&a, &b);
1368        assert_eq!(
1369            a.chunk_lengths().collect::<Vec<_>>(),
1370            b.chunk_lengths().collect::<Vec<_>>()
1371        );
1372
1373        let a = Int32Chunked::new(PlSmallStr::EMPTY, &[1, 2, 3, 4]);
1374        let mut b = Int32Chunked::new(PlSmallStr::EMPTY, &[1]);
1375        let b1 = b.clone();
1376        b.append(&b1)?;
1377        b.append(&b1)?;
1378        b.append(&b1)?;
1379        let (a, b) = align_chunks_binary(&a, &b);
1380        assert_eq!(
1381            a.chunk_lengths().collect::<Vec<_>>(),
1382            b.chunk_lengths().collect::<Vec<_>>()
1383        );
1384
1385        Ok(())
1386    }
1387}
polars_core/utils/mod.rs

polars_core/utils/
mod.rs
