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                // TODO @ cat-rework: can we implement this now?
469                // TODO!: properly implement this recursively.
470                #[cfg(feature = "dtype-categorical")]
471                polars_ensure!(!matches!(&**child_type, Categorical(_, _)), InvalidOperation: "array of categorical is not yet supported");
472
473                // cast to the physical type to avoid logical chunks.
474                let chunks = cast_chunks(ca.chunks(), &physical_type, options)?;
475                // SAFETY: we just cast so the dtype matches.
476                // we must take this path to correct for physical types.
477                unsafe {
478                    Ok(Series::from_chunks_and_dtype_unchecked(
479                        ca.name().clone(),
480                        chunks,
481                        &Array(child_type.clone(), *width),
482                    ))
483                }
484            },
485            #[cfg(feature = "dtype-u8")]
486            Binary => {
487                polars_ensure!(
488                    matches!(self.inner_dtype(), UInt8),
489                    InvalidOperation: "cannot cast List type (inner: '{:?}', to: '{:?}')",
490                    self.inner_dtype(),
491                    dtype,
492                );
493                let chunks = cast_chunks(self.chunks(), &DataType::Binary, options)?;
494
495                // SAFETY: we just cast so the dtype matches.
496                unsafe {
497                    Ok(Series::from_chunks_and_dtype_unchecked(
498                        self.name().clone(),
499                        chunks,
500                        &DataType::Binary,
501                    ))
502                }
503            },
504            _ => {
505                polars_bail!(
506                    InvalidOperation: "cannot cast List type (inner: '{:?}', to: '{:?}')",
507                    ca.inner_dtype(),
508                    dtype,
509                )
510            },
511        }
512    }
513
514    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
515        use DataType::*;
516        match dtype {
517            List(child_type) => cast_list_unchecked(self, child_type),
518            _ => self.cast_with_options(dtype, CastOptions::Overflowing),
519        }
520    }
521}
522
523/// We cannot cast anything to or from List/LargeList
524/// So this implementation casts the inner type
525#[cfg(feature = "dtype-array")]
526impl ChunkCast for ArrayChunked {
527    fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Series> {
528        let ca = self
529            .trim_lists_to_normalized_offsets()
530            .map_or(Cow::Borrowed(self), Cow::Owned);
531        let ca = ca.propagate_nulls().map_or(ca, Cow::Owned);
532
533        use DataType::*;
534        match dtype {
535            Array(child_type, width) => {
536                polars_ensure!(
537                    *width == ca.width(),
538                    InvalidOperation: "cannot cast Array to a different width"
539                );
540
541                match (ca.inner_dtype(), &**child_type) {
542                    (old, new) if old == new => Ok(ca.into_owned().into_series()),
543                    // TODO @ cat-rework: can we implement this now?
544                    #[cfg(feature = "dtype-categorical")]
545                    (dt, Categorical(_, _) | Enum(_, _)) if !matches!(dt, String) => {
546                        polars_bail!(InvalidOperation: "cannot cast Array inner type: '{:?}' to dtype: {:?}", dt, child_type)
547                    },
548                    _ => {
549                        // ensure the inner logical type bubbles up
550                        let (arr, child_type) =
551                            cast_fixed_size_list(ca.as_ref(), child_type, options)?;
552                        // SAFETY: we just cast so the dtype matches.
553                        // we must take this path to correct for physical types.
554                        unsafe {
555                            Ok(Series::from_chunks_and_dtype_unchecked(
556                                ca.name().clone(),
557                                vec![arr],
558                                &Array(Box::new(child_type), *width),
559                            ))
560                        }
561                    },
562                }
563            },
564            List(child_type) => {
565                let physical_type = dtype.to_physical();
566                // cast to the physical type to avoid logical chunks.
567                let chunks = cast_chunks(ca.chunks(), &physical_type, options)?;
568                // SAFETY: we just cast so the dtype matches.
569                // we must take this path to correct for physical types.
570                unsafe {
571                    Ok(Series::from_chunks_and_dtype_unchecked(
572                        ca.name().clone(),
573                        chunks,
574                        &List(child_type.clone()),
575                    ))
576                }
577            },
578            _ => {
579                polars_bail!(
580                    InvalidOperation: "cannot cast Array type (inner: '{:?}', to: '{:?}')",
581                    ca.inner_dtype(),
582                    dtype,
583                )
584            },
585        }
586    }
587
588    unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Series> {
589        self.cast_with_options(dtype, CastOptions::Overflowing)
590    }
591}
592
593// Returns inner data type. This is needed because a cast can instantiate the dtype inner
594// values for instance with categoricals
595fn cast_list(
596    ca: &ListChunked,
597    child_type: &DataType,
598    options: CastOptions,
599) -> PolarsResult<(ArrayRef, DataType)> {
600    // We still rechunk because we must bubble up a single data-type
601    // TODO!: consider a version that works on chunks and merges the data-types and arrays.
602    let ca = ca.rechunk();
603    let arr = ca.downcast_as_array();
604    // SAFETY: inner dtype is passed correctly
605    let s = unsafe {
606        Series::from_chunks_and_dtype_unchecked(
607            PlSmallStr::EMPTY,
608            vec![arr.values().clone()],
609            ca.inner_dtype(),
610        )
611    };
612    let new_inner = s.cast_with_options(child_type, options)?;
613
614    let inner_dtype = new_inner.dtype().clone();
615    debug_assert_eq!(&inner_dtype, child_type);
616
617    let new_values = new_inner.array_ref(0).clone();
618
619    let dtype = ListArray::<i64>::default_datatype(new_values.dtype().clone());
620    let new_arr = ListArray::<i64>::new(
621        dtype,
622        arr.offsets().clone(),
623        new_values,
624        arr.validity().cloned(),
625    );
626    Ok((new_arr.boxed(), inner_dtype))
627}
628
629unsafe fn cast_list_unchecked(ca: &ListChunked, child_type: &DataType) -> PolarsResult<Series> {
630    // TODO! add chunked, but this must correct for list offsets.
631    let ca = ca.rechunk();
632    let arr = ca.downcast_as_array();
633    // SAFETY: inner dtype is passed correctly
634    let s = unsafe {
635        Series::from_chunks_and_dtype_unchecked(
636            PlSmallStr::EMPTY,
637            vec![arr.values().clone()],
638            ca.inner_dtype(),
639        )
640    };
641    let new_inner = s.cast_unchecked(child_type)?;
642    let new_values = new_inner.array_ref(0).clone();
643
644    let dtype = ListArray::<i64>::default_datatype(new_values.dtype().clone());
645    let new_arr = ListArray::<i64>::new(
646        dtype,
647        arr.offsets().clone(),
648        new_values,
649        arr.validity().cloned(),
650    );
651    Ok(ListChunked::from_chunks_and_dtype_unchecked(
652        ca.name().clone(),
653        vec![Box::new(new_arr)],
654        DataType::List(Box::new(child_type.clone())),
655    )
656    .into_series())
657}
658
659// Returns inner data type. This is needed because a cast can instantiate the dtype inner
660// values for instance with categoricals
661#[cfg(feature = "dtype-array")]
662fn cast_fixed_size_list(
663    ca: &ArrayChunked,
664    child_type: &DataType,
665    options: CastOptions,
666) -> PolarsResult<(ArrayRef, DataType)> {
667    let ca = ca.rechunk();
668    let arr = ca.downcast_as_array();
669    // SAFETY: inner dtype is passed correctly
670    let s = unsafe {
671        Series::from_chunks_and_dtype_unchecked(
672            PlSmallStr::EMPTY,
673            vec![arr.values().clone()],
674            ca.inner_dtype(),
675        )
676    };
677    let new_inner = s.cast_with_options(child_type, options)?;
678
679    let inner_dtype = new_inner.dtype().clone();
680    debug_assert_eq!(&inner_dtype, child_type);
681
682    let new_values = new_inner.array_ref(0).clone();
683
684    let dtype = FixedSizeListArray::default_datatype(new_values.dtype().clone(), ca.width());
685    let new_arr = FixedSizeListArray::new(dtype, ca.len(), new_values, arr.validity().cloned());
686    Ok((Box::new(new_arr), inner_dtype))
687}
688
689#[cfg(test)]
690mod test {
691    use crate::chunked_array::cast::CastOptions;
692    use crate::prelude::*;
693
694    #[test]
695    fn test_cast_list() -> PolarsResult<()> {
696        let mut builder = ListPrimitiveChunkedBuilder::<Int32Type>::new(
697            PlSmallStr::from_static("a"),
698            10,
699            10,
700            DataType::Int32,
701        );
702        builder.append_opt_slice(Some(&[1i32, 2, 3]));
703        builder.append_opt_slice(Some(&[1i32, 2, 3]));
704        let ca = builder.finish();
705
706        let new = ca.cast_with_options(
707            &DataType::List(DataType::Float64.into()),
708            CastOptions::Strict,
709        )?;
710
711        assert_eq!(new.dtype(), &DataType::List(DataType::Float64.into()));
712        Ok(())
713    }
714
715    #[test]
716    #[cfg(feature = "dtype-categorical")]
717    fn test_cast_noop() {
718        // check if we can cast categorical twice without panic
719        let ca = StringChunked::new(PlSmallStr::from_static("foo"), &["bar", "ham"]);
720        let cats = Categories::global();
721        let out = ca
722            .cast_with_options(
723                &DataType::from_categories(cats.clone()),
724                CastOptions::Strict,
725            )
726            .unwrap();
727        let out = out.cast(&DataType::from_categories(cats)).unwrap();
728        assert!(matches!(out.dtype(), &DataType::Categorical(_, _)))
729    }
730}
polars_core/chunked_array/cast.rs

polars_core/chunked_array/
cast.rs
