polars_core/series/
amortized_iter.rs

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

polars_core/series/
amortized_iter.rs
