datui_lib/export/
nested_json.rs

1//! List, array and struct cells as JSON text for one-value-per-field destinations (CSV
2//! export, clipboard copies): JSON keeps struct field names and reads back anywhere
3//! (`str.json_decode`). The text is Polars' JSON writer's, as in NDJSON export.
4//!
5//! Special cases where Polars' writers panic or have no form: binary is standard base64
6//! everywhere; dates and ms/µs datetimes go in as their written text (values past the
7//! calendar as their stored number, as the table shows; ns is always in range);
8//! durations are ISO 8601 seconds as the JSON writer spells them (`PT3723.004S`,
9//! `-PT1.5S`, `P0D`), exact to the nanosecond.
10
11use base64::Engine as _;
12use polars::prelude::*;
13
14/// Whether a column is a list, array or struct, which a delimited writer
15/// takes only as JSON text.
16pub fn is_nested(dtype: &DataType) -> bool {
17    matches!(
18        dtype,
19        DataType::List(_) | DataType::Array(_, _) | DataType::Struct(_)
20    )
21}
22
23fn is_binary(dtype: &DataType) -> bool {
24    matches!(dtype, DataType::Binary | DataType::BinaryOffset)
25}
26
27/// Whether `dtype` is binary or has binary anywhere inside.
28pub fn has_binary(dtype: &DataType) -> bool {
29    match dtype {
30        DataType::List(inner) | DataType::Array(inner, _) => has_binary(inner),
31        DataType::Struct(fields) => fields.iter().any(|f| has_binary(f.dtype())),
32        dtype => is_binary(dtype),
33    }
34}
35
36/// A date or datetime the writers can panic on. Every nanosecond count is a
37/// date (1677 to 2262), so those go to the writers as they are, at no cost.
38fn is_calendar(dtype: &DataType) -> bool {
39    crate::past_calendar::can_leave_calendar(dtype)
40}
41
42/// Whether `dtype` is binary or [`is_calendar`], or has one inside: what the
43/// JSON writer is given as text.
44fn has_json_text(dtype: &DataType) -> bool {
45    match dtype {
46        DataType::List(inner) | DataType::Array(inner, _) => has_json_text(inner),
47        DataType::Struct(fields) => fields.iter().any(|f| has_json_text(f.dtype())),
48        dtype => is_binary(dtype) || is_calendar(dtype),
49    }
50}
51
52/// `dtype` with every binary, date and datetime leaf as a string, the type
53/// [`leaves_as_json_text`] returns.
54fn json_text_dtype(dtype: &DataType) -> DataType {
55    match dtype {
56        DataType::List(inner) => DataType::List(Box::new(json_text_dtype(inner))),
57        DataType::Array(inner, width) => DataType::Array(Box::new(json_text_dtype(inner)), *width),
58        DataType::Struct(fields) => DataType::Struct(
59            fields
60                .iter()
61                .map(|f| Field::new(f.name().clone(), json_text_dtype(f.dtype())))
62                .collect(),
63        ),
64        dtype if is_binary(dtype) || is_calendar(dtype) => DataType::String,
65        dtype => dtype.clone(),
66    }
67}
68
69/// `series` with every binary value, at any depth, as its base64 text, and every
70/// date and datetime as the JSON writer's text ([`calendar_as_text`]). Lists,
71/// arrays and structs keep their shape and their nulls.
72pub fn leaves_as_json_text(series: &Series) -> PolarsResult<Series> {
73    Ok(match series.dtype() {
74        dtype if !has_json_text(dtype) => series.clone(),
75        DataType::List(_) => series
76            .list()?
77            .apply_to_inner(&|inner| leaves_as_json_text(&inner))?
78            .into_series(),
79        DataType::Array(..) => series
80            .array()?
81            .apply_to_inner(&|inner| leaves_as_json_text(&inner))?
82            .into_series(),
83        DataType::Struct(_) => series
84            .struct_()?
85            .try_apply_fields(leaves_as_json_text)?
86            .into_series(),
87        dtype if is_calendar(dtype) => calendar_as_text(series, Writer::Json)?,
88        _ => {
89            let engine = base64::engine::general_purpose::STANDARD;
90            let bytes = series.cast(&DataType::Binary)?;
91            bytes
92                .binary()?
93                .iter()
94                .map(|value| value.map(|b| engine.encode(b)))
95                .collect::<StringChunked>()
96                .with_name(series.name().clone())
97                .into_series()
98        }
99    })
100}
101
102/// `lf` with every column that has binary, a date or a datetime in it as text
103/// ([`leaves_as_json_text`]), in place and under its own name: what a JSON export
104/// writes. Planned, not run.
105pub fn lazy_for_json(mut lf: LazyFrame) -> PolarsResult<LazyFrame> {
106    let schema = lf.collect_schema()?;
107    let exprs: Vec<Expr> = schema
108        .iter()
109        .filter(|(_, dtype)| has_json_text(dtype))
110        .map(|(name, _)| {
111            col(name.clone()).map(
112                |c| leaves_as_json_text(c.as_materialized_series()).map(Column::from),
113                |_, field| {
114                    Ok(Field::new(
115                        field.name().clone(),
116                        json_text_dtype(field.dtype()),
117                    ))
118                },
119            )
120        })
121        .collect();
122    Ok(if exprs.is_empty() {
123        lf
124    } else {
125        lf.with_columns(exprs)
126    })
127}
128
129/// `value` `unit`s as ISO 8601, as Polars' JSON writer (chrono's `TimeDelta`) spells it:
130/// whole seconds plus significant fraction digits, `P0D` for zero, `-` when negative.
131/// Computed here: chrono's range stops short of `i64::MIN` ms.
132pub fn duration_iso(value: i64, unit: TimeUnit, out: &mut String) {
133    use std::fmt::Write as _;
134    let nanos_per_unit: i128 = match unit {
135        TimeUnit::Nanoseconds => 1,
136        TimeUnit::Microseconds => 1_000,
137        TimeUnit::Milliseconds => 1_000_000,
138    };
139    let total = i128::from(value) * nanos_per_unit;
140    if total == 0 {
141        out.push_str("P0D");
142        return;
143    }
144    if total < 0 {
145        out.push('-');
146    }
147    let abs = total.unsigned_abs();
148    let (secs, nanos) = (abs / 1_000_000_000, abs % 1_000_000_000);
149    let _ = write!(out, "PT{secs}");
150    if nanos > 0 {
151        let (mut fraction, mut digits) = (nanos, 9);
152        while fraction % 10 == 0 {
153            fraction /= 10;
154            digits -= 1;
155        }
156        let _ = write!(out, ".{fraction:0digits$}");
157    }
158    out.push('S');
159}
160
161/// A duration column as a String column of [`duration_iso`] text, with its nulls.
162pub fn duration_as_iso(series: &Series) -> PolarsResult<Series> {
163    let DataType::Duration(unit) = series.dtype() else {
164        polars_bail!(InvalidOperation: "expected a duration, got {}", series.dtype());
165    };
166    let unit = *unit;
167    Ok(series
168        .to_physical_repr()
169        .i64()?
170        .apply_into_string_amortized(|value, out| duration_iso(value, unit, out))
171        .with_name(series.name().clone())
172        .into_series())
173}
174
175/// Whether a column needs to become text before a CSV writer takes it: it is
176/// nested, binary, a duration, or a date or datetime in ms or us, which the
177/// writer can panic on.
178pub fn needs_text(dtype: &DataType) -> bool {
179    is_nested(dtype)
180        || is_binary(dtype)
181        || is_calendar(dtype)
182        || matches!(dtype, DataType::Duration(_))
183}
184
185/// One column that [`needs_text`] as the String column a CSV cell or a copy
186/// holds: JSON, base64, ISO 8601 or the CSV writer's own text, by its type.
187fn column_as_text(column: &Column) -> PolarsResult<Column> {
188    match column.dtype() {
189        DataType::Duration(_) => duration_as_iso(column.as_materialized_series()).map(Column::from),
190        dtype if is_calendar(dtype) => {
191            calendar_as_text(column.as_materialized_series(), Writer::Csv).map(Column::from)
192        }
193        dtype if is_binary(dtype) => {
194            leaves_as_json_text(column.as_materialized_series()).map(Column::from)
195        }
196        _ => column_as_json(column),
197    }
198}
199
200/// The writer whose text for a date or datetime [`calendar_as_text`] writes.
201#[derive(Debug, Clone, Copy, PartialEq, Eq)]
202pub enum Writer {
203    Csv,
204    Json,
205}
206
207/// A date or datetime column as the text `writer` writes for it, with its nulls;
208/// a value past the calendar's range, on which the writers panic, as its stored
209/// number ([`crate::exact::out_of_range`]), as the table shows it.
210pub fn calendar_as_text(series: &Series, writer: Writer) -> PolarsResult<Series> {
211    // The writers' own defaults: the CSV writer's formats, and the JSON
212    // writer's chrono display (`to_rfc3339` with a zone).
213    let format = match (series.dtype(), writer) {
214        (DataType::Date, _) => "%Y-%m-%d",
215        (DataType::Datetime(unit, zone), Writer::Csv) => match (unit, zone.is_some()) {
216            (TimeUnit::Milliseconds, false) => "%FT%H:%M:%S.%3f",
217            (TimeUnit::Milliseconds, true) => "%FT%H:%M:%S.%3f%z",
218            (TimeUnit::Microseconds, false) => "%FT%H:%M:%S.%6f",
219            (TimeUnit::Microseconds, true) => "%FT%H:%M:%S.%6f%z",
220            (TimeUnit::Nanoseconds, false) => "%FT%H:%M:%S.%9f",
221            (TimeUnit::Nanoseconds, true) => "%FT%H:%M:%S.%9f%z",
222        },
223        (DataType::Datetime(_, None), Writer::Json) => "%Y-%m-%d %H:%M:%S%.f",
224        (DataType::Datetime(_, Some(_)), Writer::Json) => "%Y-%m-%dT%H:%M:%S%.f%:z",
225        (dtype, _) => polars_bail!(InvalidOperation: "expected a date or datetime, got {dtype}"),
226    };
227    crate::past_calendar::text_or_stored(series, |s| match s.dtype() {
228        DataType::Date => s.date()?.to_string(format),
229        _ => s.datetime()?.to_string(format),
230    })
231}
232
233/// One nested column as a String column of JSON, null where the value is null.
234pub fn column_as_json(column: &Column) -> PolarsResult<Column> {
235    let series = leaves_as_json_text(column.as_materialized_series())?;
236    let chunks = (0..series.n_chunks()).map(|i| {
237        let array = series.to_arrow(i, CompatLevel::newest());
238        // The writer spells a missing value `null`; the cell stays empty
239        // instead, as a null does everywhere else in a CSV.
240        polars_json::json::write::serialize_to_utf8(array.as_ref())
241            .with_validity(array.validity().cloned())
242    });
243    Ok(StringChunked::from_chunk_iter(series.name().clone(), chunks).into_column())
244}
245
246/// `lf` with every column a CSV writer cannot take ([`needs_text`]) replaced
247/// by its text, in place and under its own name: what a CSV export writes.
248/// Planned, not run: the text is built as the rows are collected.
249pub fn lazy_as_json(mut lf: LazyFrame) -> PolarsResult<LazyFrame> {
250    let schema = lf.collect_schema()?;
251    let exprs: Vec<Expr> = schema
252        .iter()
253        .filter(|(_, dtype)| needs_text(dtype))
254        .map(|(name, _)| {
255            col(name.clone()).map(
256                |c| column_as_text(&c),
257                |_, field| Ok(Field::new(field.name().clone(), DataType::String)),
258            )
259        })
260        .collect();
261    Ok(if exprs.is_empty() {
262        lf
263    } else {
264        lf.with_columns(exprs)
265    })
266}
267
268/// [`lazy_as_json`] for frames already in memory: a copy's rows, as the CSV
269/// writer takes them.
270pub fn frame_as_json(df: &DataFrame) -> PolarsResult<DataFrame> {
271    frame_as_text(df, needs_text)
272}
273
274/// [`frame_as_json`] with dates and datetimes kept in their own type: the cells
275/// a Markdown or HTML copy writes through [`crate::exact::value_text`], which
276/// spells one past the calendar as its stored number itself.
277pub fn frame_as_cells(df: &DataFrame) -> PolarsResult<DataFrame> {
278    frame_as_text(df, |dtype| needs_text(dtype) && !is_calendar(dtype))
279}
280
281fn frame_as_text(df: &DataFrame, converts: impl Fn(&DataType) -> bool) -> PolarsResult<DataFrame> {
282    let mut out = df.clone();
283    for column in df.columns() {
284        if converts(column.dtype()) {
285            out.with_column(column_as_text(column)?)?;
286        }
287    }
288    Ok(out)
289}
290
291#[cfg(test)]
292pub(crate) mod tests;
datui_lib/export/nested_json.rs

datui_lib/export/
nested_json.rs
