polars_core/series/
amortized_iter.rs

1use std::rc::Rc;
2
3use crate::prelude::*;
4
5/// A [`Series`] that amortizes a few allocations during iteration.
6#[derive(Clone)]
7pub struct AmortSeries {
8    container: Rc<Series>,
9}
10
11/// We don't implement Deref so that the caller is aware of converting to Series
12impl AsRef<Series> for AmortSeries {
13    fn as_ref(&self) -> &Series {
14        self.container.as_ref()
15    }
16}
17
18pub type ArrayBox = Box<dyn Array>;
19
20impl AmortSeries {
21    pub fn new(container: Rc<Series>) -> Self {
22        debug_assert_eq!(container.chunks().len(), 1);
23        AmortSeries { container }
24    }
25
26    pub fn deep_clone(&self) -> Series {
27        unsafe {
28            let s = &(*self.container);
29            debug_assert_eq!(s.chunks().len(), 1);
30            let array_ref = s.chunks().get_unchecked(0).clone();
31            let name = s.name().clone();
32            Series::from_chunks_and_dtype_unchecked(name, vec![array_ref], s.dtype())
33        }
34    }
35
36    #[inline]
37    /// Swaps inner state with the `array`. Prefer `AmortSeries::with_array` as this
38    /// restores the state.
39    /// # Safety
40    /// This swaps an underlying pointer that might be hold by other cloned series.
41    pub unsafe fn swap(&mut self, array: &mut ArrayRef) {
42        let inner = self.container.array_ref(0) as *const ArrayRef as *mut ArrayRef;
43        let inner = inner.as_mut().unwrap();
44        std::mem::swap(inner, array);
45
46        // ensure lengths are correct.
47        unsafe {
48            let ptr = Rc::as_ptr(&self.container) as *mut Series;
49            (*ptr)._get_inner_mut().compute_len();
50        }
51    }
52
53    /// Temporary swaps out the array, and restores the original state
54    /// when application of the function `f` is done.
55    ///
56    /// # Safety
57    /// Array must be from `Series` physical dtype.
58    #[inline]
59    pub unsafe fn with_array<F, T>(&mut self, array: &mut ArrayRef, f: F) -> T
60    where
61        F: Fn(&AmortSeries) -> T,
62    {
63        unsafe {
64            self.swap(array);
65            let out = f(self);
66            self.swap(array);
67            out
68        }
69    }
70}
71
72// SAFETY:
73// type must be matching
74pub(crate) unsafe fn unstable_series_container_and_ptr(
75    name: PlSmallStr,
76    inner_values: ArrayRef,
77    iter_dtype: &DataType,
78) -> (Series, *mut ArrayRef) {
79    let series_container = {
80        let mut s = Series::from_chunks_and_dtype_unchecked(name, vec![inner_values], iter_dtype);
81        s.clear_flags();
82        s
83    };
84
85    let ptr = series_container.array_ref(0) as *const ArrayRef as *mut ArrayRef;
86    (series_container, ptr)
87}
polars_core/series/amortized_iter.rs

polars_core/series/
amortized_iter.rs
