polars_core/series/ops/
json_map.rs

1//! JSON support for `Map`. A Map is a JSON object, so its keys must be strings. Readers decode
2//! a Map as `List(Struct {key: String, value})` in source order and rebuild it with
3//! [`Series::from_json_decoded`].
4#[cfg(feature = "dtype-map")]
5use polars_arrow::bitmap::{Bitmap, BitmapBuilder};
6
7use crate::chunked_array::cast::CastOptions;
8#[cfg(feature = "dtype-map")]
9use crate::chunked_array::logical::try_apply_map_entries;
10use crate::prelude::*;
11
12impl DataType {
13    /// Reject Maps with unsupported JSON key types, including nested Maps.
14    pub fn ensure_json_map_keys(&self) -> PolarsResult<()> {
15        if !self.contains_map() {
16            return Ok(());
17        }
18        self.try_visit_with(|_dtype| {
19            #[cfg(feature = "dtype-map")]
20            if let DataType::Map(key, _) = _dtype {
21                polars_ensure!(
22                    is_json_map_key(key),
23                    ComputeError:
24                    "JSON Map keys must be String, Categorical or Enum, got `{key}`.\n\nCast keys to String or use `Expr.map.entries` to write a list of structs."
25                );
26            }
27            Ok(())
28        })
29    }
30
31    /// The dtype a JSON reader decodes into before [`Series::from_json_decoded`].
32    ///
33    /// Maps become `List(Struct {key: String, value})`, Arrays containing a Map become Lists,
34    /// Extensions containing a Map become their storage and Enum/Categorical leaves become
35    /// String.
36    pub fn json_decode_dtype(&self) -> DataType {
37        use DataType as D;
38        match self {
39            #[cfg(feature = "dtype-map")]
40            D::Map(_, value) => D::Map(Box::new(D::String), Box::new(value.json_decode_dtype()))
41                .map_storage_dtype()
42                .unwrap(),
43            D::List(inner) => D::List(Box::new(inner.json_decode_dtype())),
44            #[cfg(feature = "dtype-array")]
45            D::Array(inner, _) if inner.contains_map() => {
46                D::List(Box::new(inner.json_decode_dtype()))
47            },
48            #[cfg(feature = "dtype-array")]
49            D::Array(inner, width) => D::Array(Box::new(inner.json_decode_dtype()), *width),
50            #[cfg(feature = "dtype-struct")]
51            D::Struct(fields) => D::Struct(
52                fields
53                    .iter()
54                    .map(|f| Field::new(f.name.clone(), f.dtype.json_decode_dtype()))
55                    .collect(),
56            ),
57            #[cfg(feature = "dtype-extension")]
58            D::Extension(_, storage) if storage.contains_map() => storage.json_decode_dtype(),
59            #[cfg(feature = "dtype-extension")]
60            D::Extension(ext, storage) => {
61                D::Extension(ext.clone(), Box::new(storage.json_decode_dtype()))
62            },
63            #[cfg(feature = "dtype-categorical")]
64            D::Enum(..) | D::Categorical(..) => D::String,
65            dt => dt.clone(),
66        }
67    }
68}
69
70#[cfg(feature = "dtype-map")]
71fn is_json_map_key(dtype: &DataType) -> bool {
72    match dtype {
73        DataType::String => true,
74        #[cfg(feature = "dtype-categorical")]
75        DataType::Categorical(..) | DataType::Enum(..) => true,
76        _ => false,
77    }
78}
79
80impl Series {
81    /// Build `target` from a series decoded as [`DataType::json_decode_dtype`].
82    ///
83    /// Map keys are cast to the key dtype and deduplicated with first-position/last-value
84    /// semantics. With `ignore_errors`, rows with an unknown Enum key or a wrong Array width
85    /// become null.
86    pub fn from_json_decoded(self, target: &DataType, ignore_errors: bool) -> PolarsResult<Series> {
87        #[cfg(feature = "dtype-map")]
88        if target.contains_map() {
89            target.ensure_json_map_keys()?;
90            return from_json_decoded_rec(&self, target, ignore_errors);
91        }
92        cast_leaf(self, target, ignore_errors)
93    }
94}
95
96fn cast_leaf(s: Series, target: &DataType, ignore_errors: bool) -> PolarsResult<Series> {
97    if s.dtype() == target {
98        return Ok(s);
99    }
100    let options = if ignore_errors {
101        CastOptions::NonStrict
102    } else {
103        CastOptions::Strict
104    };
105    s.cast_with_options(target, options)
106}
107
108#[cfg(feature = "dtype-map")]
109fn from_json_decoded_rec(
110    series: &Series,
111    target: &DataType,
112    ignore_errors: bool,
113) -> PolarsResult<Series> {
114    match target {
115        DataType::Map(key_dtype, value_dtype) => {
116            let mut ok_keys: Option<Bitmap> = None;
117            let storage = try_apply_map_entries(series.list()?, |key, value| {
118                // JSON object keys are never null, so a null marks an unknown Enum label.
119                let decoded = key.cast_with_options(key_dtype, CastOptions::NonStrict)?;
120                if let Some(ok) = decoded.rechunk_validity().filter(|v| v.unset_bits() > 0) {
121                    if !ignore_errors {
122                        let idx = ok.iter().position(|ok| !ok).unwrap();
123                        let key = key.str()?.get(idx).unwrap();
124                        polars_bail!(
125                            ComputeError:
126                            "JSON object key \"{key}\" is not a valid Map key of type `{key_dtype}`"
127                        );
128                    }
129                    ok_keys = Some(ok);
130                }
131                let value = from_json_decoded_rec(value, value_dtype, ignore_errors)?;
132                Ok((decoded, value))
133            })?
134            .into_series();
135
136            let storage = match ok_keys {
137                Some(ok) => {
138                    null_rows_where(&storage, |start, len| ok.null_count_range(start, len) > 0)
139                },
140                None => storage,
141            };
142            Ok(MapChunked::try_from_storage(target.clone(), storage)?.into_series())
143        },
144        DataType::List(inner) if inner.contains_map() => {
145            let ca = series.list()?;
146            let values = from_json_decoded_rec(&ca.get_inner(), inner, ignore_errors)?;
147            Ok(ca.with_inner_values(&values).into_series())
148        },
149        #[cfg(feature = "dtype-array")]
150        DataType::Array(inner, width) if inner.contains_map() => {
151            let ca = series.list()?;
152            let values = from_json_decoded_rec(&ca.get_inner(), inner, ignore_errors)?;
153            let mut list = ca.with_inner_values(&values).into_series();
154            if ignore_errors {
155                list = null_rows_where(&list, |_, len| len != *width);
156            }
157            list.cast_with_options(target, CastOptions::Strict)
158        },
159        DataType::Struct(fields) if target.contains_map() => {
160            let ca = series.struct_()?;
161            let out = ca.try_apply_fields(|field| {
162                let f = fields
163                    .iter()
164                    .find(|f| f.name() == field.name())
165                    .expect("the decode dtype keeps the target's struct fields");
166                from_json_decoded_rec(field, f.dtype(), ignore_errors)
167            })?;
168            Ok(out.into_series())
169        },
170        #[cfg(feature = "dtype-extension")]
171        DataType::Extension(typ, storage) if storage.contains_map() => {
172            Ok(from_json_decoded_rec(series, storage, ignore_errors)?.into_extension(typ.clone()))
173        },
174        _ => cast_leaf(series.clone(), target, ignore_errors),
175    }
176}
177
178/// Null the valid rows of a list series for which `bad(start, len)` holds. `start` indexes
179/// the concatenated offset windows of all chunks.
180#[cfg(feature = "dtype-map")]
181fn null_rows_where(list: &Series, mut bad: impl FnMut(usize, usize) -> bool) -> Series {
182    let ca = list.list().unwrap();
183    let mut validity = BitmapBuilder::with_capacity(ca.len());
184    let mut changed = false;
185    let mut base = 0;
186    for arr in ca.downcast_iter() {
187        let offsets = arr.offsets();
188        let first = *offsets.first() as usize;
189        for row in 0..arr.len() {
190            let (start, end) = offsets.start_end(row);
191            let valid = arr.is_valid(row);
192            let keep = valid && !bad(base + start - first, end - start);
193            changed |= valid != keep;
194            validity.push(keep);
195        }
196        base += offsets.range() as usize;
197    }
198
199    if changed {
200        list.with_validity(Some(validity.freeze()))
201    } else {
202        list.clone()
203    }
204}
polars_core/series/ops/json_map.rs

polars_core/series/ops/
json_map.rs
