polars_core/frame/group_by/

into_groups.rs

use arrow::legacy::kernels::sort_partition::{
    create_clean_partitions, partition_to_groups, partition_to_groups_amortized_varsize,
};
use polars_error::signals::try_raise_keyboard_interrupt;
use polars_utils::total_ord::{ToTotalOrd, TotalHash};

use super::*;
use crate::chunked_array::cast::CastOptions;
use crate::chunked_array::ops::row_encode::_get_rows_encoded_ca_unordered;
use crate::config::verbose;
use crate::series::BitRepr;
use crate::utils::Container;
use crate::utils::flatten::flatten_par;

/// Used to create the tuples for a group_by operation.
pub trait IntoGroupsType {
    /// Create the tuples need for a group_by operation.
    ///     * The first value in the tuple is the first index of the group.
    ///     * The second value in the tuple is the indexes of the groups including the first value.
    fn group_tuples(&self, _multithreaded: bool, _sorted: bool) -> PolarsResult<GroupsType> {
        unimplemented!()
    }
}

fn group_multithreaded<T: PolarsDataType>(ca: &ChunkedArray<T>) -> bool {
    // TODO! change to something sensible
    ca.len() > 1000 && POOL.current_num_threads() > 1
}

fn num_groups_proxy<T>(ca: &ChunkedArray<T>, multithreaded: bool, sorted: bool) -> GroupsType
where
    T: PolarsNumericType,
    T::Native: TotalHash + TotalEq + DirtyHash + ToTotalOrd,
    <T::Native as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + DirtyHash,
{
    if multithreaded && group_multithreaded(ca) {
        let n_partitions = _set_partition_size();

        // use the arrays as iterators
        if ca.null_count() == 0 {
            let keys = ca
                .downcast_iter()
                .map(|arr| arr.values().as_slice())
                .collect::<Vec<_>>();
            group_by_threaded_slice(keys, n_partitions, sorted)
        } else {
            let keys = ca
                .downcast_iter()
                .map(|arr| arr.iter().map(|o| o.copied()))
                .collect::<Vec<_>>();
            group_by_threaded_iter(&keys, n_partitions, sorted)
        }
    } else if !ca.has_nulls() {
        group_by(ca.into_no_null_iter(), sorted)
    } else {
        group_by(ca.iter(), sorted)
    }
}

impl<T> ChunkedArray<T>
where
    T: PolarsNumericType,
    T::Native: NumCast,
{
    fn create_groups_from_sorted(&self, multithreaded: bool) -> GroupsSlice {
        if verbose() {
            eprintln!("group_by keys are sorted; running sorted key fast path");
        }
        let arr = self.downcast_iter().next().unwrap();
        if arr.is_empty() {
            return GroupsSlice::default();
        }
        let mut values = arr.values().as_slice();
        let null_count = arr.null_count();
        let length = values.len();

        // all nulls
        if null_count == length {
            return vec![[0, length as IdxSize]];
        }

        let mut nulls_first = false;
        if null_count > 0 {
            nulls_first = arr.get(0).is_none()
        }

        if nulls_first {
            values = &values[null_count..];
        } else {
            values = &values[..length - null_count];
        };

        let n_threads = POOL.current_num_threads();
        if multithreaded && n_threads > 1 {
            let parts =
                create_clean_partitions(values, n_threads, self.is_sorted_descending_flag());
            let n_parts = parts.len();

            let first_ptr = &values[0] as *const T::Native as usize;
            let groups = parts.par_iter().enumerate().map(|(i, part)| {
                // we go via usize as *const is not send
                let first_ptr = first_ptr as *const T::Native;

                let part_first_ptr = &part[0] as *const T::Native;
                let mut offset = unsafe { part_first_ptr.offset_from(first_ptr) } as IdxSize;

                // nulls first: only add the nulls at the first partition
                if nulls_first && i == 0 {
                    partition_to_groups(part, null_count as IdxSize, true, offset)
                }
                // nulls last: only compute at the last partition
                else if !nulls_first && i == n_parts - 1 {
                    partition_to_groups(part, null_count as IdxSize, false, offset)
                }
                // other partitions
                else {
                    if nulls_first {
                        offset += null_count as IdxSize;
                    };

                    partition_to_groups(part, 0, false, offset)
                }
            });
            let groups = POOL.install(|| groups.collect::<Vec<_>>());
            flatten_par(&groups)
        } else {
            partition_to_groups(values, null_count as IdxSize, nulls_first, 0)
        }
    }
}

