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

polars_core/chunked_array/
cast.rs
