polars_core/frame/group_by/
position.rs

1use std::mem::ManuallyDrop;
2use std::ops::{Deref, DerefMut};
3
4use polars_arrow::offset::OffsetsBuffer;
5use polars_utils::idx_vec::IdxVec;
6use rayon::iter::plumbing::UnindexedConsumer;
7use rayon::prelude::*;
8
9use crate::prelude::*;
10use crate::runtime::RAYON;
11use crate::utils::{NoNull, flatten, slice_slice};
12
13/// Indexes of the groups, the first index is stored separately.
14/// this make sorting fast.
15#[derive(Debug, Clone, PartialEq, Eq, Default)]
16pub struct GroupsIdx {
17    pub(crate) sorted_by_first_idx: bool,
18    /// Positions of the start of each group.
19    first: Vec<IdxSize>,
20    /// Global positions of all elements of all groups.
21    all: Vec<IdxVec>,
22}
23
24pub type IdxItem = (IdxSize, IdxVec);
25pub type BorrowIdxItem<'a> = (IdxSize, &'a IdxVec);
26
27impl Drop for GroupsIdx {
28    fn drop(&mut self) {
29        let v = std::mem::take(&mut self.all);
30        // ~65k took approximately 1ms on local machine, so from that point we drop on other thread
31        // to stop query from being blocked
32        #[cfg(not(target_family = "wasm"))]
33        if v.len() > 1 << 16 {
34            std::thread::spawn(move || drop(v));
35        } else {
36            drop(v);
37        }
38
39        #[cfg(target_family = "wasm")]
40        drop(v);
41    }
42}
43
44impl From<Vec<IdxItem>> for GroupsIdx {
45    fn from(v: Vec<IdxItem>) -> Self {
46        v.into_iter().collect()
47    }
48}
49
50impl From<Vec<Vec<IdxItem>>> for GroupsIdx {
51    fn from(v: Vec<Vec<IdxItem>>) -> Self {
52        // single threaded flatten: 10% faster than `iter().flatten().collect()
53        // this is the multi-threaded impl of that
54        let (cap, offsets) = flatten::cap_and_offsets(&v);
55        let mut first = Vec::with_capacity(cap);
56        let first_ptr = first.as_ptr() as usize;
57        let mut all = Vec::with_capacity(cap);
58        let all_ptr = all.as_ptr() as usize;
59
60        RAYON.install(|| {
61            v.into_par_iter()
62                .zip(offsets)
63                .for_each(|(mut inner, offset)| {
64                    unsafe {
65                        let first = (first_ptr as *const IdxSize as *mut IdxSize).add(offset);
66                        let all = (all_ptr as *const IdxVec as *mut IdxVec).add(offset);
67
68                        let inner_ptr = inner.as_mut_ptr();
69                        for i in 0..inner.len() {
70                            let (first_val, vals) = std::ptr::read(inner_ptr.add(i));
71                            std::ptr::write(first.add(i), first_val);
72                            std::ptr::write(all.add(i), vals);
73                        }
74                        // set len to 0 so that the contents will not get dropped
75                        // they are moved to `first` and `all`
76                        inner.set_len(0);
77                    }
78                });
79        });
80        unsafe {
81            all.set_len(cap);
82            first.set_len(cap);
83        }
84        GroupsIdx {
85            sorted_by_first_idx: false,
86            first,
87            all,
88        }
89    }
90}
91
92impl GroupsIdx {
93    pub fn new(first: Vec<IdxSize>, all: Vec<IdxVec>, sorted_by_first_idx: bool) -> Self {
94        Self {
95            sorted_by_first_idx,
96            first,
97            all,
98        }
99    }
100
101    pub fn sort_by_first_idx(&mut self) {
102        if self.sorted_by_first_idx {
103            return;
104        }
105        let mut idx = 0;
106        let first = std::mem::take(&mut self.first);
107        // store index and values so that we can sort those
108        let mut idx_vals = first
109            .into_iter()
110            .map(|v| {
111                let out = [idx, v];
112                idx += 1;
113                out
114            })
115            .collect_trusted::<Vec<_>>();
116        idx_vals.sort_unstable_by_key(|v| v[1]);
117
118        let take_first = || idx_vals.iter().map(|v| v[1]).collect_trusted::<Vec<_>>();
119        let take_all = || {
120            idx_vals
121                .iter()
122                .map(|v| unsafe {
123                    let idx = v[0] as usize;
124                    std::mem::take(self.all.get_unchecked_mut(idx))
125                })
126                .collect_trusted::<Vec<_>>()
127        };
128        let (first, all) = RAYON.install(|| rayon::join(take_first, take_all));
129        self.first = first;
130        self.all = all;
131        self.sorted_by_first_idx = true
132    }
133    pub fn is_sorted_by_first_idx(&self) -> bool {
134        self.sorted_by_first_idx
135    }
136
137    pub fn iter(
138        &self,
139    ) -> std::iter::Zip<
140        std::iter::Copied<std::slice::Iter<'_, IdxSize>>,
141        std::slice::Iter<'_, IdxVec>,
142    > {
143        self.into_iter()
144    }
145
146    #[inline(always)]
147    pub fn all(&self) -> &[IdxVec] {
148        &self.all
149    }
150
151    #[inline(always)]
152    pub fn first(&self) -> &[IdxSize] {
153        &self.first
154    }
155
156    #[inline(always)]
157    pub(crate) fn len(&self) -> usize {
158        self.first.len()
159    }
160
161    #[inline(always)]
162    pub(crate) unsafe fn get_unchecked(&self, index: usize) -> BorrowIdxItem<'_> {
163        let first = *self.first.get_unchecked(index);
164        let all = self.all.get_unchecked(index);
165        (first, all)
166    }
167
168    // Create an 'empty group', containing 1 group of length 0
169    pub fn new_empty() -> Self {
170        Self {
171            sorted_by_first_idx: false,
172            first: vec![0],
173            all: vec![vec![].into()],
174        }
175    }
176}
177
178impl FromIterator<IdxItem> for GroupsIdx {
179    fn from_iter<T: IntoIterator<Item = IdxItem>>(iter: T) -> Self {
180        let (first, all) = iter.into_iter().unzip();
181        GroupsIdx {
182            sorted_by_first_idx: false,
183            first,
184            all,
185        }
186    }
187}
188
189impl<'a> IntoIterator for &'a GroupsIdx {
190    type Item = BorrowIdxItem<'a>;
191    type IntoIter = std::iter::Zip<
192        std::iter::Copied<std::slice::Iter<'a, IdxSize>>,
193        std::slice::Iter<'a, IdxVec>,
194    >;
195
196    fn into_iter(self) -> Self::IntoIter {
197        self.first.iter().copied().zip(self.all.iter())
198    }
199}
200
201impl IntoIterator for GroupsIdx {
202    type Item = IdxItem;
203    type IntoIter = std::iter::Zip<std::vec::IntoIter<IdxSize>, std::vec::IntoIter<IdxVec>>;
204
205    fn into_iter(mut self) -> Self::IntoIter {
206        let first = std::mem::take(&mut self.first);
207        let all = std::mem::take(&mut self.all);
208        first.into_iter().zip(all)
209    }
210}
211
212impl FromParallelIterator<IdxItem> for GroupsIdx {
213    fn from_par_iter<I>(par_iter: I) -> Self
214    where
215        I: IntoParallelIterator<Item = IdxItem>,
216    {
217        let (first, all) = par_iter.into_par_iter().unzip();
218        GroupsIdx {
219            sorted_by_first_idx: false,
220            first,
221            all,
222        }
223    }
224}
225
226impl<'a> IntoParallelIterator for &'a GroupsIdx {
227    type Iter = rayon::iter::Zip<
228        rayon::iter::Copied<rayon::slice::Iter<'a, IdxSize>>,
229        rayon::slice::Iter<'a, IdxVec>,
230    >;
231    type Item = BorrowIdxItem<'a>;
232
233    fn into_par_iter(self) -> Self::Iter {
234        self.first.par_iter().copied().zip(self.all.par_iter())
235    }
236}
237
238impl IntoParallelIterator for GroupsIdx {
239    type Iter = rayon::iter::Zip<rayon::vec::IntoIter<IdxSize>, rayon::vec::IntoIter<IdxVec>>;
240    type Item = IdxItem;
241
242    fn into_par_iter(mut self) -> Self::Iter {
243        let first = std::mem::take(&mut self.first);
244        let all = std::mem::take(&mut self.all);
245        first.into_par_iter().zip(all.into_par_iter())
246    }
247}
248
249/// Every group is indicated by an array where the
250///  - first value is an index to the start of the group
251///  - second value is the length of the group
252///
253/// Only used when group values are stored together
254///
255/// This type should have the invariant that it is always sorted in ascending
256/// order by the start indices.
257pub type GroupsSlice = Vec<[IdxSize; 2]>;
258
259#[derive(Debug, Clone, PartialEq, Eq)]
260pub enum GroupsType {
261    Idx(GroupsIdx),
262    Slice {
263        // the groups slices
264        groups: GroupsSlice,
265        /// Indicates if the groups may overlap, i.e., at least one index MAY be
266        /// included in more than one group slice.
267        overlapping: bool,
268        /// Indicates if the groups are rolling, i.e. for every consecutive group
269        /// slice (offset, len), and given start = offset and end = offset + len,
270        /// then both new_start >= start AND new_end >= end MUST be true.
271        monotonic: bool,
272    },
273}
274
275impl Default for GroupsType {
276    fn default() -> Self {
277        GroupsType::Idx(GroupsIdx::default())
278    }
279}
280
281/// Returns whether `groups` is monotonic:
282/// True if for every consecutive pair of group slices both the
283/// start and the end offset are non-decreasing.
284pub fn slice_groups_are_monotonic(groups: &GroupsSlice) -> bool {
285    if groups.len() < 2 {
286        return true;
287    }
288
289    let (offset, len) = (groups[0][0], groups[0][1]);
290    let mut prev_start = offset;
291    let mut prev_end = offset + len;
292
293    for g in &groups[1..] {
294        let start = g[0];
295        let end = g[0] + g[1];
296
297        if start < prev_start || end < prev_end {
298            return false;
299        }
300
301        prev_start = start;
302        prev_end = end;
303    }
304    true
305}
306
307impl GroupsType {
308    pub fn new_slice(groups: GroupsSlice, overlapping: bool, monotonic: bool) -> Self {
309        #[cfg(debug_assertions)]
310        {
311            fn groups_overlap(groups: &GroupsSlice) -> bool {
312                if groups.len() < 2 {
313                    return false;
314                }
315                let mut groups = groups.clone();
316                groups.sort();
317                let mut prev_end = groups[0][1];
318
319                for g in &groups[1..] {
320                    let start = g[0];
321                    let end = g[1];
322                    if start < prev_end {
323                        return true;
324                    }
325                    if end > prev_end {
326                        prev_end = end;
327                    }
328                }
329                false
330            }
331
332            assert!(overlapping || !groups_overlap(&groups));
333            assert!(!monotonic || slice_groups_are_monotonic(&groups));
334        }
335
336        Self::Slice {
337            groups,
338            overlapping,
339            monotonic,
340        }
341    }
342
343    pub fn into_idx(self) -> GroupsIdx {
344        match self {
345            GroupsType::Idx(groups) => groups,
346            GroupsType::Slice { groups, .. } => {
347                polars_warn!(
348                    "Had to reallocate groups, missed an optimization opportunity. Please open an issue."
349                );
350                groups
351                    .iter()
352                    .map(|&[first, len]| (first, (first..first + len).collect::<IdxVec>()))
353                    .collect()
354            },
355        }
356    }
357
358    pub(crate) fn prepare_list_agg(
359        &self,
360        total_len: usize,
361    ) -> (Option<IdxCa>, OffsetsBuffer<i64>, bool) {
362        let mut can_fast_explode = true;
363        match self {
364            GroupsType::Idx(groups) => {
365                let mut list_offset = Vec::with_capacity(self.len() + 1);
366                let mut gather_offsets = Vec::with_capacity(total_len);
367
368                let mut len_so_far = 0i64;
369                list_offset.push(len_so_far);
370
371                for idx in groups {
372                    let idx = idx.1;
373                    gather_offsets.extend_from_slice(idx);
374                    len_so_far += idx.len() as i64;
375                    list_offset.push(len_so_far);
376                    can_fast_explode &= !idx.is_empty();
377                }
378                unsafe {
379                    (
380                        Some(IdxCa::from_vec(PlSmallStr::EMPTY, gather_offsets)),
381                        OffsetsBuffer::new_unchecked(list_offset.into()),
382                        can_fast_explode,
383                    )
384                }
385            },
386            GroupsType::Slice { groups, .. } => {
387                let mut list_offset = Vec::with_capacity(self.len() + 1);
388                let mut gather_offsets = Vec::with_capacity(total_len);
389                let mut len_so_far = 0i64;
390                list_offset.push(len_so_far);
391
392                for g in groups {
393                    let len = g[1];
394                    let offset = g[0];
395                    gather_offsets.extend(offset..offset + len);
396
397                    len_so_far += len as i64;
398                    list_offset.push(len_so_far);
399                    can_fast_explode &= len > 0;
400                }
401
402                unsafe {
403                    (
404                        Some(IdxCa::from_vec(PlSmallStr::EMPTY, gather_offsets)),
405                        OffsetsBuffer::new_unchecked(list_offset.into()),
406                        can_fast_explode,
407                    )
408                }
409            },
410        }
411    }
412
413    pub fn iter(&self) -> GroupsTypeIter<'_> {
414        GroupsTypeIter::new(self)
415    }
416
417    pub fn sort_by_first_idx(&mut self) {
418        match self {
419            GroupsType::Idx(groups) => {
420                if !groups.is_sorted_by_first_idx() {
421                    groups.sort_by_first_idx()
422                }
423            },
424            GroupsType::Slice { .. } => {
425                // Only reached for non-overlapping slices, which are ordered by start.
426            },
427        }
428    }
429
430    pub(crate) fn is_sorted_by_first_idx(&self) -> bool {
431        match self {
432            GroupsType::Idx(groups) => groups.is_sorted_by_first_idx(),
433            // Overlapping slices can step back, e.g. after a per-window `slice` offset.
434            GroupsType::Slice {
435                overlapping,
436                monotonic,
437                ..
438            } => !*overlapping || *monotonic,
439        }
440    }
441
442    pub fn is_overlapping(&self) -> bool {
443        matches!(
444            self,
445            GroupsType::Slice {
446                overlapping: true,
447                ..
448            }
449        )
450    }
451
452    pub fn is_monotonic(&self) -> bool {
453        matches!(
454            self,
455            GroupsType::Slice {
456                monotonic: true,
457                ..
458            }
459        )
460    }
461
462    pub fn take_group_firsts(self) -> Vec<IdxSize> {
463        match self {
464            GroupsType::Idx(mut groups) => std::mem::take(&mut groups.first),
465            GroupsType::Slice { groups, .. } => {
466                groups.into_iter().map(|[first, _len]| first).collect()
467            },
468        }
469    }
470
471    /// Checks if groups are of equal length. The caller is responsible for
472    /// updating the groups by calling `groups()` prior to calling this method.
473    pub fn check_lengths(self: &GroupsType, other: &GroupsType) -> PolarsResult<()> {
474        if std::ptr::eq(self, other) {
475            return Ok(());
476        }
477        polars_ensure!(self.iter().zip(other.iter()).all(|(a, b)| {
478            a.len() == b.len()
479        }), ShapeMismatch: "expressions must have matching group lengths");
480        Ok(())
481    }
482
483    /// # Safety
484    /// This will not do any bounds checks. The caller must ensure
485    /// all groups have members.
486    pub unsafe fn take_group_lasts(self) -> Vec<IdxSize> {
487        match self {
488            GroupsType::Idx(groups) => groups
489                .all
490                .iter()
491                .map(|idx| *idx.get_unchecked(idx.len() - 1))
492                .collect(),
493            GroupsType::Slice { groups, .. } => groups
494                .into_iter()
495                .map(|[first, len]| first + len - 1)
496                .collect(),
497        }
498    }
499
500    pub fn par_iter(&self) -> GroupsTypeParIter<'_> {
501        GroupsTypeParIter::new(self)
502    }
503
504    /// Get a reference to the `GroupsIdx`.
505    ///
506    /// # Panic
507    ///
508    /// panics if the groups are a slice.
509    pub fn unwrap_idx(&self) -> &GroupsIdx {
510        match self {
511            GroupsType::Idx(groups) => groups,
512            GroupsType::Slice { .. } => panic!("groups are slices not index"),
513        }
514    }
515
516    /// Get a reference to the `GroupsSlice`.
517    ///
518    /// # Panic
519    ///
520    /// panics if the groups are an idx.
521    pub fn unwrap_slice(&self) -> &GroupsSlice {
522        match self {
523            GroupsType::Slice { groups, .. } => groups,
524            GroupsType::Idx(_) => panic!("groups are index not slices"),
525        }
526    }
527
528    pub fn get(&self, index: usize) -> GroupsIndicator<'_> {
529        match self {
530            GroupsType::Idx(groups) => {
531                let first = groups.first[index];
532                let all = &groups.all[index];
533                GroupsIndicator::Idx((first, all))
534            },
535            GroupsType::Slice { groups, .. } => GroupsIndicator::Slice(groups[index]),
536        }
537    }
538
539    /// Get a mutable reference to the `GroupsIdx`.
540    ///
541    /// # Panic
542    ///
543    /// panics if the groups are a slice.
544    pub fn idx_mut(&mut self) -> &mut GroupsIdx {
545        match self {
546            GroupsType::Idx(groups) => groups,
547            GroupsType::Slice { .. } => panic!("groups are slices not index"),
548        }
549    }
550
551    pub fn len(&self) -> usize {
552        match self {
553            GroupsType::Idx(groups) => groups.len(),
554            GroupsType::Slice { groups, .. } => groups.len(),
555        }
556    }
557
558    pub fn is_empty(&self) -> bool {
559        self.len() == 0
560    }
561
562    pub fn group_count(&self) -> IdxCa {
563        match self {
564            GroupsType::Idx(groups) => {
565                let ca: NoNull<IdxCa> = groups
566                    .iter()
567                    .map(|(_first, idx)| idx.len() as IdxSize)
568                    .collect_trusted();
569                ca.into_inner()
570            },
571            GroupsType::Slice { groups, .. } => {
572                let ca: NoNull<IdxCa> = groups.iter().map(|[_first, len]| *len).collect_trusted();
573                ca.into_inner()
574            },
575        }
576    }
577
578    pub fn as_list_chunked(&self) -> ListChunked {
579        match self {
580            GroupsType::Idx(groups) => groups
581                .iter()
582                .map(|(_first, idx)| {
583                    let ca: NoNull<IdxCa> = idx.iter().map(|&v| v as IdxSize).collect();
584                    ca.into_inner().into_series()
585                })
586                .collect_trusted(),
587            GroupsType::Slice { groups, .. } => groups
588                .iter()
589                .map(|&[first, len]| {
590                    let ca: NoNull<IdxCa> = (first..first + len).collect_trusted();
591                    ca.into_inner().into_series()
592                })
593                .collect_trusted(),
594        }
595    }
596
597    pub fn into_sliceable(self) -> GroupPositions {
598        let len = self.len();
599        slice_groups(Arc::new(self), 0, len)
600    }
601
602    pub fn num_elements(&self) -> usize {
603        match self {
604            GroupsType::Idx(i) => i.all().iter().map(|v| v.len()).sum(),
605            GroupsType::Slice {
606                groups,
607                overlapping: _,
608                monotonic: _,
609            } => groups.iter().map(|[_, l]| *l as usize).sum(),
610        }
611    }
612}
613
614impl From<GroupsIdx> for GroupsType {
615    fn from(groups: GroupsIdx) -> Self {
616        GroupsType::Idx(groups)
617    }
618}
619
620pub enum GroupsIndicator<'a> {
621    Idx(BorrowIdxItem<'a>),
622    Slice([IdxSize; 2]),
623}
624
625impl GroupsIndicator<'_> {
626    pub fn len(&self) -> usize {
627        match self {
628            GroupsIndicator::Idx(g) => g.1.len(),
629            GroupsIndicator::Slice([_, len]) => *len as usize,
630        }
631    }
632    pub fn first(&self) -> IdxSize {
633        match self {
634            GroupsIndicator::Idx(g) => g.0,
635            GroupsIndicator::Slice([first, _]) => *first,
636        }
637    }
638    pub fn is_empty(&self) -> bool {
639        self.len() == 0
640    }
641}
642
643pub struct GroupsTypeIter<'a> {
644    vals: &'a GroupsType,
645    len: usize,
646    idx: usize,
647}
648
649impl<'a> GroupsTypeIter<'a> {
650    fn new(vals: &'a GroupsType) -> Self {
651        let len = vals.len();
652        let idx = 0;
653        GroupsTypeIter { vals, len, idx }
654    }
655}
656
657impl<'a> Iterator for GroupsTypeIter<'a> {
658    type Item = GroupsIndicator<'a>;
659
660    fn nth(&mut self, n: usize) -> Option<Self::Item> {
661        self.idx = self.idx.saturating_add(n);
662        self.next()
663    }
664
665    fn next(&mut self) -> Option<Self::Item> {
666        if self.idx >= self.len {
667            return None;
668        }
669
670        let out = unsafe {
671            match self.vals {
672                GroupsType::Idx(groups) => {
673                    let item = groups.get_unchecked(self.idx);
674                    Some(GroupsIndicator::Idx(item))
675                },
676                GroupsType::Slice { groups, .. } => {
677                    Some(GroupsIndicator::Slice(*groups.get_unchecked(self.idx)))
678                },
679            }
680        };
681        self.idx += 1;
682        out
683    }
684}
685
686pub struct GroupsTypeParIter<'a> {
687    vals: &'a GroupsType,
688    len: usize,
689}
690
691impl<'a> GroupsTypeParIter<'a> {
692    fn new(vals: &'a GroupsType) -> Self {
693        let len = vals.len();
694        GroupsTypeParIter { vals, len }
695    }
696}
697
698impl<'a> ParallelIterator for GroupsTypeParIter<'a> {
699    type Item = GroupsIndicator<'a>;
700
701    fn drive_unindexed<C>(self, consumer: C) -> C::Result
702    where
703        C: UnindexedConsumer<Self::Item>,
704    {
705        (0..self.len)
706            .into_par_iter()
707            .map(|i| unsafe {
708                match self.vals {
709                    GroupsType::Idx(groups) => GroupsIndicator::Idx(groups.get_unchecked(i)),
710                    GroupsType::Slice { groups, .. } => {
711                        GroupsIndicator::Slice(*groups.get_unchecked(i))
712                    },
713                }
714            })
715            .drive_unindexed(consumer)
716    }
717}
718
719#[derive(Debug)]
720pub struct GroupPositions {
721    // SAFETY: sliced is a shallow clone of original
722    // It emulates a shared reference, not an exclusive reference
723    // Its data must not be mutated through direct access
724    sliced: ManuallyDrop<GroupsType>,
725    // Unsliced buffer
726    original: Arc<GroupsType>,
727    offset: i64,
728    len: usize,
729}
730
731impl Clone for GroupPositions {
732    fn clone(&self) -> Self {
733        let sliced = slice_groups_inner(&self.original, self.offset, self.len);
734
735        Self {
736            sliced,
737            original: self.original.clone(),
738            offset: self.offset,
739            len: self.len,
740        }
741    }
742}
743
744impl AsRef<GroupsType> for GroupPositions {
745    fn as_ref(&self) -> &GroupsType {
746        self.sliced.deref()
747    }
748}
749
750impl Deref for GroupPositions {
751    type Target = GroupsType;
752
753    fn deref(&self) -> &Self::Target {
754        self.sliced.deref()
755    }
756}
757
758impl Default for GroupPositions {
759    fn default() -> Self {
760        GroupsType::default().into_sliceable()
761    }
762}
763
764impl GroupPositions {
765    pub fn slice(&self, offset: i64, len: usize) -> Self {
766        let offset = self.offset + offset;
767        slice_groups(self.original.clone(), offset, len)
768    }
769
770    pub fn sort_by_first_idx(&mut self) {
771        if !self.as_ref().is_sorted_by_first_idx() {
772            let original = Arc::make_mut(&mut self.original);
773            original.sort_by_first_idx();
774
775            self.sliced = slice_groups_inner(original, self.offset, self.len);
776        }
777    }
778
779    pub fn unroll(mut self) -> GroupPositions {
780        match self.sliced.deref_mut() {
781            GroupsType::Idx(_) => self,
782            GroupsType::Slice {
783                overlapping: false, ..
784            } => self,
785            GroupsType::Slice { groups, .. } => {
786                // SAFETY: sliced is a shallow partial clone of original.
787                // A new owning Vec is required per GH issue #21859
788                let mut cum_offset = 0 as IdxSize;
789                let groups: Vec<_> = groups
790                    .iter()
791                    .map(|[_, len]| {
792                        let new = [cum_offset, *len];
793                        cum_offset += *len;
794                        new
795                    })
796                    .collect();
797
798                GroupsType::new_slice(groups, false, true).into_sliceable()
799            },
800        }
801    }
802
803    pub fn as_unrolled_slice(&self) -> Option<&GroupsSlice> {
804        match &*self.sliced {
805            GroupsType::Idx(_) => None,
806            GroupsType::Slice {
807                groups: _,
808                overlapping: true,
809                monotonic: _,
810            } => None,
811            GroupsType::Slice {
812                groups,
813                overlapping: false,
814                monotonic: _,
815            } => Some(groups),
816        }
817    }
818
819    /// Compare groups based on inner pointer.
820    pub fn is_same(&self, other: &Self) -> bool {
821        Arc::ptr_eq(&self.original, &other.original)
822            && self.offset == other.offset
823            && self.len == other.len
824    }
825}
826
827fn slice_groups_inner(g: &GroupsType, offset: i64, len: usize) -> ManuallyDrop<GroupsType> {
828    // SAFETY:
829    // we create new `Vec`s from the sliced groups. But we wrap them in ManuallyDrop
830    // so that we never call drop on them.
831    // These groups lifetimes are bounded to the `g`. This must remain valid
832    // for the scope of the aggregation.
833    match g {
834        GroupsType::Idx(groups) => {
835            let first = unsafe {
836                let first = slice_slice(groups.first(), offset, len);
837                let ptr = first.as_ptr() as *mut _;
838                Vec::from_raw_parts(ptr, first.len(), first.len())
839            };
840
841            let all = unsafe {
842                let all = slice_slice(groups.all(), offset, len);
843                let ptr = all.as_ptr() as *mut _;
844                Vec::from_raw_parts(ptr, all.len(), all.len())
845            };
846            ManuallyDrop::new(GroupsType::Idx(GroupsIdx::new(
847                first,
848                all,
849                groups.is_sorted_by_first_idx(),
850            )))
851        },
852        GroupsType::Slice {
853            groups,
854            overlapping,
855            monotonic,
856        } => {
857            let groups = unsafe {
858                let groups = slice_slice(groups, offset, len);
859                let ptr = groups.as_ptr() as *mut _;
860                Vec::from_raw_parts(ptr, groups.len(), groups.len())
861            };
862
863            ManuallyDrop::new(GroupsType::new_slice(groups, *overlapping, *monotonic))
864        },
865    }
866}
867
868fn slice_groups(g: Arc<GroupsType>, offset: i64, len: usize) -> GroupPositions {
869    let sliced = slice_groups_inner(g.as_ref(), offset, len);
870
871    GroupPositions {
872        sliced,
873        original: g,
874        offset,
875        len,
876    }
877}
polars_core/frame/group_by/position.rs

polars_core/frame/group_by/
position.rs
