polars_core/series/ops/
canonicalize_maps.rs

1#[cfg(feature = "dtype-map")]
2use crate::chunked_array::logical::canonicalize_map_storage;
3use crate::prelude::*;
4
5impl Series {
6    /// Canonicalize all nested `Map`s bottom-up using first-position/last-value
7    /// semantics. Returns `None` if unchanged.
8    pub fn canonicalize_maps(&self) -> PolarsResult<Option<Series>> {
9        #[cfg(feature = "dtype-map")]
10        {
11            canonicalize_maps_rec(self)
12        }
13        #[cfg(not(feature = "dtype-map"))]
14        {
15            Ok(None)
16        }
17    }
18}
19
20#[cfg(feature = "dtype-map")]
21fn canonicalize_maps_rec(series: &Series) -> PolarsResult<Option<Series>> {
22    if !series.dtype().contains_map() {
23        return Ok(None);
24    }
25
26    match series.dtype() {
27        DataType::Map(_, _) => {
28            let map = series.map().unwrap();
29
30            // Parent keys are row-encoded, so canonicalize nested maps first.
31            let nested = canonicalize_maps_rec(map.storage())?;
32            let storage = nested.as_ref().unwrap_or(map.storage());
33            let deduped = canonicalize_map_storage(storage)?;
34
35            match deduped.or(nested) {
36                None => Ok(None),
37                Some(storage) => Ok(Some(
38                    unsafe { MapChunked::from_storage_unchecked(map.dtype().clone(), storage) }
39                        .into_series(),
40                )),
41            }
42        },
43        DataType::List(_) => {
44            let ca = series.list().unwrap();
45            Ok(canonicalize_maps_rec(&ca.get_inner())?
46                .map(|values| ca.with_inner_values(&values).into_series()))
47        },
48        #[cfg(feature = "dtype-array")]
49        DataType::Array(_, _) => {
50            let ca = series.array().unwrap();
51            Ok(canonicalize_maps_rec(&ca.get_inner())?
52                .map(|values| ca.with_inner_values(&values).into_series()))
53        },
54        #[cfg(feature = "dtype-struct")]
55        DataType::Struct(_) => {
56            let ca = series.struct_().unwrap();
57
58            // Avoid rebuilding the struct when no field changes.
59            let fields = ca.fields_as_series();
60            let mut new_fields = Vec::with_capacity(fields.len());
61            let mut changed = false;
62            for field in &fields {
63                let new_field = canonicalize_maps_rec(field)?;
64                changed |= new_field.is_some();
65                new_fields.push(new_field);
66            }
67
68            if !changed {
69                return Ok(None);
70            }
71
72            // `try_apply_fields` preserves the outer validity.
73            let mut new_fields = new_fields.into_iter();
74            let out = ca.try_apply_fields(|field| {
75                Ok(new_fields.next().unwrap().unwrap_or_else(|| field.clone()))
76            })?;
77
78            Ok(Some(out.into_series()))
79        },
80        #[cfg(feature = "dtype-extension")]
81        DataType::Extension(typ, _) => {
82            let ext = series.ext().unwrap();
83            Ok(canonicalize_maps_rec(ext.storage())?.map(|s| s.into_extension(typ.clone())))
84        },
85        _ => Ok(None),
86    }
87}
polars_core/series/ops/canonicalize_maps.rs

polars_core/series/ops/
canonicalize_maps.rs
