polars_core/frame/group_by/
into_groups.rs

1use arrow::legacy::kernels::sort_partition::{
2    create_clean_partitions, partition_to_groups, partition_to_groups_amortized_varsize,
3};
4use polars_error::signals::try_raise_keyboard_interrupt;
5use polars_utils::total_ord::{ToTotalOrd, TotalHash};
6
7use super::*;
8use crate::chunked_array::cast::CastOptions;
9use crate::chunked_array::ops::row_encode::_get_rows_encoded_ca_unordered;
10use crate::config::verbose;
11use crate::series::BitRepr;
12use crate::utils::Container;
13use crate::utils::flatten::flatten_par;
14
15/// Used to create the tuples for a group_by operation.
16pub trait IntoGroupsType {
17    /// Create the tuples need for a group_by operation.
18    ///     * The first value in the tuple is the first index of the group.
19    ///     * The second value in the tuple is the indexes of the groups including the first value.
20    fn group_tuples(&self, _multithreaded: bool, _sorted: bool) -> PolarsResult<GroupsType> {
21        unimplemented!()
22    }
23}
24
25fn group_multithreaded<T: PolarsDataType>(ca: &ChunkedArray<T>) -> bool {
26    // TODO! change to something sensible
27    ca.len() > 1000 && POOL.current_num_threads() > 1
28}
29
30fn num_groups_proxy<T>(ca: &ChunkedArray<T>, multithreaded: bool, sorted: bool) -> GroupsType
31where
32    T: PolarsNumericType,
33    T::Native: TotalHash + TotalEq + DirtyHash + ToTotalOrd,
34    <T::Native as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + DirtyHash,
35{
36    if multithreaded && group_multithreaded(ca) {
37        let n_partitions = _set_partition_size();
38
39        // use the arrays as iterators
40        if ca.null_count() == 0 {
41            let keys = ca
42                .downcast_iter()
43                .map(|arr| arr.values().as_slice())
44                .collect::<Vec<_>>();
45            group_by_threaded_slice(keys, n_partitions, sorted)
46        } else {
47            let keys = ca
48                .downcast_iter()
49                .map(|arr| arr.iter().map(|o| o.copied()))
50                .collect::<Vec<_>>();
51            group_by_threaded_iter(&keys, n_partitions, sorted)
52        }
53    } else if !ca.has_nulls() {
54        group_by(ca.into_no_null_iter(), sorted)
55    } else {
56        group_by(ca.iter(), sorted)
57    }
58}
59
60impl<T> ChunkedArray<T>
61where
62    T: PolarsNumericType,
63    T::Native: NumCast,
64{
65    fn create_groups_from_sorted(&self, multithreaded: bool) -> GroupsSlice {
66        if verbose() {
67            eprintln!("group_by keys are sorted; running sorted key fast path");
68        }
69        let arr = self.downcast_iter().next().unwrap();
70        if arr.is_empty() {
71            return GroupsSlice::default();
72        }
73        let mut values = arr.values().as_slice();
74        let null_count = arr.null_count();
75        let length = values.len();
76
77        // all nulls
78        if null_count == length {
79            return vec![[0, length as IdxSize]];
80        }
81
82        let mut nulls_first = false;
83        if null_count > 0 {
84            nulls_first = arr.get(0).is_none()
85        }
86
87        if nulls_first {
88            values = &values[null_count..];
89        } else {
90            values = &values[..length - null_count];
91        };
92
93        let n_threads = POOL.current_num_threads();
94        if multithreaded && n_threads > 1 {
95            let parts =
96                create_clean_partitions(values, n_threads, self.is_sorted_descending_flag());
97            let n_parts = parts.len();
98
99            let first_ptr = &values[0] as *const T::Native as usize;
100            let groups = parts.par_iter().enumerate().map(|(i, part)| {
101                // we go via usize as *const is not send
102                let first_ptr = first_ptr as *const T::Native;
103
104                let part_first_ptr = &part[0] as *const T::Native;
105                let mut offset = unsafe { part_first_ptr.offset_from(first_ptr) } as IdxSize;
106
107                // nulls first: only add the nulls at the first partition
108                if nulls_first && i == 0 {
109                    partition_to_groups(part, null_count as IdxSize, true, offset)
110                }
111                // nulls last: only compute at the last partition
112                else if !nulls_first && i == n_parts - 1 {
113                    partition_to_groups(part, null_count as IdxSize, false, offset)
114                }
115                // other partitions
116                else {
117                    if nulls_first {
118                        offset += null_count as IdxSize;
119                    };
120
121                    partition_to_groups(part, 0, false, offset)
122                }
123            });
124            let groups = POOL.install(|| groups.collect::<Vec<_>>());
125            flatten_par(&groups)
126        } else {
127            partition_to_groups(values, null_count as IdxSize, nulls_first, 0)
128        }
129    }
130}
131
132#[cfg(all(feature = "dtype-categorical", feature = "performant"))]
133impl<T: PolarsCategoricalType> IntoGroupsType for CategoricalChunked<T>
134where
135    ChunkedArray<T::PolarsPhysical>: IntoGroupsType,
136{
137    fn group_tuples(&self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
138        self.phys.group_tuples(multithreaded, sorted)
139    }
140}
141
142impl<T> IntoGroupsType for ChunkedArray<T>
143where
144    T: PolarsNumericType,
145    T::Native: TotalHash + TotalEq + DirtyHash + ToTotalOrd,
146    <T::Native as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + DirtyHash,
147{
148    fn group_tuples(&self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
149        // sorted path
150        if self.is_sorted_ascending_flag() || self.is_sorted_descending_flag() {
151            // don't have to pass `sorted` arg, GroupSlice is always sorted.
152            return Ok(GroupsType::Slice {
153                groups: self.rechunk().create_groups_from_sorted(multithreaded),
154                rolling: false,
155            });
156        }
157
158        let out = match self.dtype() {
159            DataType::Float32 => {
160                // Convince the compiler that we are this type.
161                let ca: &Float32Chunked = unsafe {
162                    &*(self as *const ChunkedArray<T> as *const ChunkedArray<Float32Type>)
163                };
164                num_groups_proxy(ca, multithreaded, sorted)
165            },
166            DataType::Float64 => {
167                // Convince the compiler that we are this type.
168                let ca: &Float64Chunked = unsafe {
169                    &*(self as *const ChunkedArray<T> as *const ChunkedArray<Float64Type>)
170                };
171                num_groups_proxy(ca, multithreaded, sorted)
172            },
173            _ => match self.to_bit_repr() {
174                BitRepr::U8(ca) => num_groups_proxy(&ca, multithreaded, sorted),
175                BitRepr::U16(ca) => num_groups_proxy(&ca, multithreaded, sorted),
176                BitRepr::U32(ca) => num_groups_proxy(&ca, multithreaded, sorted),
177                BitRepr::U64(ca) => num_groups_proxy(&ca, multithreaded, sorted),
178                #[cfg(feature = "dtype-i128")]
179                BitRepr::I128(ca) => num_groups_proxy(&ca, multithreaded, sorted),
180            },
181        };
182        try_raise_keyboard_interrupt();
183        Ok(out)
184    }
185}
186impl IntoGroupsType for BooleanChunked {
187    fn group_tuples(&self, mut multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
188        multithreaded &= POOL.current_num_threads() > 1;
189
190        #[cfg(feature = "performant")]
191        {
192            let ca = self
193                .cast_with_options(&DataType::UInt8, CastOptions::Overflowing)
194                .unwrap();
195            let ca = ca.u8().unwrap();
196            ca.group_tuples(multithreaded, sorted)
197        }
198        #[cfg(not(feature = "performant"))]
199        {
200            let ca = self
201                .cast_with_options(&DataType::UInt32, CastOptions::Overflowing)
202                .unwrap();
203            let ca = ca.u32().unwrap();
204            ca.group_tuples(multithreaded, sorted)
205        }
206    }
207}
208
209impl IntoGroupsType for StringChunked {
210    #[allow(clippy::needless_lifetimes)]
211    fn group_tuples<'a>(&'a self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
212        self.as_binary().group_tuples(multithreaded, sorted)
213    }
214}
215
216impl IntoGroupsType for BinaryChunked {
217    #[allow(clippy::needless_lifetimes)]
218    fn group_tuples<'a>(
219        &'a self,
220        mut multithreaded: bool,
221        sorted: bool,
222    ) -> PolarsResult<GroupsType> {
223        if self.is_sorted_any() && !self.has_nulls() && self.n_chunks() == 1 {
224            let arr = self.downcast_get(0).unwrap();
225            let values = arr.values_iter();
226            let mut out = Vec::with_capacity(values.len() / 30);
227            partition_to_groups_amortized_varsize(values, arr.len() as _, 0, false, 0, &mut out);
228            return Ok(GroupsType::Slice {
229                groups: out,
230                rolling: false,
231            });
232        }
233
234        multithreaded &= POOL.current_num_threads() > 1;
235        let bh = self.to_bytes_hashes(multithreaded, Default::default());
236
237        let out = if multithreaded {
238            let n_partitions = bh.len();
239            // Take slices so that the vecs are not cloned.
240            let bh = bh.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
241            group_by_threaded_slice(bh, n_partitions, sorted)
242        } else {
243            group_by(bh[0].iter(), sorted)
244        };
245        try_raise_keyboard_interrupt();
246        Ok(out)
247    }
248}
249
250impl IntoGroupsType for BinaryOffsetChunked {
251    #[allow(clippy::needless_lifetimes)]
252    fn group_tuples<'a>(
253        &'a self,
254        mut multithreaded: bool,
255        sorted: bool,
256    ) -> PolarsResult<GroupsType> {
257        if self.is_sorted_any() && !self.has_nulls() && self.n_chunks() == 1 {
258            let arr = self.downcast_get(0).unwrap();
259            let values = arr.values_iter();
260            let mut out = Vec::with_capacity(values.len() / 30);
261            partition_to_groups_amortized_varsize(values, arr.len() as _, 0, false, 0, &mut out);
262            return Ok(GroupsType::Slice {
263                groups: out,
264                rolling: false,
265            });
266        }
267        multithreaded &= POOL.current_num_threads() > 1;
268        let bh = self.to_bytes_hashes(multithreaded, Default::default());
269
270        let out = if multithreaded {
271            let n_partitions = bh.len();
272            // Take slices so that the vecs are not cloned.
273            let bh = bh.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
274            group_by_threaded_slice(bh, n_partitions, sorted)
275        } else {
276            group_by(bh[0].iter(), sorted)
277        };
278        Ok(out)
279    }
280}
281
282impl IntoGroupsType for ListChunked {
283    #[allow(clippy::needless_lifetimes)]
284    #[allow(unused_variables)]
285    fn group_tuples<'a>(
286        &'a self,
287        mut multithreaded: bool,
288        sorted: bool,
289    ) -> PolarsResult<GroupsType> {
290        multithreaded &= POOL.current_num_threads() > 1;
291        let by = &[self.clone().into_column()];
292        let ca = if multithreaded {
293            encode_rows_vertical_par_unordered(by).unwrap()
294        } else {
295            _get_rows_encoded_ca_unordered(PlSmallStr::EMPTY, by).unwrap()
296        };
297
298        ca.group_tuples(multithreaded, sorted)
299    }
300}
301
302#[cfg(feature = "dtype-array")]
303impl IntoGroupsType for ArrayChunked {
304    #[allow(clippy::needless_lifetimes)]
305    #[allow(unused_variables)]
306    fn group_tuples<'a>(
307        &'a self,
308        mut multithreaded: bool,
309        sorted: bool,
310    ) -> PolarsResult<GroupsType> {
311        multithreaded &= POOL.current_num_threads() > 1;
312        let by = &[self.clone().into_column()];
313        let ca = if multithreaded {
314            encode_rows_vertical_par_unordered(by).unwrap()
315        } else {
316            _get_rows_encoded_ca_unordered(PlSmallStr::EMPTY, by).unwrap()
317        };
318        ca.group_tuples(multithreaded, sorted)
319    }
320}
321
322#[cfg(feature = "object")]
323impl<T> IntoGroupsType for ObjectChunked<T>
324where
325    T: PolarsObject,
326{
327    fn group_tuples(&self, _multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
328        Ok(group_by(self.into_iter(), sorted))
329    }
330}
polars_core/frame/group_by/into_groups.rs

polars_core/frame/group_by/
into_groups.rs
