polars_core/frame/group_by/
into_groups.rs

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

polars_core/frame/group_by/
into_groups.rs
