polars_core/series/ops/
downcast.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2use crate::prelude::*;
3use crate::series::implementations::null::NullChunked;
4
5macro_rules! unpack_chunked_err {
6    ($series:expr => $name:expr) => {
7        polars_err!(SchemaMismatch: "invalid series dtype: expected `{}`, got `{}` for series with name `{}`", $name, $series.dtype(), $series.name())
8    };
9}
10
11macro_rules! try_unpack_chunked {
12    ($series:expr, $expected:pat $(if $guard: expr)? => $ca:ty) => {
13        match $series.dtype() {
14            $expected $(if $guard)? => {
15                // Check downcast in debug compiles
16                #[cfg(debug_assertions)]
17                {
18                    Some($series.as_ref().as_any().downcast_ref::<$ca>().unwrap())
19                }
20                #[cfg(not(debug_assertions))]
21                unsafe {
22                    Some(&*($series.as_ref() as *const dyn SeriesTrait as *const $ca))
23                }
24            },
25            _ => None,
26        }
27    };
28}
29
30impl Series {
31    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int8`]
32    pub fn try_i8(&self) -> Option<&Int8Chunked> {
33        try_unpack_chunked!(self, DataType::Int8 => Int8Chunked)
34    }
35
36    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int16`]
37    pub fn try_i16(&self) -> Option<&Int16Chunked> {
38        try_unpack_chunked!(self, DataType::Int16 => Int16Chunked)
39    }
40
41    /// Unpack to [`ChunkedArray`]
42    /// ```
43    /// # use polars_core::prelude::*;
44    /// let s = Series::new("foo".into(), [1i32 ,2, 3]);
45    /// let s_squared: Series = s.i32()
46    ///     .unwrap()
47    ///     .into_iter()
48    ///     .map(|opt_v| {
49    ///         match opt_v {
50    ///             Some(v) => Some(v * v),
51    ///             None => None, // null value
52    ///         }
53    /// }).collect();
54    /// ```
55    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int32`]
56    pub fn try_i32(&self) -> Option<&Int32Chunked> {
57        try_unpack_chunked!(self, DataType::Int32 => Int32Chunked)
58    }
59
60    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int64`]
61    pub fn try_i64(&self) -> Option<&Int64Chunked> {
62        try_unpack_chunked!(self, DataType::Int64 => Int64Chunked)
63    }
64
65    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int128`]
66    #[cfg(feature = "dtype-i128")]
67    pub fn try_i128(&self) -> Option<&Int128Chunked> {
68        try_unpack_chunked!(self, DataType::Int128 => Int128Chunked)
69    }
70
71    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Float16`]
72    #[cfg(feature = "dtype-f16")]
73    pub fn try_f16(&self) -> Option<&Float16Chunked> {
74        try_unpack_chunked!(self, DataType::Float16 => Float16Chunked)
75    }
76
77    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Float32`]
78    pub fn try_f32(&self) -> Option<&Float32Chunked> {
79        try_unpack_chunked!(self, DataType::Float32 => Float32Chunked)
80    }
81
82    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Float64`]
83    pub fn try_f64(&self) -> Option<&Float64Chunked> {
84        try_unpack_chunked!(self, DataType::Float64 => Float64Chunked)
85    }
86
87    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt8`]
88    pub fn try_u8(&self) -> Option<&UInt8Chunked> {
89        try_unpack_chunked!(self, DataType::UInt8 => UInt8Chunked)
90    }
91
92    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt16`]
93    pub fn try_u16(&self) -> Option<&UInt16Chunked> {
94        try_unpack_chunked!(self, DataType::UInt16 => UInt16Chunked)
95    }
96
97    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt32`]
98    pub fn try_u32(&self) -> Option<&UInt32Chunked> {
99        try_unpack_chunked!(self, DataType::UInt32 => UInt32Chunked)
100    }
101
102    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt64`]
103    pub fn try_u64(&self) -> Option<&UInt64Chunked> {
104        try_unpack_chunked!(self, DataType::UInt64 => UInt64Chunked)
105    }
106
107    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt128`]
108    #[cfg(feature = "dtype-u128")]
109    pub fn try_u128(&self) -> Option<&UInt128Chunked> {
110        try_unpack_chunked!(self, DataType::UInt128 => UInt128Chunked)
111    }
112
113    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Boolean`]
114    pub fn try_bool(&self) -> Option<&BooleanChunked> {
115        try_unpack_chunked!(self, DataType::Boolean => BooleanChunked)
116    }
117
118    /// Unpack to [`ChunkedArray`] of dtype [`DataType::String`]
119    pub fn try_str(&self) -> Option<&StringChunked> {
120        try_unpack_chunked!(self, DataType::String => StringChunked)
121    }
122
123    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Binary`]
124    pub fn try_binary(&self) -> Option<&BinaryChunked> {
125        try_unpack_chunked!(self, DataType::Binary => BinaryChunked)
126    }
127
128    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Binary`]
129    pub fn try_binary_offset(&self) -> Option<&BinaryOffsetChunked> {
130        try_unpack_chunked!(self, DataType::BinaryOffset => BinaryOffsetChunked)
131    }
132
133    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Time`]
134    #[cfg(feature = "dtype-time")]
135    pub fn try_time(&self) -> Option<&TimeChunked> {
136        try_unpack_chunked!(self, DataType::Time => TimeChunked)
137    }
138
139    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Date`]
140    #[cfg(feature = "dtype-date")]
141    pub fn try_date(&self) -> Option<&DateChunked> {
142        try_unpack_chunked!(self, DataType::Date => DateChunked)
143    }
144
145    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Datetime`]
146    #[cfg(feature = "dtype-datetime")]
147    pub fn try_datetime(&self) -> Option<&DatetimeChunked> {
148        try_unpack_chunked!(self, DataType::Datetime(_, _) => DatetimeChunked)
149    }
150
151    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Duration`]
152    #[cfg(feature = "dtype-duration")]
153    pub fn try_duration(&self) -> Option<&DurationChunked> {
154        try_unpack_chunked!(self, DataType::Duration(_) => DurationChunked)
155    }
156
157    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Decimal`]
158    #[cfg(feature = "dtype-decimal")]
159    pub fn try_decimal(&self) -> Option<&DecimalChunked> {
160        try_unpack_chunked!(self, DataType::Decimal(_, _) => DecimalChunked)
161    }
162
163    /// Unpack to [`ChunkedArray`] of dtype list
164    pub fn try_list(&self) -> Option<&ListChunked> {
165        try_unpack_chunked!(self, DataType::List(_) => ListChunked)
166    }
167
168    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Array`]
169    #[cfg(feature = "dtype-array")]
170    pub fn try_array(&self) -> Option<&ArrayChunked> {
171        try_unpack_chunked!(self, DataType::Array(_, _) => ArrayChunked)
172    }
173
174    #[cfg(feature = "dtype-categorical")]
175    pub fn try_cat<T: PolarsCategoricalType>(&self) -> Option<&CategoricalChunked<T>> {
176        try_unpack_chunked!(self, dt @ DataType::Enum(_, _) | dt @ DataType::Categorical(_, _) if dt.cat_physical().unwrap() == T::physical() => CategoricalChunked<T>)
177    }
178
179    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Categorical`] or [`DataType::Enum`] with a physical type of UInt8.
180    #[cfg(feature = "dtype-categorical")]
181    pub fn try_cat8(&self) -> Option<&Categorical8Chunked> {
182        self.try_cat::<Categorical8Type>()
183    }
184
185    #[cfg(feature = "dtype-categorical")]
186    pub fn try_cat16(&self) -> Option<&Categorical16Chunked> {
187        self.try_cat::<Categorical16Type>()
188    }
189
190    #[cfg(feature = "dtype-categorical")]
191    pub fn try_cat32(&self) -> Option<&Categorical32Chunked> {
192        self.try_cat::<Categorical32Type>()
193    }
194
195    /// Unpack to [`ExtensionChunked`] of dtype [`DataType::Extension`].
196    #[cfg(feature = "dtype-extension")]
197    pub fn try_ext(&self) -> Option<&ExtensionChunked> {
198        try_unpack_chunked!(self, DataType::Extension(_, _) => ExtensionChunked)
199    }
200
201    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Struct`]
202    #[cfg(feature = "dtype-struct")]
203    pub fn try_struct(&self) -> Option<&StructChunked> {
204        #[cfg(debug_assertions)]
205        {
206            if let DataType::Struct(_) = self.dtype() {
207                let any = self.as_any();
208                assert!(any.is::<StructChunked>());
209            }
210        }
211        try_unpack_chunked!(self, DataType::Struct(_) => StructChunked)
212    }
213
214    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Null`]
215    pub fn try_null(&self) -> Option<&NullChunked> {
216        try_unpack_chunked!(self, DataType::Null => NullChunked)
217    }
218    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int8`]
219    pub fn i8(&self) -> PolarsResult<&Int8Chunked> {
220        self.try_i8()
221            .ok_or_else(|| unpack_chunked_err!(self => "Int8"))
222    }
223
224    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int16`]
225    pub fn i16(&self) -> PolarsResult<&Int16Chunked> {
226        self.try_i16()
227            .ok_or_else(|| unpack_chunked_err!(self => "Int16"))
228    }
229
230    /// Unpack to [`ChunkedArray`]
231    /// ```
232    /// # use polars_core::prelude::*;
233    /// let s = Series::new("foo".into(), [1i32 ,2, 3]);
234    /// let s_squared: Series = s.i32()
235    ///     .unwrap()
236    ///     .into_iter()
237    ///     .map(|opt_v| {
238    ///         match opt_v {
239    ///             Some(v) => Some(v * v),
240    ///             None => None, // null value
241    ///         }
242    /// }).collect();
243    /// ```
244    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int32`]
245    pub fn i32(&self) -> PolarsResult<&Int32Chunked> {
246        self.try_i32()
247            .ok_or_else(|| unpack_chunked_err!(self => "Int32"))
248    }
249
250    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int64`]
251    pub fn i64(&self) -> PolarsResult<&Int64Chunked> {
252        self.try_i64()
253            .ok_or_else(|| unpack_chunked_err!(self => "Int64"))
254    }
255
256    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Int128`]
257    #[cfg(feature = "dtype-i128")]
258    pub fn i128(&self) -> PolarsResult<&Int128Chunked> {
259        self.try_i128()
260            .ok_or_else(|| unpack_chunked_err!(self => "Int128"))
261    }
262
263    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Float16`]
264    #[cfg(feature = "dtype-f16")]
265    pub fn f16(&self) -> PolarsResult<&Float16Chunked> {
266        self.try_f16()
267            .ok_or_else(|| unpack_chunked_err!(self => "Float16"))
268    }
269
270    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Float32`]
271    pub fn f32(&self) -> PolarsResult<&Float32Chunked> {
272        self.try_f32()
273            .ok_or_else(|| unpack_chunked_err!(self => "Float32"))
274    }
275
276    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Float64`]
277    pub fn f64(&self) -> PolarsResult<&Float64Chunked> {
278        self.try_f64()
279            .ok_or_else(|| unpack_chunked_err!(self => "Float64"))
280    }
281
282    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt8`]
283    pub fn u8(&self) -> PolarsResult<&UInt8Chunked> {
284        self.try_u8()
285            .ok_or_else(|| unpack_chunked_err!(self => "UInt8"))
286    }
287
288    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt16`]
289    pub fn u16(&self) -> PolarsResult<&UInt16Chunked> {
290        self.try_u16()
291            .ok_or_else(|| unpack_chunked_err!(self => "UInt16"))
292    }
293
294    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt32`]
295    pub fn u32(&self) -> PolarsResult<&UInt32Chunked> {
296        self.try_u32()
297            .ok_or_else(|| unpack_chunked_err!(self => "UInt32"))
298    }
299
300    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt64`]
301    pub fn u64(&self) -> PolarsResult<&UInt64Chunked> {
302        self.try_u64()
303            .ok_or_else(|| unpack_chunked_err!(self => "UInt64"))
304    }
305
306    /// Unpack to [`ChunkedArray`] of dtype [`DataType::UInt128`]
307    #[cfg(feature = "dtype-u128")]
308    pub fn u128(&self) -> PolarsResult<&UInt128Chunked> {
309        self.try_u128()
310            .ok_or_else(|| unpack_chunked_err!(self => "UInt128"))
311    }
312
313    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Boolean`]
314    pub fn bool(&self) -> PolarsResult<&BooleanChunked> {
315        self.try_bool()
316            .ok_or_else(|| unpack_chunked_err!(self => "Boolean"))
317    }
318
319    /// Unpack to [`ChunkedArray`] of dtype [`DataType::String`]
320    pub fn str(&self) -> PolarsResult<&StringChunked> {
321        self.try_str()
322            .ok_or_else(|| unpack_chunked_err!(self => "String"))
323    }
324
325    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Binary`]
326    pub fn binary(&self) -> PolarsResult<&BinaryChunked> {
327        self.try_binary()
328            .ok_or_else(|| unpack_chunked_err!(self => "Binary"))
329    }
330
331    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Binary`]
332    pub fn binary_offset(&self) -> PolarsResult<&BinaryOffsetChunked> {
333        self.try_binary_offset()
334            .ok_or_else(|| unpack_chunked_err!(self => "BinaryOffset"))
335    }
336
337    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Time`]
338    #[cfg(feature = "dtype-time")]
339    pub fn time(&self) -> PolarsResult<&TimeChunked> {
340        self.try_time()
341            .ok_or_else(|| unpack_chunked_err!(self => "Time"))
342    }
343
344    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Date`]
345    #[cfg(feature = "dtype-date")]
346    pub fn date(&self) -> PolarsResult<&DateChunked> {
347        self.try_date()
348            .ok_or_else(|| unpack_chunked_err!(self => "Date"))
349    }
350
351    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Datetime`]
352    #[cfg(feature = "dtype-datetime")]
353    pub fn datetime(&self) -> PolarsResult<&DatetimeChunked> {
354        self.try_datetime()
355            .ok_or_else(|| unpack_chunked_err!(self => "Datetime"))
356    }
357
358    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Duration`]
359    #[cfg(feature = "dtype-duration")]
360    pub fn duration(&self) -> PolarsResult<&DurationChunked> {
361        self.try_duration()
362            .ok_or_else(|| unpack_chunked_err!(self => "Duration"))
363    }
364
365    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Decimal`]
366    #[cfg(feature = "dtype-decimal")]
367    pub fn decimal(&self) -> PolarsResult<&DecimalChunked> {
368        self.try_decimal()
369            .ok_or_else(|| unpack_chunked_err!(self => "Decimal"))
370    }
371
372    /// Unpack to [`ChunkedArray`] of dtype list
373    pub fn list(&self) -> PolarsResult<&ListChunked> {
374        self.try_list()
375            .ok_or_else(|| unpack_chunked_err!(self => "List"))
376    }
377
378    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Array`]
379    #[cfg(feature = "dtype-array")]
380    pub fn array(&self) -> PolarsResult<&ArrayChunked> {
381        self.try_array()
382            .ok_or_else(|| unpack_chunked_err!(self => "Array"))
383    }
384
385    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Categorical`] or [`DataType::Enum`].
386    #[cfg(feature = "dtype-categorical")]
387    pub fn cat<T: PolarsCategoricalType>(&self) -> PolarsResult<&CategoricalChunked<T>> {
388        self.try_cat::<T>()
389            .ok_or_else(|| unpack_chunked_err!(self => "Enum | Categorical"))
390    }
391
392    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Categorical`] or [`DataType::Enum`] with a physical type of UInt8.
393    #[cfg(feature = "dtype-categorical")]
394    pub fn cat8(&self) -> PolarsResult<&CategoricalChunked<Categorical8Type>> {
395        self.try_cat8()
396            .ok_or_else(|| unpack_chunked_err!(self => "Enum8 | Categorical8"))
397    }
398
399    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Categorical`] or [`DataType::Enum`] with a physical type of UInt16.
400    #[cfg(feature = "dtype-categorical")]
401    pub fn cat16(&self) -> PolarsResult<&CategoricalChunked<Categorical16Type>> {
402        self.try_cat16()
403            .ok_or_else(|| unpack_chunked_err!(self => "Enum16 | Categorical16"))
404    }
405
406    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Categorical`] or [`DataType::Enum`] with a physical type of UInt32.
407    #[cfg(feature = "dtype-categorical")]
408    pub fn cat32(&self) -> PolarsResult<&CategoricalChunked<Categorical32Type>> {
409        self.try_cat32()
410            .ok_or_else(|| unpack_chunked_err!(self => "Enum32 | Categorical32"))
411    }
412
413    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Struct`]
414    #[cfg(feature = "dtype-struct")]
415    pub fn struct_(&self) -> PolarsResult<&StructChunked> {
416        #[cfg(debug_assertions)]
417        {
418            if let DataType::Struct(_) = self.dtype() {
419                let any = self.as_any();
420                assert!(any.is::<StructChunked>());
421            }
422        }
423
424        self.try_struct()
425            .ok_or_else(|| unpack_chunked_err!(self => "Struct"))
426    }
427
428    /// Unpack to [`ExtensionChunked`] of dtype [`DataType::Extension`].
429    #[cfg(feature = "dtype-extension")]
430    pub fn ext(&self) -> PolarsResult<&ExtensionChunked> {
431        self.try_ext()
432            .ok_or_else(|| unpack_chunked_err!(self => "Extension"))
433    }
434
435    /// Unpack to [`ChunkedArray`] of dtype [`DataType::Null`]
436    pub fn null(&self) -> PolarsResult<&NullChunked> {
437        self.try_null()
438            .ok_or_else(|| unpack_chunked_err!(self => "Null"))
439    }
440}