#[cfg(all(feature = "dtype-categorical", feature = "performant"))]
impl<T: PolarsCategoricalType> IntoGroupsType for CategoricalChunked<T>
where
    ChunkedArray<T::PolarsPhysical>: IntoGroupsType,
{
    fn group_tuples(&self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
        self.phys.group_tuples(multithreaded, sorted)
    }
}

impl<T> IntoGroupsType for ChunkedArray<T>
where
    T: PolarsNumericType,
    T::Native: TotalHash + TotalEq + DirtyHash + ToTotalOrd,
    <T::Native as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + DirtyHash,
{
    fn group_tuples(&self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
        // sorted path
        if self.is_sorted_ascending_flag() || self.is_sorted_descending_flag() {
            // don't have to pass `sorted` arg, GroupSlice is always sorted.
            return Ok(GroupsType::Slice {
                groups: self.rechunk().create_groups_from_sorted(multithreaded),
                overlapping: false,
            });
        }

        let out = match self.dtype() {
            DataType::Float32 => {
                // Convince the compiler that we are this type.
                let ca: &Float32Chunked = unsafe {
                    &*(self as *const ChunkedArray<T> as *const ChunkedArray<Float32Type>)
                };
                num_groups_proxy(ca, multithreaded, sorted)
            },
            DataType::Float64 => {
                // Convince the compiler that we are this type.
                let ca: &Float64Chunked = unsafe {
                    &*(self as *const ChunkedArray<T> as *const ChunkedArray<Float64Type>)
                };
                num_groups_proxy(ca, multithreaded, sorted)
            },
            _ => match self.to_bit_repr() {
                BitRepr::U8(ca) => num_groups_proxy(&ca, multithreaded, sorted),
                BitRepr::U16(ca) => num_groups_proxy(&ca, multithreaded, sorted),
                BitRepr::U32(ca) => num_groups_proxy(&ca, multithreaded, sorted),
                BitRepr::U64(ca) => num_groups_proxy(&ca, multithreaded, sorted),
                #[cfg(feature = "dtype-u128")]
                BitRepr::U128(ca) => num_groups_proxy(&ca, multithreaded, sorted),
            },
        };
        try_raise_keyboard_interrupt();
        Ok(out)
    }
}
impl IntoGroupsType for BooleanChunked {
    fn group_tuples(&self, mut multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
        multithreaded &= POOL.current_num_threads() > 1;

        #[cfg(feature = "performant")]
        {
            let ca = self
                .cast_with_options(&DataType::UInt8, CastOptions::Overflowing)
                .unwrap();
            let ca = ca.u8().unwrap();
            ca.group_tuples(multithreaded, sorted)
        }
        #[cfg(not(feature = "performant"))]
        {
            let ca = self
                .cast_with_options(&DataType::UInt32, CastOptions::Overflowing)
                .unwrap();
            let ca = ca.u32().unwrap();
            ca.group_tuples(multithreaded, sorted)
        }
    }
}

impl IntoGroupsType for StringChunked {
    #[allow(clippy::needless_lifetimes)]
    fn group_tuples<'a>(&'a self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
        self.as_binary().group_tuples(multithreaded, sorted)
    }
}

impl IntoGroupsType for BinaryChunked {
    #[allow(clippy::needless_lifetimes)]
    fn group_tuples<'a>(
        &'a self,
        mut multithreaded: bool,
        sorted: bool,
    ) -> PolarsResult<GroupsType> {
        if self.is_sorted_any() && !self.has_nulls() && self.n_chunks() == 1 {
            let arr = self.downcast_get(0).unwrap();
            let values = arr.values_iter();
            let mut out = Vec::with_capacity(values.len() / 30);
            partition_to_groups_amortized_varsize(values, arr.len() as _, 0, false, 0, &mut out);
            return Ok(GroupsType::Slice {
                groups: out,
                overlapping: false,
            });
        }

        multithreaded &= POOL.current_num_threads() > 1;
        let bh = self.to_bytes_hashes(multithreaded, Default::default());

        let out = if multithreaded {
            let n_partitions = bh.len();
            // Take slices so that the vecs are not cloned.
            let bh = bh.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
            group_by_threaded_slice(bh, n_partitions, sorted)
        } else {
            group_by(bh[0].iter(), sorted)
        };
        try_raise_keyboard_interrupt();
        Ok(out)
    }
}

