polars_core/series/implementations/
categorical.rs

1use super::*;
2use crate::chunked_array::comparison::*;
3use crate::prelude::*;
4
5unsafe impl IntoSeries for CategoricalChunked {
6    fn into_series(self) -> Series {
7        Series(Arc::new(SeriesWrap(self)))
8    }
9}
10
11impl SeriesWrap<CategoricalChunked> {
12    fn finish_with_state(&self, keep_fast_unique: bool, cats: UInt32Chunked) -> CategoricalChunked {
13        let mut out = unsafe {
14            CategoricalChunked::from_cats_and_rev_map_unchecked(
15                cats,
16                self.0.get_rev_map().clone(),
17                self.0.is_enum(),
18                self.0.get_ordering(),
19            )
20        };
21        if keep_fast_unique && self.0._can_fast_unique() {
22            out.set_fast_unique(true)
23        }
24        out
25    }
26
27    fn with_state<F>(&self, keep_fast_unique: bool, apply: F) -> CategoricalChunked
28    where
29        F: Fn(&UInt32Chunked) -> UInt32Chunked,
30    {
31        let cats = apply(self.0.physical());
32        self.finish_with_state(keep_fast_unique, cats)
33    }
34
35    fn try_with_state<'a, F>(
36        &'a self,
37        keep_fast_unique: bool,
38        apply: F,
39    ) -> PolarsResult<CategoricalChunked>
40    where
41        F: for<'b> Fn(&'a UInt32Chunked) -> PolarsResult<UInt32Chunked>,
42    {
43        let cats = apply(self.0.physical())?;
44        Ok(self.finish_with_state(keep_fast_unique, cats))
45    }
46}
47
48impl private::PrivateSeries for SeriesWrap<CategoricalChunked> {
49    fn compute_len(&mut self) {
50        self.0.physical_mut().compute_len()
51    }
52    fn _field(&self) -> Cow<Field> {
53        Cow::Owned(self.0.field())
54    }
55    fn _dtype(&self) -> &DataType {
56        self.0.dtype()
57    }
58    fn _get_flags(&self) -> StatisticsFlags {
59        self.0.get_flags()
60    }
61    fn _set_flags(&mut self, flags: StatisticsFlags) {
62        self.0.set_flags(flags)
63    }
64
65    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
66        self.0.physical().equal_element(idx_self, idx_other, other)
67    }
68
69    #[cfg(feature = "zip_with")]
70    fn zip_with_same_type(&self, mask: &BooleanChunked, other: &Series) -> PolarsResult<Series> {
71        self.0
72            .zip_with(mask, other.categorical()?)
73            .map(|ca| ca.into_series())
74    }
75    fn into_total_ord_inner<'a>(&'a self) -> Box<dyn TotalOrdInner + 'a> {
76        if self.0.uses_lexical_ordering() {
77            (&self.0).into_total_ord_inner()
78        } else {
79            self.0.physical().into_total_ord_inner()
80        }
81    }
82    fn into_total_eq_inner<'a>(&'a self) -> Box<dyn TotalEqInner + 'a> {
83        invalid_operation_panic!(into_total_eq_inner, self)
84    }
85
86    fn vec_hash(
87        &self,
88        random_state: PlSeedableRandomStateQuality,
89        buf: &mut Vec<u64>,
90    ) -> PolarsResult<()> {
91        self.0.physical().vec_hash(random_state, buf)?;
92        Ok(())
93    }
94
95    fn vec_hash_combine(
96        &self,
97        build_hasher: PlSeedableRandomStateQuality,
98        hashes: &mut [u64],
99    ) -> PolarsResult<()> {
100        self.0.physical().vec_hash_combine(build_hasher, hashes)?;
101        Ok(())
102    }
103
104    #[cfg(feature = "algorithm_group_by")]
105    unsafe fn agg_list(&self, groups: &GroupsType) -> Series {
106        // we cannot cast and dispatch as the inner type of the list would be incorrect
107        let list = self.0.physical().agg_list(groups);
108        let mut list = list.list().unwrap().clone();
109        unsafe { list.to_logical(self.dtype().clone()) };
110        list.into_series()
111    }
112
113    #[cfg(feature = "algorithm_group_by")]
114    fn group_tuples(&self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
115        #[cfg(feature = "performant")]
116        {
117            Ok(self.0.group_tuples_perfect(multithreaded, sorted))
118        }
119        #[cfg(not(feature = "performant"))]
120        {
121            self.0.physical().group_tuples(multithreaded, sorted)
122        }
123    }
124
125    fn arg_sort_multiple(
126        &self,
127        by: &[Column],
128        options: &SortMultipleOptions,
129    ) -> PolarsResult<IdxCa> {
130        self.0.arg_sort_multiple(by, options)
131    }
132}
133
134impl SeriesTrait for SeriesWrap<CategoricalChunked> {
135    fn rename(&mut self, name: PlSmallStr) {
136        self.0.physical_mut().rename(name);
137    }
138
139    fn chunk_lengths(&self) -> ChunkLenIter {
140        self.0.physical().chunk_lengths()
141    }
142    fn name(&self) -> &PlSmallStr {
143        self.0.physical().name()
144    }
145
146    fn chunks(&self) -> &Vec<ArrayRef> {
147        self.0.physical().chunks()
148    }
149    unsafe fn chunks_mut(&mut self) -> &mut Vec<ArrayRef> {
150        self.0.physical_mut().chunks_mut()
151    }
152    fn shrink_to_fit(&mut self) {
153        self.0.physical_mut().shrink_to_fit()
154    }
155
156    fn slice(&self, offset: i64, length: usize) -> Series {
157        self.with_state(false, |cats| cats.slice(offset, length))
158            .into_series()
159    }
160    fn split_at(&self, offset: i64) -> (Series, Series) {
161        let (a, b) = self.0.physical().split_at(offset);
162        let a = self.finish_with_state(false, a).into_series();
163        let b = self.finish_with_state(false, b).into_series();
164        (a, b)
165    }
166
167    fn append(&mut self, other: &Series) -> PolarsResult<()> {
168        polars_ensure!(self.0.dtype() == other.dtype(), append);
169        self.0.append(other.categorical().unwrap())
170    }
171    fn append_owned(&mut self, mut other: Series) -> PolarsResult<()> {
172        polars_ensure!(self.0.dtype() == other.dtype(), append);
173        let other = other
174            ._get_inner_mut()
175            .as_any_mut()
176            .downcast_mut::<CategoricalChunked>()
177            .unwrap();
178        self.0.append_owned(std::mem::take(other))
179    }
180
181    fn extend(&mut self, other: &Series) -> PolarsResult<()> {
182        polars_ensure!(self.0.dtype() == other.dtype(), extend);
183        let other_ca = other.categorical().unwrap();
184        // Fast path for globals of the same source
185        let rev_map_self = self.0.get_rev_map();
186        let rev_map_other = other_ca.get_rev_map();
187        match (&**rev_map_self, &**rev_map_other) {
188            (RevMapping::Global(_, _, idl), RevMapping::Global(_, _, idr)) if idl == idr => {
189                let mut rev_map_merger = GlobalRevMapMerger::new(rev_map_self.clone());
190                rev_map_merger.merge_map(rev_map_other)?;
191                self.0.physical_mut().extend(other_ca.physical())?;
192                // SAFETY: rev_maps are merged
193                unsafe { self.0.set_rev_map(rev_map_merger.finish(), false) };
194                Ok(())
195            },
196            _ => self.0.append(other_ca),
197        }
198    }
199
200    fn filter(&self, filter: &BooleanChunked) -> PolarsResult<Series> {
201        self.try_with_state(false, |cats| cats.filter(filter))
202            .map(|ca| ca.into_series())
203    }
204
205    fn take(&self, indices: &IdxCa) -> PolarsResult<Series> {
206        self.try_with_state(false, |cats| cats.take(indices))
207            .map(|ca| ca.into_series())
208    }
209
210    unsafe fn take_unchecked(&self, indices: &IdxCa) -> Series {
211        self.with_state(false, |cats| cats.take_unchecked(indices))
212            .into_series()
213    }
214
215    fn take_slice(&self, indices: &[IdxSize]) -> PolarsResult<Series> {
216        self.try_with_state(false, |cats| cats.take(indices))
217            .map(|ca| ca.into_series())
218    }
219
220    unsafe fn take_slice_unchecked(&self, indices: &[IdxSize]) -> Series {
221        self.with_state(false, |cats| cats.take_unchecked(indices))
222            .into_series()
223    }
224
225    fn len(&self) -> usize {
226        self.0.len()
227    }
228
229    fn rechunk(&self) -> Series {
230        self.with_state(true, |ca| ca.rechunk().into_owned())
231            .into_series()
232    }
233
234    fn new_from_index(&self, index: usize, length: usize) -> Series {
235        self.with_state(false, |cats| cats.new_from_index(index, length))
236            .into_series()
237    }
238
239    fn cast(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
240        self.0.cast_with_options(dtype, options)
241    }
242
243    #[inline]
244    unsafe fn get_unchecked(&self, index: usize) -> AnyValue {
245        self.0.get_any_value_unchecked(index)
246    }
247
248    fn sort_with(&self, options: SortOptions) -> PolarsResult<Series> {
249        Ok(self.0.sort_with(options).into_series())
250    }
251
252    fn arg_sort(&self, options: SortOptions) -> IdxCa {
253        self.0.arg_sort(options)
254    }
255
256    fn null_count(&self) -> usize {
257        self.0.physical().null_count()
258    }
259
260    fn has_nulls(&self) -> bool {
261        self.0.physical().has_nulls()
262    }
263
264    #[cfg(feature = "algorithm_group_by")]
265    fn unique(&self) -> PolarsResult<Series> {
266        self.0.unique().map(|ca| ca.into_series())
267    }
268
269    #[cfg(feature = "algorithm_group_by")]
270    fn n_unique(&self) -> PolarsResult<usize> {
271        self.0.n_unique()
272    }
273
274    #[cfg(feature = "algorithm_group_by")]
275    fn arg_unique(&self) -> PolarsResult<IdxCa> {
276        self.0.physical().arg_unique()
277    }
278
279    fn is_null(&self) -> BooleanChunked {
280        self.0.physical().is_null()
281    }
282
283    fn is_not_null(&self) -> BooleanChunked {
284        self.0.physical().is_not_null()
285    }
286
287    fn reverse(&self) -> Series {
288        self.with_state(true, |cats| cats.reverse()).into_series()
289    }
290
291    fn as_single_ptr(&mut self) -> PolarsResult<usize> {
292        self.0.physical_mut().as_single_ptr()
293    }
294
295    fn shift(&self, periods: i64) -> Series {
296        self.with_state(false, |ca| ca.shift(periods)).into_series()
297    }
298
299    fn clone_inner(&self) -> Arc<dyn SeriesTrait> {
300        Arc::new(SeriesWrap(Clone::clone(&self.0)))
301    }
302
303    fn min_reduce(&self) -> PolarsResult<Scalar> {
304        Ok(ChunkAggSeries::min_reduce(&self.0))
305    }
306
307    fn max_reduce(&self) -> PolarsResult<Scalar> {
308        Ok(ChunkAggSeries::max_reduce(&self.0))
309    }
310
311    fn as_any(&self) -> &dyn Any {
312        &self.0
313    }
314
315    fn as_any_mut(&mut self) -> &mut dyn Any {
316        &mut self.0
317    }
318
319    fn as_phys_any(&self) -> &dyn Any {
320        self.0.physical()
321    }
322
323    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
324        self as _
325    }
326}
327
328impl private::PrivateSeriesNumeric for SeriesWrap<CategoricalChunked> {
329    fn bit_repr(&self) -> Option<BitRepr> {
330        Some(BitRepr::Small(self.0.physical().clone()))
331    }
332}