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

polars_core/utils/
mod.rs
