polars_core/series/
amortized_iter.rs

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

polars_core/series/
amortized_iter.rs
