polars_utils/
functions.rs

1use std::mem::MaybeUninit;
2use std::ops::Range;
3use std::sync::Arc;
4
5// The ith portion of a range split in k (as equal as possible) parts.
6#[inline(always)]
7pub fn range_portion(i: usize, k: usize, r: Range<usize>) -> Range<usize> {
8    // Each portion having size n / k leaves n % k elements unaccounted for.
9    // Make the first n % k portions have 1 extra element.
10    let n = r.len();
11    let base_size = n / k;
12    let num_one_larger = n % k;
13    let num_before = base_size * i + i.min(num_one_larger);
14    let our_size = base_size + (i < num_one_larger) as usize;
15    r.start + num_before..r.start + num_before + our_size
16}
17
18pub fn arc_map<T: Clone, F: FnMut(T) -> T>(mut arc: Arc<T>, mut f: F) -> Arc<T> {
19    unsafe {
20        // Make the Arc unique (cloning if necessary).
21        Arc::make_mut(&mut arc);
22
23        // If f panics we must be able to drop the Arc without assuming it is initialized.
24        let mut uninit_arc = Arc::from_raw(Arc::into_raw(arc).cast::<MaybeUninit<T>>());
25
26        // Replace the value inside the arc.
27        let ptr = Arc::get_mut(&mut uninit_arc).unwrap_unchecked() as *mut MaybeUninit<T>;
28        *ptr = MaybeUninit::new(f(ptr.read().assume_init()));
29
30        // Now the Arc is properly initialized again.
31        Arc::from_raw(Arc::into_raw(uninit_arc).cast::<T>())
32    }
33}
34
35pub fn try_arc_map<T: Clone, E, F: FnMut(T) -> Result<T, E>>(
36    mut arc: Arc<T>,
37    mut f: F,
38) -> Result<Arc<T>, E> {
39    unsafe {
40        // Make the Arc unique (cloning if necessary).
41        Arc::make_mut(&mut arc);
42
43        // If f panics we must be able to drop the Arc without assuming it is initialized.
44        let mut uninit_arc = Arc::from_raw(Arc::into_raw(arc).cast::<MaybeUninit<T>>());
45
46        // Replace the value inside the arc.
47        let ptr = Arc::get_mut(&mut uninit_arc).unwrap_unchecked() as *mut MaybeUninit<T>;
48        *ptr = MaybeUninit::new(f(ptr.read().assume_init())?);
49
50        // Now the Arc is properly initialized again.
51        Ok(Arc::from_raw(Arc::into_raw(uninit_arc).cast::<T>()))
52    }
53}