polars_core/chunked_array/
cast.rs

1//! Implementations of the ChunkCast Trait.
2
3use std::borrow::Cow;
4
5use polars_compute::cast::CastOptionsImpl;
6#[cfg(feature = "serde-lazy")]
7use serde::{Deserialize, Serialize};
8
9use super::flags::StatisticsFlags;
10#[cfg(feature = "dtype-datetime")]
11use crate::prelude::DataType::Datetime;
12use crate::prelude::*;
13use crate::utils::{handle_array_casting_failures, handle_casting_failures};
14
15#[derive(Copy, Clone, Debug, Default, PartialEq, Hash, Eq)]
16#[cfg_attr(feature = "serde-lazy", derive(Serialize, Deserialize))]
17#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
18#[repr(u8)]
19pub enum CastOptions {
20    /// Raises on overflow
21    #[default]
22    Strict,
23    /// Overflow is replaced with null
24    NonStrict,
25    /// Allows wrapping overflow
26    Overflowing,
27}
28
29impl CastOptions {
30    pub fn is_strict(&self) -> bool {
31        matches!(self, CastOptions::Strict)
32    }
33}
34
35impl From<CastOptions> for CastOptionsImpl {
36    fn from(value: CastOptions) -> Self {
37        let wrapped = match value {
38            CastOptions::Strict | CastOptions::NonStrict => false,
39            CastOptions::Overflowing => true,
40        };
41        CastOptionsImpl {
42            wrapped,
43            partial: false,
44        }
45    }
46}
47
48pub(crate) fn cast_chunks(
49    chunks: &[ArrayRef],
50    dtype: &DataType,
51    options: CastOptions,
52) -> PolarsResult<Vec<ArrayRef>> {
53    let check_nulls = matches!(options, CastOptions::Strict);
54    let options = options.into();
55
56    let arrow_dtype = dtype.try_to_arrow(CompatLevel::newest())?;
57    chunks
58        .iter()
59        .map(|arr| {
60            let out = polars_compute::cast::cast(arr.as_ref(), &arrow_dtype, options);
61            if check_nulls {
62                out.and_then(|new| {
63                    if arr.null_count() != new.null_count() {
64                        handle_array_casting_failures(&**arr, &*new)?;
65                    }
66                    Ok(new)
67                })
68            } else {
69                out
70            }
71        })
72        .collect::<PolarsResult<Vec<_>>>()
73}
74
75fn cast_impl_inner(
76    name: PlSmallStr,
77    chunks: &[ArrayRef],
78    dtype: &DataType,
79    options: CastOptions,
80) -> PolarsResult<Series> {
81    let chunks = match dtype {
82        #[cfg(feature = "dtype-decimal")]
83        DataType::Decimal(_, _) => {
84            let mut chunks = cast_chunks(chunks, dtype, options)?;
85            // @NOTE: We cannot cast here as that will lower the scale.
86            for chunk in chunks.iter_mut() {
87                *chunk = std::mem::take(
88                    chunk
89                        .as_any_mut()
90                        .downcast_mut::<PrimitiveArray<i128>>()
91                        .unwrap(),
92                )
93                .to(ArrowDataType::Int128)
94                .to_boxed();
95            }
96            chunks
97        },
98        _ => cast_chunks(chunks, &dtype.to_physical(), options)?,
99    };
100
101    let out = Series::try_from((name, chunks))?;
102    use DataType::*;
103    let out = match dtype {
104        Date => out.into_date(),
105        Datetime(tu, tz) => match tz {
106            #[cfg(feature = "timezones")]
107            Some(tz) => {
108                TimeZone::validate_time_zone(tz)?;
109                out.into_datetime(*tu, Some(tz.clone()))
110            },
111            _ => out.into_datetime(*tu, None),
112        },
113        Duration(tu) => out.into_duration(*tu),
114        #[cfg(feature = "dtype-time")]
115        Time => out.into_time(),
116        #[cfg(feature = "dtype-decimal")]
117        Decimal(precision, scale) => out.into_decimal(*precision, *scale)?,
118        #[cfg(feature = "dtype-extension")]
119        Extension(typ, _) => out.into_extension(typ.clone()),
120        _ => out,
121    };
122
123    Ok(out)
124}
125
126fn cast_impl(
127    name: PlSmallStr,
128    chunks: &[ArrayRef],
129    dtype: &DataType,
130    options: CastOptions,
131) -> PolarsResult<Series> {
132    cast_impl_inner(name, chunks, dtype, options)
133}
134
135#[cfg(feature = "dtype-struct")]
136fn cast_single_to_struct(
137    name: PlSmallStr,
138    chunks: &[ArrayRef],
139    fields: &[Field],
140    options: CastOptions,
141) -> PolarsResult<Series> {
142    polars_ensure!(fields.len() == 1, InvalidOperation: "must specify one field in the struct");
143    let mut new_fields = Vec::with_capacity(fields.len());
144    // cast to first field dtype
145    let mut fields = fields.iter();
146    let fld = fields.next().unwrap();
147    let s = cast_impl_inner(fld.name.clone(), chunks, &fld.dtype, options)?;
148    let length = s.len();
149    new_fields.push(s);
150
151    for fld in fields {
152        new_fields.push(Series::full_null(fld.name.clone(), length, &fld.dtype));
153    }
154
155    StructChunked::from_series(name, length, new_fields.iter()).map(|ca| ca.into_series())
156}
157
158impl<T> ChunkedArray<T>
159where
160    T: PolarsNumericType,
161{
162    fn cast_impl(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
163        if self.dtype() == dtype {
164            // SAFETY: chunks are correct dtype
165            let mut out = unsafe {
166                Series::from_chunks_and_dtype_unchecked(
167                    self.name().clone(),
168                    self.chunks.clone(),
169                    dtype,
170                )
171            };
172            out.set_sorted_flag(self.is_sorted_flag());
173            return Ok(out);
174        }
175        match dtype {
176            #[cfg(feature = "dtype-categorical")]
177            DataType::Categorical(..) | DataType::Enum(..) => {
178                polars_bail!(
179                    ComputeError:
180                    "casting from {} to {dtype} is not supported.\n\
181                    Instead of `.cast({dtype:?}`, use `.cat.to({dtype:?})`.",
182                    T::get_static_dtype()
183                );
184            },
185
186            #[cfg(feature = "dtype-struct")]
187            DataType::Struct(fields) => {
188                cast_single_to_struct(self.name().clone(), &self.chunks, fields, options)
189            },
190            _ => cast_impl_inner(self.name().clone(), &self.chunks, dtype, options).map(|mut s| {
191                // maintain sorted if data types
192                // - remain signed
193                // - unsigned -> signed
194                // this may still fail with overflow?
195                let to_signed = dtype.is_signed_integer();
196                let unsigned2unsigned =
197                    self.dtype().is_unsigned_integer() && dtype.is_unsigned_integer();
198                let allowed = to_signed || unsigned2unsigned;
199
200                if (allowed)
201                    && (s.null_count() == self.null_count())
202                    // physical to logicals
203                    || (self.dtype().to_physical() == dtype.to_physical())
204                {
205                    let is_sorted = self.is_sorted_flag();
206                    s.set_sorted_flag(is_sorted)
207                }
208                s
209            }),
210        }
211    }
212}
213
214impl<T> ChunkCast for ChunkedArray<T>
215where
216    T: PolarsNumericType,
217{
218    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
219        self.cast_impl(dtype, options)
220    }
221
222    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
223        match dtype {
224            // LEGACY
225            // TODO @ cat-rework: remove after exposing to/from physical functions.
226            #[cfg(feature = "dtype-categorical")]
227            DataType::Categorical(cats, _mapping) => {
228                polars_ensure!(self.dtype() == &cats.physical().dtype(), ComputeError: "cannot cast numeric types to 'Categorical'");
229                with_match_categorical_physical_type!(cats.physical(), |$C| {
230                    // SAFETY: we are guarded by the type system.
231                    type PhysCa = ChunkedArray<<$C as PolarsCategoricalType>::PolarsPhysical>;
232                    let ca = unsafe { &*(self as *const ChunkedArray<T> as *const PhysCa) };
233                    Ok(CategoricalChunked::<$C>::from_cats_and_dtype_unchecked(ca.clone(), dtype.clone())
234                        .into_series())
235                })
236            },
237
238            // LEGACY
239            // TODO @ cat-rework: remove after exposing to/from physical functions.
240            #[cfg(feature = "dtype-categorical")]
241            DataType::Enum(fcats, _mapping) => {
242                polars_ensure!(self.dtype() == &fcats.physical().dtype(), ComputeError: "cannot cast numeric types to 'Enum'");
243                with_match_categorical_physical_type!(fcats.physical(), |$C| {
244                    // SAFETY: we are guarded by the type system.
245                    type PhysCa = ChunkedArray<<$C as PolarsCategoricalType>::PolarsPhysical>;
246                    let ca = unsafe { &*(self as *const ChunkedArray<T> as *const PhysCa) };
247                    Ok(CategoricalChunked::<$C>::from_cats_and_dtype_unchecked(ca.clone(), dtype.clone()).into_series())
248                })
249            },
250
251            _ => self.cast_impl(dtype, CastOptions::Overflowing),
252        }
253    }
254}
255
256impl ChunkCast for StringChunked {
257    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
258        match dtype {
259            #[cfg(feature = "dtype-categorical")]
260            DataType::Categorical(cats, _mapping) => {
261                with_match_categorical_physical_type!(cats.physical(), |$C| {
262                    Ok(CategoricalChunked::<$C>::from_str_iter(self.name().clone(), dtype.clone(), self.iter())?
263                        .into_series())
264                })
265            },
266            #[cfg(feature = "dtype-categorical")]
267            DataType::Enum(fcats, _mapping) => {
268                let ret = with_match_categorical_physical_type!(fcats.physical(), |$C| {
269                    CategoricalChunked::<$C>::from_str_iter(self.name().clone(), dtype.clone(), self.iter())?
270                        .into_series()
271                });
272
273                if options.is_strict() && self.null_count() != ret.null_count() {
274                    handle_casting_failures(&self.clone().into_series(), &ret)?;
275                }
276
277                Ok(ret)
278            },
279            #[cfg(feature = "dtype-struct")]
280            DataType::Struct(fields) => {
281                cast_single_to_struct(self.name().clone(), &self.chunks, fields, options)
282            },
283            #[cfg(feature = "dtype-decimal")]
284            DataType::Decimal(precision, scale) => {
285                let chunks = self.downcast_iter().map(|arr| {
286                    polars_compute::cast::binview_to_decimal(&arr.to_binview(), *precision, *scale)
287                        .to(ArrowDataType::Int128)
288                });
289                let ca = Int128Chunked::from_chunk_iter(self.name().clone(), chunks);
290                Ok(ca.into_decimal_unchecked(*precision, *scale).into_series())
291            },
292            #[cfg(feature = "dtype-date")]
293            DataType::Date => {
294                let result = cast_chunks(&self.chunks, dtype, options)?;
295                let out = Series::try_from((self.name().clone(), result))?;
296                Ok(out)
297            },
298            #[cfg(feature = "dtype-datetime")]
299            DataType::Datetime(time_unit, time_zone) => match time_zone {
300                #[cfg(feature = "timezones")]
301                Some(time_zone) => {
302                    TimeZone::validate_time_zone(time_zone)?;
303                    let result = cast_chunks(
304                        &self.chunks,
305                        &Datetime(time_unit.to_owned(), Some(time_zone.clone())),
306                        options,
307                    )?;
308                    Series::try_from((self.name().clone(), result))
309                },
310                _ => {
311                    let result =
312                        cast_chunks(&self.chunks, &Datetime(time_unit.to_owned(), None), options)?;
313                    Series::try_from((self.name().clone(), result))
314                },
315            },
316            _ => cast_impl(self.name().clone(), &self.chunks, dtype, options),
317        }
318    }
319
320    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
321        self.cast_with_options(dtype, CastOptions::Overflowing)
322    }
323}
324
325impl BinaryChunked {
326    /// # Safety
327    /// String is not validated
328    pub unsafe fn to_string_unchecked(&self) -> StringChunked {
329        let chunks = self
330            .downcast_iter()
331            .map(|arr| unsafe { arr.to_utf8view_unchecked() }.boxed())
332            .collect();
333        let field = Arc::new(Field::new(self.name().clone(), DataType::String));
334
335        let mut ca = StringChunked::new_with_compute_len(field, chunks);
336
337        use StatisticsFlags as F;
338        ca.retain_flags_from(self, F::IS_SORTED_ANY | F::CAN_FAST_EXPLODE_LIST);
339        ca
340    }
341}
342
343impl StringChunked {
344    pub fn as_binary(&self) -> BinaryChunked {
345        let chunks = self
346            .downcast_iter()
347            .map(|arr| arr.to_binview().boxed())
348            .collect();
349        let field = Arc::new(Field::new(self.name().clone(), DataType::Binary));
350
351        let mut ca = BinaryChunked::new_with_compute_len(field, chunks);
352
353        use StatisticsFlags as F;
354        ca.retain_flags_from(self, F::IS_SORTED_ANY | F::CAN_FAST_EXPLODE_LIST);
355        ca
356    }
357}
358
359impl ChunkCast for BinaryChunked {
360    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
361        match dtype {
362            #[cfg(feature = "dtype-struct")]
363            DataType::Struct(fields) => {
364                cast_single_to_struct(self.name().clone(), &self.chunks, fields, options)
365            },
366            _ => cast_impl(self.name().clone(), &self.chunks, dtype, options),
367        }
368    }
369
370    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
371        match dtype {
372            DataType::String => unsafe { Ok(self.to_string_unchecked().into_series()) },
373            _ => self.cast_with_options(dtype, CastOptions::Overflowing),
374        }
375    }
376}
377
378impl ChunkCast for BinaryOffsetChunked {
379    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
380        match dtype {
381            #[cfg(feature = "dtype-struct")]
382            DataType::Struct(fields) => {
383                cast_single_to_struct(self.name().clone(), &self.chunks, fields, options)
384            },
385            _ => cast_impl(self.name().clone(), &self.chunks, dtype, options),
386        }
387    }
388
389    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
390        self.cast_with_options(dtype, CastOptions::Overflowing)
391    }
392}
393
394impl ChunkCast for BooleanChunked {
395    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
396        match dtype {
397            #[cfg(feature = "dtype-struct")]
398            DataType::Struct(fields) => {
399                cast_single_to_struct(self.name().clone(), &self.chunks, fields, options)
400            },
401            #[cfg(feature = "dtype-categorical")]
402            DataType::Categorical(_, _) | DataType::Enum(_, _) => {
403                polars_bail!(InvalidOperation: "cannot cast Boolean to Categorical");
404            },
405            _ => cast_impl(self.name().clone(), &self.chunks, dtype, options),
406        }
407    }
408
409    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
410        self.cast_with_options(dtype, CastOptions::Overflowing)
411    }
412}
413
414impl ChunkCast for ListChunked {
415    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
416        let ca = self
417            .trim_lists_to_normalized_offsets()
418            .map_or(Cow::Borrowed(self), Cow::Owned);
419        let ca = ca.propagate_nulls().map_or(ca, Cow::Owned);
420
421        use DataType::*;
422        match dtype {
423            List(child_type) => {
424                match (ca.inner_dtype(), &**child_type) {
425                    (old, new) if old == new => Ok(ca.into_owned().into_series()),
426                    // TODO @ cat-rework: can we implement this now?
427                    #[cfg(feature = "dtype-categorical")]
428                    (dt, Categorical(_, _) | Enum(_, _))
429                        if !matches!(dt, Categorical(_, _) | Enum(_, _) | String | Null) =>
430                    {
431                        polars_bail!(InvalidOperation: "cannot cast List inner type: '{:?}' to Categorical", dt)
432                    },
433                    _ => {
434                        // ensure the inner logical type bubbles up
435                        let (arr, child_type) = cast_list(ca.as_ref(), child_type, options)?;
436                        // SAFETY: we just cast so the dtype matches.
437                        // we must take this path to correct for physical types.
438                        unsafe {
439                            Ok(Series::from_chunks_and_dtype_unchecked(
440                                ca.name().clone(),
441                                vec![arr],
442                                &List(Box::new(child_type)),
443                            ))
444                        }
445                    },
446                }
447            },
448            #[cfg(feature = "dtype-array")]
449            Array(child_type, width) => {
450                let physical_type = dtype.to_physical();
451
452                // cast to the physical type to avoid logical chunks.
453                let chunks = cast_chunks(ca.chunks(), &physical_type, options)?;
454                // SAFETY: we just cast so the dtype matches.
455                // we must take this path to correct for physical types.
456                unsafe {
457                    Ok(Series::from_chunks_and_dtype_unchecked(
458                        ca.name().clone(),
459                        chunks,
460                        &Array(child_type.clone(), *width),
461                    ))
462                }
463            },
464            #[cfg(feature = "dtype-u8")]
465            Binary => {
466                polars_ensure!(
467                    matches!(self.inner_dtype(), UInt8),
468                    InvalidOperation: "cannot cast List type (inner: '{:?}', to: '{:?}')",
469                    self.inner_dtype(),
470                    dtype,
471                );
472                let chunks = cast_chunks(self.chunks(), &DataType::Binary, options)?;
473
474                // SAFETY: we just cast so the dtype matches.
475                unsafe {
476                    Ok(Series::from_chunks_and_dtype_unchecked(
477                        self.name().clone(),
478                        chunks,
479                        &DataType::Binary,
480                    ))
481                }
482            },
483            #[cfg(feature = "dtype-map")]
484            Map(to_key, to_value) => {
485                let storage = if ca.inner_dtype().is_nested_null() {
486                    // Every row is empty, so there are no entry children to transform.
487                    ca.cast_with_options(&dtype.map_storage_dtype().unwrap(), options)?
488                } else {
489                    try_apply_map_entries(ca.as_ref(), |key, value| {
490                        Ok((
491                            key.cast_with_options(to_key, options)?,
492                            value.cast_with_options(to_value, options)?,
493                        ))
494                    })?
495                    .into_series()
496                };
497
498                Ok(MapChunked::try_from_storage(dtype.clone(), storage)?.into_series())
499            },
500            _ => {
501                polars_bail!(
502                    InvalidOperation: "cannot cast List type (inner: '{:?}', to: '{:?}')",
503                    ca.inner_dtype(),
504                    dtype,
505                )
506            },
507        }
508    }
509
510    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
511        use DataType::*;
512        match dtype {
513            List(child_type) => cast_list_unchecked(self, child_type),
514            _ => self.cast_with_options(dtype, CastOptions::Overflowing),
515        }
516    }
517}
518
519/// We cannot cast anything to or from List/LargeList
520/// So this implementation casts the inner type
521#[cfg(feature = "dtype-array")]
522impl ChunkCast for ArrayChunked {
523    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
524        let ca = self
525            .trim_lists_to_normalized_offsets()
526            .map_or(Cow::Borrowed(self), Cow::Owned);
527        let ca = ca.propagate_nulls().map_or(ca, Cow::Owned);
528
529        use DataType::*;
530        match dtype {
531            Array(child_type, width) => {
532                polars_ensure!(
533                    *width == ca.width(),
534                    InvalidOperation: "cannot cast Array to a different width"
535                );
536
537                match (ca.inner_dtype(), &**child_type) {
538                    (old, new) if old == new => Ok(ca.into_owned().into_series()),
539                    // TODO @ cat-rework: can we implement this now?
540                    #[cfg(feature = "dtype-categorical")]
541                    (dt, Categorical(_, _) | Enum(_, _)) if !matches!(dt, String) => {
542                        polars_bail!(InvalidOperation: "cannot cast Array inner type: '{:?}' to dtype: {:?}", dt, child_type)
543                    },
544                    _ => {
545                        // ensure the inner logical type bubbles up
546                        let (arr, child_type) =
547                            cast_fixed_size_list(ca.as_ref(), child_type, options)?;
548                        // SAFETY: we just cast so the dtype matches.
549                        // we must take this path to correct for physical types.
550                        unsafe {
551                            Ok(Series::from_chunks_and_dtype_unchecked(
552                                ca.name().clone(),
553                                vec![arr],
554                                &Array(Box::new(child_type), *width),
555                            ))
556                        }
557                    },
558                }
559            },
560            List(child_type) => {
561                let physical_type = dtype.to_physical();
562                // cast to the physical type to avoid logical chunks.
563                let chunks = cast_chunks(ca.chunks(), &physical_type, options)?;
564                // SAFETY: we just cast so the dtype matches.
565                // we must take this path to correct for physical types.
566                unsafe {
567                    Ok(Series::from_chunks_and_dtype_unchecked(
568                        ca.name().clone(),
569                        chunks,
570                        &List(child_type.clone()),
571                    ))
572                }
573            },
574            _ => {
575                polars_bail!(
576                    InvalidOperation: "cannot cast Array type (inner: '{:?}', to: '{:?}')",
577                    ca.inner_dtype(),
578                    dtype,
579                )
580            },
581        }
582    }
583
584    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
585        self.cast_with_options(dtype, CastOptions::Overflowing)
586    }
587}
588
589// Returns inner data type. This is needed because a cast can instantiate the dtype inner
590// values for instance with categoricals
591fn cast_list(
592    ca: &ListChunked,
593    child_type: &DataType,
594    options: CastOptions,
595) -> PolarsResult<(ArrayRef, DataType)> {
596    // We still rechunk because we must bubble up a single data-type
597    // TODO!: consider a version that works on chunks and merges the data-types and arrays.
598    let ca = ca.rechunk();
599    let arr = ca.downcast_as_array();
600    // SAFETY: inner dtype is passed correctly
601    let s = unsafe {
602        Series::from_chunks_and_dtype_unchecked(
603            PlSmallStr::EMPTY,
604            vec![arr.values().clone()],
605            ca.inner_dtype(),
606        )
607    };
608    let new_inner = s.cast_with_options(child_type, options)?;
609
610    let inner_dtype = new_inner.dtype().clone();
611    debug_assert_eq!(&inner_dtype, child_type);
612
613    let new_values = new_inner.array_ref(0).clone();
614
615    let dtype = ListArray::<i64>::default_datatype(new_values.dtype().clone());
616    let new_arr = ListArray::<i64>::new(
617        dtype,
618        arr.offsets().clone(),
619        new_values,
620        arr.validity().cloned(),
621    );
622    Ok((new_arr.boxed(), inner_dtype))
623}
624
625unsafe fn cast_list_unchecked(ca: &ListChunked, child_type: &DataType) -> PolarsResult<Series> {
626    // TODO! add chunked, but this must correct for list offsets.
627    let ca = ca.rechunk();
628    let arr = ca.downcast_as_array();
629    // SAFETY: inner dtype is passed correctly
630    let s = unsafe {
631        Series::from_chunks_and_dtype_unchecked(
632            PlSmallStr::EMPTY,
633            vec![arr.values().clone()],
634            ca.inner_dtype(),
635        )
636    };
637    let new_inner = s.cast_unchecked(child_type)?;
638    let new_values = new_inner.array_ref(0).clone();
639
640    let dtype = ListArray::<i64>::default_datatype(new_values.dtype().clone());
641    let new_arr = ListArray::<i64>::new(
642        dtype,
643        arr.offsets().clone(),
644        new_values,
645        arr.validity().cloned(),
646    );
647    Ok(ListChunked::from_chunks_and_dtype_unchecked(
648        ca.name().clone(),
649        vec![Box::new(new_arr)],
650        DataType::List(Box::new(child_type.clone())),
651    )
652    .into_series())
653}
654
655// Returns inner data type. This is needed because a cast can instantiate the dtype inner
656// values for instance with categoricals
657#[cfg(feature = "dtype-array")]
658fn cast_fixed_size_list(
659    ca: &ArrayChunked,
660    child_type: &DataType,
661    options: CastOptions,
662) -> PolarsResult<(ArrayRef, DataType)> {
663    let ca = ca.rechunk();
664    let arr = ca.downcast_as_array();
665    // SAFETY: inner dtype is passed correctly
666    let s = unsafe {
667        Series::from_chunks_and_dtype_unchecked(
668            PlSmallStr::EMPTY,
669            vec![arr.values().clone()],
670            ca.inner_dtype(),
671        )
672    };
673    let new_inner = s.cast_with_options(child_type, options)?;
674
675    let inner_dtype = new_inner.dtype().clone();
676    debug_assert_eq!(&inner_dtype, child_type);
677
678    let new_values = new_inner.array_ref(0).clone();
679
680    let dtype = FixedSizeListArray::default_datatype(new_values.dtype().clone(), ca.width());
681    let new_arr = FixedSizeListArray::new(dtype, ca.len(), new_values, arr.validity().cloned());
682    Ok((Box::new(new_arr), inner_dtype))
683}
684
685#[cfg(test)]
686mod test {
687    use crate::chunked_array::cast::CastOptions;
688    use crate::prelude::*;
689
690    #[test]
691    fn test_cast_list() -> PolarsResult<()> {
692        let mut builder = ListPrimitiveChunkedBuilder::<Int32Type>::new(
693            PlSmallStr::from_static("a"),
694            10,
695            10,
696            DataType::Int32,
697        );
698        builder.append_opt_slice(Some(&[1i32, 2, 3]));
699        builder.append_opt_slice(Some(&[1i32, 2, 3]));
700        let ca = builder.finish();
701
702        let new = ca.cast_with_options(
703            &DataType::List(DataType::Float64.into()),
704            CastOptions::Strict,
705        )?;
706
707        assert_eq!(new.dtype(), &DataType::List(DataType::Float64.into()));
708        Ok(())
709    }
710
711    #[test]
712    #[cfg(feature = "dtype-categorical")]
713    fn test_cast_noop() {
714        // check if we can cast categorical twice without panic
715        let ca = StringChunked::new(PlSmallStr::from_static("foo"), &["bar", "ham"]);
716        let cats = Categories::global();
717        let out = ca
718            .cast_with_options(
719                &DataType::from_categories(cats.clone()),
720                CastOptions::Strict,
721            )
722            .unwrap();
723        let out = out.cast(&DataType::from_categories(cats)).unwrap();
724        assert!(matches!(out.dtype(), &DataType::Categorical(_, _)))
725    }
726}
polars_core/chunked_array/cast.rs

polars_core/chunked_array/
cast.rs