impl IntoGroupsType for BinaryOffsetChunked {
    #[allow(clippy::needless_lifetimes)]
    fn group_tuples<'a>(
        &'a self,
        mut multithreaded: bool,
        sorted: bool,
    ) -> PolarsResult<GroupsType> {
        if self.is_sorted_any() && !self.has_nulls() && self.n_chunks() == 1 {
            let arr = self.downcast_get(0).unwrap();
            let values = arr.values_iter();
            let mut out = Vec::with_capacity(values.len() / 30);
            partition_to_groups_amortized_varsize(values, arr.len() as _, 0, false, 0, &mut out);
            return Ok(GroupsType::Slice {
                groups: out,
                overlapping: false,
            });
        } else if self.is_sorted_any() {
            let mut groups = Vec::new();

            let Some(y) = self.chunks().iter().position(|k| !k.as_ref().is_empty()) else {
                return Ok(GroupsType::Slice {
                    groups,
                    overlapping: false,
                });
            };

            let mut start_idx = 0;
            let mut i = 1;
            let mut x = 1;
            let mut start_value = self.downcast_chunks().get(y).unwrap().get(0);

            for keys in self.downcast_iter().skip(y) {
                if keys.has_nulls() {
                    for k in keys.iter().skip(x) {
                        if k != start_value {
                            groups.push([start_idx, i - start_idx]);
                            start_idx = i;
                            start_value = k;
                        }
                        i += 1;
                    }
                } else {
                    for k in keys.values_iter().skip(x) {
                        if Some(k) != start_value {
                            groups.push([start_idx, i - start_idx]);
                            start_idx = i;
                            start_value = Some(k);
                        }
                        i += 1;
                    }
                }
                x = 0;
            }

            groups.push([start_idx, i - start_idx]);
            return Ok(GroupsType::Slice {
                groups,
                overlapping: false,
            });
        }

        multithreaded &= POOL.current_num_threads() > 1;
        let bh = self.to_bytes_hashes(multithreaded, Default::default());

        let out = if multithreaded {
            let n_partitions = bh.len();
            // Take slices so that the vecs are not cloned.
            let bh = bh.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
            group_by_threaded_slice(bh, n_partitions, sorted)
        } else {
            group_by(bh[0].iter(), sorted)
        };
        Ok(out)
    }
}

impl IntoGroupsType for ListChunked {
    #[allow(clippy::needless_lifetimes)]
    #[allow(unused_variables)]
    fn group_tuples<'a>(
        &'a self,
        mut multithreaded: bool,
        sorted: bool,
    ) -> PolarsResult<GroupsType> {
        multithreaded &= POOL.current_num_threads() > 1;
        let by = &[self.clone().into_column()];
        let ca = if multithreaded {
            encode_rows_vertical_par_unordered(by).unwrap()
        } else {
            _get_rows_encoded_ca_unordered(PlSmallStr::EMPTY, by).unwrap()
        };

        ca.group_tuples(multithreaded, sorted)
    }
}

#[cfg(feature = "dtype-array")]
impl IntoGroupsType for ArrayChunked {
    #[allow(clippy::needless_lifetimes)]
    #[allow(unused_variables)]
    fn group_tuples<'a>(
        &'a self,
        mut multithreaded: bool,
        sorted: bool,
    ) -> PolarsResult<GroupsType> {
        multithreaded &= POOL.current_num_threads() > 1;
        let by = &[self.clone().into_column()];
        let ca = if multithreaded {
            encode_rows_vertical_par_unordered(by).unwrap()
        } else {
            _get_rows_encoded_ca_unordered(PlSmallStr::EMPTY, by).unwrap()
        };
        ca.group_tuples(multithreaded, sorted)
    }
}

#[cfg(feature = "object")]
impl<T> IntoGroupsType for ObjectChunked<T>
where
    T: PolarsObject,
{
    fn group_tuples(&self, _multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
        Ok(group_by(self.into_iter(), sorted))
    }
}