Skip to main content

datui_lib/
nested_json.rs

1//! List, array and struct cells as JSON text, for the destinations that hold
2//! one value per field: a CSV export and every clipboard copy. JSON rather than
3//! the table's own rendering (`[1, 2]`, `{1,"a"}`) because it keeps struct field
4//! names and reads back with any JSON parser, Polars' `str.json_decode` included.
5//!
6//! The text is Polars' own JSON writer, the one NDJSON export uses, so a list
7//! reads the same in a CSV as in a `.jsonl` written from the same view.
8//!
9//! Binary has no JSON or CSV form, and Polars' JSON writer panics on it, so it
10//! is written as standard base64 text wherever it sits: a CSV, JSON or NDJSON
11//! export and a copy all spell the same bytes the same way.
12//!
13//! Polars' writers also panic on a date or datetime past the calendar's range
14//! (a sentinel like `i64::MIN + 1` microseconds), so dates and millisecond and
15//! microsecond datetimes are given to them as the text they would write, and
16//! such a value as its stored number, as the table shows it. A nanosecond
17//! count is always a date and goes to the writers as it is.
18//!
19//! A duration has no CSV form either, and is written as the JSON writer spells
20//! it: ISO 8601 in seconds (`PT3723.004S`, `-PT1.5S`, `P0D`). That is exact to
21//! the nanosecond in every unit, and reads the same alone in a CSV cell or a
22//! copy, inside a list, and in a JSON export.
23
24use base64::Engine as _;
25use polars::prelude::*;
26
27/// Whether a column is a list, array or struct, which a delimited writer
28/// takes only as JSON text.
29pub fn is_nested(dtype: &DataType) -> bool {
30    matches!(
31        dtype,
32        DataType::List(_) | DataType::Array(_, _) | DataType::Struct(_)
33    )
34}
35
36fn is_binary(dtype: &DataType) -> bool {
37    matches!(dtype, DataType::Binary | DataType::BinaryOffset)
38}
39
40/// Whether `dtype` is binary or has binary anywhere inside.
41pub fn has_binary(dtype: &DataType) -> bool {
42    match dtype {
43        DataType::List(inner) | DataType::Array(inner, _) => has_binary(inner),
44        DataType::Struct(fields) => fields.iter().any(|f| has_binary(f.dtype())),
45        dtype => is_binary(dtype),
46    }
47}
48
49/// A date or datetime the writers can panic on. Every nanosecond count is a
50/// date (1677 to 2262), so those go to the writers as they are, at no cost.
51fn is_calendar(dtype: &DataType) -> bool {
52    crate::past_calendar::can_leave_calendar(dtype)
53}
54
55/// Whether `dtype` is binary or [`is_calendar`], or has one inside: what the
56/// JSON writer is given as text.
57fn has_json_text(dtype: &DataType) -> bool {
58    match dtype {
59        DataType::List(inner) | DataType::Array(inner, _) => has_json_text(inner),
60        DataType::Struct(fields) => fields.iter().any(|f| has_json_text(f.dtype())),
61        dtype => is_binary(dtype) || is_calendar(dtype),
62    }
63}
64
65/// `dtype` with every binary, date and datetime leaf as a string, the type
66/// [`leaves_as_json_text`] returns.
67fn json_text_dtype(dtype: &DataType) -> DataType {
68    match dtype {
69        DataType::List(inner) => DataType::List(Box::new(json_text_dtype(inner))),
70        DataType::Array(inner, width) => DataType::Array(Box::new(json_text_dtype(inner)), *width),
71        DataType::Struct(fields) => DataType::Struct(
72            fields
73                .iter()
74                .map(|f| Field::new(f.name().clone(), json_text_dtype(f.dtype())))
75                .collect(),
76        ),
77        dtype if is_binary(dtype) || is_calendar(dtype) => DataType::String,
78        dtype => dtype.clone(),
79    }
80}
81
82/// `series` with every binary value, at any depth, as its base64 text, and every
83/// date and datetime as the JSON writer's text ([`calendar_as_text`]). Lists,
84/// arrays and structs keep their shape and their nulls.
85pub fn leaves_as_json_text(series: &Series) -> PolarsResult<Series> {
86    Ok(match series.dtype() {
87        dtype if !has_json_text(dtype) => series.clone(),
88        DataType::List(_) => series
89            .list()?
90            .apply_to_inner(&|inner| leaves_as_json_text(&inner))?
91            .into_series(),
92        DataType::Array(..) => series
93            .array()?
94            .apply_to_inner(&|inner| leaves_as_json_text(&inner))?
95            .into_series(),
96        DataType::Struct(_) => series
97            .struct_()?
98            .try_apply_fields(leaves_as_json_text)?
99            .into_series(),
100        dtype if is_calendar(dtype) => calendar_as_text(series, Writer::Json)?,
101        _ => {
102            let engine = base64::engine::general_purpose::STANDARD;
103            let bytes = series.cast(&DataType::Binary)?;
104            bytes
105                .binary()?
106                .iter()
107                .map(|value| value.map(|b| engine.encode(b)))
108                .collect::<StringChunked>()
109                .with_name(series.name().clone())
110                .into_series()
111        }
112    })
113}
114
115/// `lf` with every column that has binary, a date or a datetime in it as text
116/// ([`leaves_as_json_text`]), in place and under its own name: what a JSON export
117/// writes. Planned, not run.
118pub fn lazy_for_json(mut lf: LazyFrame) -> PolarsResult<LazyFrame> {
119    let schema = lf.collect_schema()?;
120    let exprs: Vec<Expr> = schema
121        .iter()
122        .filter(|(_, dtype)| has_json_text(dtype))
123        .map(|(name, _)| {
124            col(name.clone()).map(
125                |c| leaves_as_json_text(c.as_materialized_series()).map(Column::from),
126                |_, field| {
127                    Ok(Field::new(
128                        field.name().clone(),
129                        json_text_dtype(field.dtype()),
130                    ))
131                },
132            )
133        })
134        .collect();
135    Ok(if exprs.is_empty() {
136        lf
137    } else {
138        lf.with_columns(exprs)
139    })
140}
141
142/// `value` `unit`s as ISO 8601 text, the way Polars' JSON writer spells a
143/// duration (chrono's `TimeDelta` display): whole seconds and the fraction's
144/// significant digits, `P0D` for zero, a leading `-` when negative. Computed
145/// here rather than through chrono, whose range ends short of `i64::MIN` ms.
146pub fn duration_iso(value: i64, unit: TimeUnit, out: &mut String) {
147    use std::fmt::Write as _;
148    let nanos_per_unit: i128 = match unit {
149        TimeUnit::Nanoseconds => 1,
150        TimeUnit::Microseconds => 1_000,
151        TimeUnit::Milliseconds => 1_000_000,
152    };
153    let total = i128::from(value) * nanos_per_unit;
154    if total == 0 {
155        out.push_str("P0D");
156        return;
157    }
158    if total < 0 {
159        out.push('-');
160    }
161    let abs = total.unsigned_abs();
162    let (secs, nanos) = (abs / 1_000_000_000, abs % 1_000_000_000);
163    let _ = write!(out, "PT{secs}");
164    if nanos > 0 {
165        let (mut fraction, mut digits) = (nanos, 9);
166        while fraction % 10 == 0 {
167            fraction /= 10;
168            digits -= 1;
169        }
170        let _ = write!(out, ".{fraction:0digits$}");
171    }
172    out.push('S');
173}
174
175/// A duration column as a String column of [`duration_iso`] text, with its nulls.
176pub fn duration_as_iso(series: &Series) -> PolarsResult<Series> {
177    let DataType::Duration(unit) = series.dtype() else {
178        polars_bail!(InvalidOperation: "expected a duration, got {}", series.dtype());
179    };
180    let unit = *unit;
181    Ok(series
182        .to_physical_repr()
183        .i64()?
184        .apply_into_string_amortized(|value, out| duration_iso(value, unit, out))
185        .with_name(series.name().clone())
186        .into_series())
187}
188
189/// Whether a column needs to become text before a CSV writer takes it: it is
190/// nested, binary, a duration, or a date or datetime in ms or us, which the
191/// writer can panic on.
192pub fn needs_text(dtype: &DataType) -> bool {
193    is_nested(dtype)
194        || is_binary(dtype)
195        || is_calendar(dtype)
196        || matches!(dtype, DataType::Duration(_))
197}
198
199/// One column that [`needs_text`] as the String column a CSV cell or a copy
200/// holds: JSON, base64, ISO 8601 or the CSV writer's own text, by its type.
201fn column_as_text(column: &Column) -> PolarsResult<Column> {
202    match column.dtype() {
203        DataType::Duration(_) => duration_as_iso(column.as_materialized_series()).map(Column::from),
204        dtype if is_calendar(dtype) => {
205            calendar_as_text(column.as_materialized_series(), Writer::Csv).map(Column::from)
206        }
207        dtype if is_binary(dtype) => {
208            leaves_as_json_text(column.as_materialized_series()).map(Column::from)
209        }
210        _ => column_as_json(column),
211    }
212}
213
214/// The writer whose text for a date or datetime [`calendar_as_text`] writes.
215#[derive(Debug, Clone, Copy, PartialEq, Eq)]
216pub enum Writer {
217    Csv,
218    Json,
219}
220
221/// A date or datetime column as the text `writer` writes for it, with its nulls;
222/// a value past the calendar's range, on which the writers panic, as its stored
223/// number ([`crate::exact::out_of_range`]), as the table shows it.
224pub fn calendar_as_text(series: &Series, writer: Writer) -> PolarsResult<Series> {
225    // The writers' own defaults: the CSV writer's formats, and the JSON
226    // writer's chrono display (`to_rfc3339` with a zone).
227    let format = match (series.dtype(), writer) {
228        (DataType::Date, _) => "%Y-%m-%d",
229        (DataType::Datetime(unit, zone), Writer::Csv) => match (unit, zone.is_some()) {
230            (TimeUnit::Milliseconds, false) => "%FT%H:%M:%S.%3f",
231            (TimeUnit::Milliseconds, true) => "%FT%H:%M:%S.%3f%z",
232            (TimeUnit::Microseconds, false) => "%FT%H:%M:%S.%6f",
233            (TimeUnit::Microseconds, true) => "%FT%H:%M:%S.%6f%z",
234            (TimeUnit::Nanoseconds, false) => "%FT%H:%M:%S.%9f",
235            (TimeUnit::Nanoseconds, true) => "%FT%H:%M:%S.%9f%z",
236        },
237        (DataType::Datetime(_, None), Writer::Json) => "%Y-%m-%d %H:%M:%S%.f",
238        (DataType::Datetime(_, Some(_)), Writer::Json) => "%Y-%m-%dT%H:%M:%S%.f%:z",
239        (dtype, _) => polars_bail!(InvalidOperation: "expected a date or datetime, got {dtype}"),
240    };
241    crate::past_calendar::text_or_stored(series, |s| match s.dtype() {
242        DataType::Date => s.date()?.to_string(format),
243        _ => s.datetime()?.to_string(format),
244    })
245}
246
247/// One nested column as a String column of JSON, null where the value is null.
248pub fn column_as_json(column: &Column) -> PolarsResult<Column> {
249    let series = leaves_as_json_text(column.as_materialized_series())?;
250    let chunks = (0..series.n_chunks()).map(|i| {
251        let array = series.to_arrow(i, CompatLevel::newest());
252        // The writer spells a missing value `null`; the cell stays empty
253        // instead, as a null does everywhere else in a CSV.
254        polars_json::json::write::serialize_to_utf8(array.as_ref())
255            .with_validity(array.validity().cloned())
256    });
257    Ok(StringChunked::from_chunk_iter(series.name().clone(), chunks).into_column())
258}
259
260/// `lf` with every column a CSV writer cannot take ([`needs_text`]) replaced
261/// by its text, in place and under its own name: what a CSV export writes.
262/// Planned, not run: the text is built as the rows are collected.
263pub fn lazy_as_json(mut lf: LazyFrame) -> PolarsResult<LazyFrame> {
264    let schema = lf.collect_schema()?;
265    let exprs: Vec<Expr> = schema
266        .iter()
267        .filter(|(_, dtype)| needs_text(dtype))
268        .map(|(name, _)| {
269            col(name.clone()).map(
270                |c| column_as_text(&c),
271                |_, field| Ok(Field::new(field.name().clone(), DataType::String)),
272            )
273        })
274        .collect();
275    Ok(if exprs.is_empty() {
276        lf
277    } else {
278        lf.with_columns(exprs)
279    })
280}
281
282/// [`lazy_as_json`] for frames already in memory: a copy's rows, as the CSV
283/// writer takes them.
284pub fn frame_as_json(df: &DataFrame) -> PolarsResult<DataFrame> {
285    frame_as_text(df, needs_text)
286}
287
288/// [`frame_as_json`] with dates and datetimes kept in their own type: the cells
289/// a Markdown or HTML copy writes through [`crate::exact::value_text`], which
290/// spells one past the calendar as its stored number itself.
291pub fn frame_as_cells(df: &DataFrame) -> PolarsResult<DataFrame> {
292    frame_as_text(df, |dtype| needs_text(dtype) && !is_calendar(dtype))
293}
294
295fn frame_as_text(df: &DataFrame, converts: impl Fn(&DataType) -> bool) -> PolarsResult<DataFrame> {
296    let mut out = df.clone();
297    for column in df.columns() {
298        if converts(column.dtype()) {
299            out.with_column(column_as_text(column)?)?;
300        }
301    }
302    Ok(out)
303}
304
305#[cfg(test)]
306pub(crate) mod tests {
307    use super::*;
308
309    /// The raw values behind [`durations`], one row each.
310    const DURATION_VALUES: [Option<i64>; 8] = [
311        Some(3_723_004),
312        None,
313        Some(-1_500),
314        Some(0),
315        Some(1),
316        Some(-1),
317        Some(i64::MAX),
318        Some(-i64::MAX),
319    ];
320
321    /// A duration column per unit over the same raw values, with a null, zero,
322    /// negatives and the ends of the range.
323    pub(crate) fn durations() -> DataFrame {
324        let values = Series::new("".into(), DURATION_VALUES);
325        let columns = [
326            ("ms", TimeUnit::Milliseconds),
327            ("us", TimeUnit::Microseconds),
328            ("ns", TimeUnit::Nanoseconds),
329        ]
330        .map(|(name, unit)| {
331            values
332                .cast(&DataType::Duration(unit))
333                .unwrap()
334                .with_name(name.into())
335                .into_column()
336        });
337        DataFrame::new_infer_height(columns.to_vec()).unwrap()
338    }
339
340    /// [`durations`] as text, by column; a null is None.
341    pub(crate) fn duration_text() -> [(&'static str, [Option<&'static str>; 8]); 3] {
342        [
343            (
344                "ms",
345                [
346                    Some("PT3723.004S"),
347                    None,
348                    Some("-PT1.5S"),
349                    Some("P0D"),
350                    Some("PT0.001S"),
351                    Some("-PT0.001S"),
352                    Some("PT9223372036854775.807S"),
353                    Some("-PT9223372036854775.807S"),
354                ],
355            ),
356            (
357                "us",
358                [
359                    Some("PT3.723004S"),
360                    None,
361                    Some("-PT0.0015S"),
362                    Some("P0D"),
363                    Some("PT0.000001S"),
364                    Some("-PT0.000001S"),
365                    Some("PT9223372036854.775807S"),
366                    Some("-PT9223372036854.775807S"),
367                ],
368            ),
369            (
370                "ns",
371                [
372                    Some("PT0.003723004S"),
373                    None,
374                    Some("-PT0.0000015S"),
375                    Some("P0D"),
376                    Some("PT0.000000001S"),
377                    Some("-PT0.000000001S"),
378                    Some("PT9223372036.854775807S"),
379                    Some("-PT9223372036.854775807S"),
380                ],
381            ),
382        ]
383    }
384
385    /// A duration is ISO 8601 text in every unit, with its nulls, the same
386    /// text in memory and planned, and the same text an NDJSON export writes.
387    #[test]
388    fn durations_are_iso_8601_as_the_json_writer_spells_them() {
389        let df = durations();
390        let cells = frame_as_json(&df).unwrap();
391        for (name, expected) in duration_text() {
392            let text = cells.column(name).unwrap().str().unwrap();
393            assert_eq!(text.iter().collect::<Vec<_>>(), expected, "{name}");
394        }
395        let lazy = lazy_as_json(df.clone().lazy()).unwrap().collect().unwrap();
396        assert!(cells.equals_missing(&lazy), "{cells}\n{lazy}");
397
398        let mut ndjson = Vec::new();
399        JsonWriter::new(&mut ndjson)
400            .with_json_format(JsonFormat::JsonLines)
401            .finish(&mut df.clone())
402            .unwrap();
403        let json = |name: &str, i: usize| {
404            cells
405                .column(name)
406                .unwrap()
407                .str()
408                .unwrap()
409                .get(i)
410                .map_or("null".to_string(), |s| format!("\"{s}\""))
411        };
412        let rebuilt: String = (0..cells.height())
413            .map(|i| {
414                format!(
415                    "{{\"ms\":{},\"us\":{},\"ns\":{}}}\n",
416                    json("ms", i),
417                    json("us", i),
418                    json("ns", i)
419                )
420            })
421            .collect();
422        assert_eq!(rebuilt, String::from_utf8(ndjson).unwrap());
423    }
424
425    /// Past the end of chrono's range, where the JSON writer gives up, the
426    /// text is still exact.
427    #[test]
428    fn the_longest_negative_duration_is_exact() {
429        let mut text = String::new();
430        duration_iso(i64::MIN, TimeUnit::Milliseconds, &mut text);
431        assert_eq!(text, "-PT9223372036854775.808S");
432        text.clear();
433        duration_iso(i64::MIN, TimeUnit::Nanoseconds, &mut text);
434        assert_eq!(text, "-PT9223372036.854775808S");
435    }
436
437    /// Dates and datetimes in every unit, with and without a zone, and their
438    /// nulls, years before 0 and past 9999 among them; `past` adds the ends of
439    /// the stored range, which no writer takes.
440    pub(crate) fn calendar(past: bool) -> DataFrame {
441        let mut stamps = vec![
442            Some(0i64),
443            None,
444            Some(-1),
445            Some(1_700_000_000_123),
446            Some(-62_000_000_000_000),
447            Some(-100_000_000_000_000),
448            Some(300_000_000_000_000),
449        ];
450        let mut days = vec![
451            Some(0i32),
452            None,
453            Some(-1),
454            Some(19_724),
455            Some(-800_000),
456            Some(-1_000_000),
457            Some(3_000_000),
458        ];
459        if past {
460            stamps.extend([Some(i64::MIN + 1), Some(i64::MAX)]);
461            days.extend([Some(i32::MIN), Some(i32::MAX)]);
462        }
463        let paris = TimeZone::opt_try_new(Some("Europe/Paris")).unwrap();
464        let mut columns = vec![
465            Series::new("d".into(), days)
466                .cast(&DataType::Date)
467                .unwrap()
468                .into_column(),
469        ];
470        for (unit, name) in [
471            (TimeUnit::Milliseconds, "ms"),
472            (TimeUnit::Microseconds, "us"),
473            (TimeUnit::Nanoseconds, "ns"),
474        ] {
475            for (zone, suffix) in [(None, ""), (paris.clone(), "_tz")] {
476                columns.push(
477                    Series::new(format!("{name}{suffix}").into(), &stamps)
478                        .cast(&DataType::Datetime(unit, zone))
479                        .unwrap()
480                        .into_column(),
481                );
482            }
483        }
484        DataFrame::new_infer_height(columns).unwrap()
485    }
486
487    /// Given as text, dates and datetimes read exactly as the CSV and JSON
488    /// writers write them.
489    #[test]
490    fn dates_as_text_are_what_the_writers_write() {
491        let df = calendar(false);
492        let mut csv = Vec::new();
493        CsvWriter::new(&mut csv).finish(&mut df.clone()).unwrap();
494        let mut as_text = Vec::new();
495        CsvWriter::new(&mut as_text)
496            .finish(&mut frame_as_json(&df).unwrap())
497            .unwrap();
498        assert_eq!(
499            String::from_utf8(as_text).unwrap(),
500            String::from_utf8(csv).unwrap()
501        );
502        let lazy = lazy_as_json(df.clone().lazy()).unwrap().collect().unwrap();
503        assert!(frame_as_json(&df).unwrap().equals_missing(&lazy));
504
505        let mut json = Vec::new();
506        JsonWriter::new(&mut json)
507            .with_json_format(JsonFormat::JsonLines)
508            .finish(&mut df.clone())
509            .unwrap();
510        let mut prepared = lazy_for_json(df.clone().lazy()).unwrap().collect().unwrap();
511        // Nanosecond datetimes go to the writer as they are.
512        assert!(
513            prepared
514                .columns()
515                .iter()
516                .all(|c| { (c.dtype() == &DataType::String) != c.name().starts_with("ns") })
517        );
518        let mut as_text = Vec::new();
519        JsonWriter::new(&mut as_text)
520            .with_json_format(JsonFormat::JsonLines)
521            .finish(&mut prepared)
522            .unwrap();
523        assert_eq!(
524            String::from_utf8(as_text).unwrap(),
525            String::from_utf8(json).unwrap()
526        );
527    }
528
529    /// The writers panic on a date past the calendar; as text it is its stored
530    /// number, alone or in a list, and the rest of its column is unchanged.
531    #[test]
532    fn a_date_past_the_calendar_is_written_as_its_stored_number() {
533        let df = calendar(true);
534        let fine = frame_as_json(&calendar(false)).unwrap();
535        let cells = frame_as_json(&df).unwrap();
536        assert!(cells.slice(0, fine.height()).equals_missing(&fine));
537        let text = |name: &str, row: usize| {
538            cells
539                .column(name)
540                .unwrap()
541                .str()
542                .unwrap()
543                .get(row)
544                .map(str::to_string)
545        };
546        let last = df.height() - 1;
547        assert_eq!(
548            text("d", last - 1).as_deref(),
549            Some("-2147483648 days since 1970-01-01")
550        );
551        assert_eq!(
552            text("ms_tz", last).as_deref(),
553            Some("9223372036854775807 ms since 1970-01-01 UTC")
554        );
555        assert_eq!(
556            text("us", last - 1).as_deref(),
557            Some("-9223372036854775807 us since 1970-01-01 UTC")
558        );
559        // Every nanosecond count is a date: the writer takes the column as it is.
560        assert_eq!(
561            cells.column("ns_tz").unwrap().dtype(),
562            df.column("ns_tz").unwrap().dtype()
563        );
564        let mut csv = Vec::new();
565        CsvWriter::new(&mut csv).finish(&mut cells.clone()).unwrap();
566        assert!(
567            String::from_utf8(csv)
568                .unwrap()
569                .contains(",2262-04-11T23:47:16.854775807,")
570        );
571        let lazy = lazy_as_json(df.clone().lazy()).unwrap().collect().unwrap();
572        assert!(cells.equals_missing(&lazy));
573
574        let mut json = Vec::new();
575        JsonWriter::new(&mut json)
576            .with_json_format(JsonFormat::JsonLines)
577            .finish(&mut lazy_for_json(df.clone().lazy()).unwrap().collect().unwrap())
578            .unwrap();
579        let json = String::from_utf8(json).unwrap();
580        assert!(
581            json.lines()
582                .last()
583                .unwrap()
584                .contains(r#""us_tz":"9223372036854775807 us since 1970-01-01 UTC""#),
585            "{json}"
586        );
587
588        let listed = df
589            .clone()
590            .lazy()
591            .select([col("us").implode(true), col("d").implode(true)])
592            .collect()
593            .unwrap();
594        let cells = frame_as_json(&listed).unwrap();
595        let us = cells
596            .column("us")
597            .unwrap()
598            .str()
599            .unwrap()
600            .get(0)
601            .unwrap()
602            .to_string();
603        assert!(
604            us.starts_with(r#"["1970-01-01 00:00:00",null,"#)
605                && us.ends_with(r#""-9223372036854775807 us since 1970-01-01 UTC","9223372036854775807 us since 1970-01-01 UTC"]"#),
606            "{us}"
607        );
608    }
609
610    fn nested() -> DataFrame {
611        let ids = Series::new("ids".into(), [1i64, 2, 3]);
612        let lists = Series::new(
613            "xs".into(),
614            [
615                Some(Series::new("".into(), ["a", "b"])),
616                None,
617                Some(Series::new("".into(), ["say \"hi\", then go"])),
618            ],
619        );
620        let point = StructChunked::from_series(
621            "point".into(),
622            3,
623            [
624                Series::new("x".into(), [1i64, 2, 3]),
625                Series::new("y".into(), [Some("a"), None, Some("c")]),
626            ]
627            .iter(),
628        )
629        .unwrap()
630        .into_series();
631        let pairs = Series::new(
632            "pair".into(),
633            [
634                Some(Series::new("".into(), [1.5f64, f64::NAN])),
635                Some(Series::new("".into(), [0.0f64, 2.0])),
636                None,
637            ],
638        )
639        .cast(&DataType::Array(Box::new(DataType::Float64), 2))
640        .unwrap();
641        DataFrame::new_infer_height(vec![ids.into(), lists.into(), point.into(), pairs.into()])
642            .unwrap()
643    }
644
645    #[test]
646    fn nested_columns_become_json_and_the_rest_stay() {
647        let df = frame_as_json(&nested()).unwrap();
648        assert_eq!(df.column("ids").unwrap().dtype(), &DataType::Int64);
649        let xs = df.column("xs").unwrap().str().unwrap().clone();
650        assert_eq!(xs.get(0), Some(r#"["a","b"]"#));
651        assert_eq!(xs.get(1), None, "a null list stays null");
652        assert_eq!(xs.get(2), Some(r#"["say \"hi\", then go"]"#));
653        let point = df.column("point").unwrap().str().unwrap().clone();
654        assert_eq!(point.get(0), Some(r#"{"x":1,"y":"a"}"#));
655        assert_eq!(point.get(1), Some(r#"{"x":2,"y":null}"#));
656        let pair = df.column("pair").unwrap().str().unwrap().clone();
657        assert_eq!(pair.get(0), Some("[1.5,null]"), "NaN has no JSON spelling");
658        assert_eq!(pair.get(1), Some("[0.0,2.0]"));
659        assert_eq!(pair.get(2), None);
660    }
661
662    #[test]
663    fn lazy_and_in_memory_agree() {
664        let eager = frame_as_json(&nested()).unwrap();
665        let lazy = lazy_as_json(nested().lazy()).unwrap().collect().unwrap();
666        assert!(eager.equals_missing(&lazy), "{eager}\n{lazy}");
667    }
668
669    /// Dates, datetimes and categoricals inside a struct or list read the same
670    /// in a CSV cell as in an NDJSON export of the same frame.
671    #[test]
672    fn cells_match_an_ndjson_export() {
673        let df = df!(
674            "d" => [Some("2024-01-02"), None],
675            "c" => [Some("a"), None],
676        )
677        .unwrap()
678        .lazy()
679        .with_columns([
680            col("d").str().to_date(StrptimeOptions::default()),
681            col("c").cast(DataType::from_categories(Categories::global())),
682        ])
683        .with_columns([col("d")
684            .cast(DataType::Datetime(
685                TimeUnit::Microseconds,
686                Some(TimeZone::UTC),
687            ))
688            .alias("dt")])
689        .select([
690            as_struct(vec![col("d"), col("dt"), col("c")]).alias("s"),
691            col("c").implode(true).alias("lc"),
692        ])
693        .collect()
694        .unwrap();
695        let mut ndjson = Vec::new();
696        JsonWriter::new(&mut ndjson)
697            .with_json_format(JsonFormat::JsonLines)
698            .finish(&mut df.clone())
699            .unwrap();
700        let cells = frame_as_json(&df).unwrap();
701        let (s, lc) = (
702            cells.column("s").unwrap().str().unwrap(),
703            cells.column("lc").unwrap().str().unwrap(),
704        );
705        let rebuilt: String = (0..cells.height())
706            .map(|i| {
707                format!(
708                    "{{\"s\":{},\"lc\":{}}}\n",
709                    s.get(i).unwrap(),
710                    lc.get(i).unwrap()
711                )
712            })
713            .collect();
714        assert_eq!(rebuilt, String::from_utf8(ndjson).unwrap());
715        assert!(rebuilt.contains(r#""dt":"2024-01-02T00:00:00+00:00","c":"a""#));
716    }
717
718    /// Binary is base64 at any depth, with its nulls, and the same in a copy,
719    /// a CSV cell and an NDJSON export.
720    #[test]
721    fn binary_is_base64_everywhere() {
722        let blob = Series::new("blob".into(), [Some(b"hi\xff".as_slice()), None]);
723        let blobs = Series::new(
724            "blobs".into(),
725            [Some(Series::new("".into(), [b"x".as_slice()])), None],
726        );
727        let pair = Series::new(
728            "pair".into(),
729            [
730                Some(Series::new("".into(), [b"a".as_slice(), b"b".as_slice()])),
731                None,
732            ],
733        )
734        .cast(&DataType::Array(Box::new(DataType::Binary), 2))
735        .unwrap();
736        let meta = StructChunked::from_series(
737            "meta".into(),
738            2,
739            [Series::new(
740                "raw".into(),
741                [b"ab".as_slice(), b"".as_slice()],
742            )]
743            .iter(),
744        )
745        .unwrap()
746        .with_outer_validity(Some([true, false].into_iter().collect()))
747        .into_series();
748        let df =
749            DataFrame::new_infer_height(vec![blob.into(), blobs.into(), pair.into(), meta.into()])
750                .unwrap();
751
752        for column in df.columns() {
753            let text = leaves_as_json_text(column.as_materialized_series()).unwrap();
754            assert_eq!(text.dtype(), &json_text_dtype(column.dtype()));
755            assert_eq!(text.null_count(), 1, "{text}");
756        }
757
758        let cells = frame_as_json(&df).unwrap();
759        let cell = |name: &str| cells.column(name).unwrap().str().unwrap().get(0);
760        assert_eq!(cell("blob"), Some("aGn/"));
761        assert_eq!(cell("blobs"), Some(r#"["eA=="]"#));
762        assert_eq!(cell("pair"), Some(r#"["YQ==","Yg=="]"#));
763        assert_eq!(cell("meta"), Some(r#"{"raw":"YWI="}"#));
764        let lazy = lazy_as_json(df.clone().lazy()).unwrap().collect().unwrap();
765        assert!(cells.equals_missing(&lazy), "{cells}\n{lazy}");
766
767        let mut ndjson = Vec::new();
768        JsonWriter::new(&mut ndjson)
769            .with_json_format(JsonFormat::JsonLines)
770            .finish(&mut lazy_for_json(df.clone().lazy()).unwrap().collect().unwrap())
771            .unwrap();
772        assert_eq!(
773            String::from_utf8(ndjson).unwrap().lines().next(),
774            Some(
775                r#"{"blob":"aGn/","blobs":["eA=="],"pair":["YQ==","Yg=="],"meta":{"raw":"YWI="}}"#
776            )
777        );
778
779        let copy =
780            crate::clipboard::tabular_payload(&df, crate::clipboard::CopyFormat::Tsv, true, true)
781                .unwrap();
782        let row: Vec<&str> = copy.text.lines().nth(1).unwrap().split('\t').collect();
783        assert_eq!(
784            row,
785            [
786                "aGn/",
787                r#""[""eA==""]""#,
788                r#""[""YQ=="",""Yg==""]""#,
789                r#""{""raw"":""YWI=""}""#
790            ]
791        );
792    }
793}