rich_data/
jsonl.rs

1//! JSON Lines: one JSON value per line, as rows.
2//!
3//! Each non-blank line is a record; an object's keys are columns, and any
4//! other value is one `value` column. [`read`] reads text whole, so every
5//! record decides the columns and schema; [`source`] streams from a reader,
6//! deciding them from the first [`DEFAULT_SAMPLE`] records (see
7//! [`RecordSource`] for what happens to keys seen later). A line longer than
8//! [`MAX_LINE`] bytes is refused rather than read into memory.
9//!
10//! ```
11//! use rich_data::{jsonl, RowSource, Value};
12//!
13//! let rows = jsonl::read("{\"name\": \"web\", \"up\": true}\n\n{\"name\": \"db\", \"p99\": 35}\n").unwrap();
14//! assert_eq!(rows.columns(), ["name", "up", "p99"]);
15//! assert_eq!(rows.rows()[1], [Value::from("db"), Value::Null, Value::Int(35)]);
16//! let schema = rows.schema().unwrap();
17//! assert_eq!(schema.field("up").unwrap().data_type().to_string(), "boolean");
18//! assert!(schema.field("p99").unwrap().is_nullable());
19//! ```
20
21use std::io::{BufRead, Read};
22
23use serde_json::Value as Json;
24
25use crate::{DataError, RecordSource, RowSource, Rows};
26
27pub use crate::record::{DEFAULT_SAMPLE, MAX_DEPTH, MAX_UNKNOWN_KEYS};
28
29/// The longest line read, in bytes (64 MiB).
30pub const MAX_LINE: usize = 64 << 20;
31
32/// Read JSON Lines text, sampling every record for the columns and schema.
33pub fn read(text: &str) -> Result<Rows, DataError> {
34    RecordSource::new(Lines::new(text.as_bytes()), usize::MAX)?.collect_rows()
35}
36
37/// Stream JSON Lines from `reader`, sampling the first [`DEFAULT_SAMPLE`]
38/// records for the columns and schema.
39pub fn source<R: BufRead>(reader: R) -> Result<JsonlSource<R>, DataError> {
40    source_with_sample(reader, DEFAULT_SAMPLE)
41}
42
43/// Stream JSON Lines from `reader`, sampling the first `sample` records.
44pub fn source_with_sample<R: BufRead>(
45    reader: R,
46    sample: usize,
47) -> Result<JsonlSource<R>, DataError> {
48    RecordSource::new(Lines::new(reader), sample)
49}
50
51/// A streaming JSON Lines [`RowSource`].
52pub type JsonlSource<R> = RecordSource<Lines<R>>;
53
54/// The records of a JSON Lines stream, with their line numbers in errors.
55#[derive(Debug)]
56pub struct Lines<R> {
57    reader: R,
58    line: usize,
59    buffer: Vec<u8>,
60    done: bool,
61}
62
63impl<R: BufRead> Lines<R> {
64    /// Records from `reader`.
65    pub fn new(reader: R) -> Self {
66        Lines {
67            reader,
68            line: 0,
69            buffer: Vec::new(),
70            done: false,
71        }
72    }
73}
74
75impl<R: BufRead> Iterator for Lines<R> {
76    type Item = Result<Json, DataError>;
77
78    fn next(&mut self) -> Option<Self::Item> {
79        while !self.done {
80            self.line += 1;
81            self.buffer.clear();
82            let read = (&mut self.reader)
83                .take(MAX_LINE as u64 + 1)
84                .read_until(b'\n', &mut self.buffer);
85            match read {
86                Err(error) => {
87                    self.done = true;
88                    return Some(Err(DataError::at(self.line, error.to_string())));
89                }
90                Ok(0) => self.done = true,
91                Ok(_) => {
92                    let mut line = self.buffer.as_slice();
93                    if line.len() > MAX_LINE && !line.ends_with(b"\n") {
94                        self.done = true;
95                        return Some(Err(DataError::at(
96                            self.line,
97                            format!("longer than {MAX_LINE} bytes"),
98                        )));
99                    }
100                    if self.line == 1 {
101                        line = line.strip_prefix("\u{feff}".as_bytes()).unwrap_or(line);
102                    }
103                    if line.iter().all(u8::is_ascii_whitespace) {
104                        continue;
105                    }
106                    return Some(
107                        serde_json::from_slice(line)
108                            .map_err(|e| DataError::at(self.line, format!("not JSON: {e}"))),
109                    );
110                }
111            }
112        }
113        None
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn errors_carry_the_line_number() {
123        let error = read("{\"a\": 1}\n\n{\"a\": \n").unwrap_err();
124        assert_eq!(error.line(), Some(3));
125        assert!(error.message().starts_with("not JSON"));
126    }
127
128    #[test]
129    fn a_stream_samples_ahead_then_yields_rows() {
130        let text = "\u{feff}{\"a\": 1}\r\n{\"a\": 2, \"b\": 3}\n[1, 2]\n";
131        let mut source = source_with_sample(text.as_bytes(), 1).unwrap();
132        assert_eq!(source.columns(), ["a"]);
133        let mut count = 0;
134        while let Some(row) = source.next_row() {
135            row.unwrap();
136            count += 1;
137        }
138        assert_eq!(count, 3);
139        assert_eq!(source.unknown_keys().collect::<Vec<_>>(), ["b", "value"]);
140    }
141
142    /// A record of many keys, at the top level or nested, is read in time
143    /// linear in its keys: each key used to be looked up by scanning every
144    /// column, so 40,000 keys (1.7 MB) took 16 s and a 64 MiB line days.
145    #[test]
146    fn a_record_of_many_keys_reads_in_linear_time() {
147        let keys = 200_000;
148        let object: String = (0..keys)
149            .map(|i| format!("\"k{i}\":{i}"))
150            .collect::<Vec<_>>()
151            .join(",");
152        let text = format!("{{{object}}}\n{{\"nested\":{{{object}}}}}\n{{{object}}}\n");
153        let started = std::time::Instant::now();
154        let mut source = source(text.as_bytes()).unwrap();
155        assert_eq!(source.columns().len(), keys + 1);
156        let mut rows = 0;
157        while let Some(row) = source.next_row() {
158            assert_eq!(row.unwrap().len(), keys + 1);
159            rows += 1;
160        }
161        assert_eq!(rows, 3);
162        // Every key is absent from the nested record, so nullable.
163        assert!(source.schema().unwrap().fields()[0].is_nullable());
164        assert!(
165            started.elapsed() < std::time::Duration::from_secs(30),
166            "took {:?}",
167            started.elapsed()
168        );
169    }
170
171    /// A stream whose every record past the sample has a new key keeps only
172    /// the first `MAX_UNKNOWN_KEYS` names: it used to keep them all, so the
173    /// memory grew with the input and `rich profile` printed every one.
174    #[test]
175    fn keys_after_the_sample_are_named_up_to_a_bound() {
176        let text: String = (0..MAX_UNKNOWN_KEYS + 50)
177            .map(|i| format!("{{\"id\": {i}, \"k{i:04}\": 1}}\n"))
178            .collect();
179        let mut source = source_with_sample(text.as_bytes(), 1).unwrap();
180        while let Some(row) = source.next_row() {
181            row.unwrap();
182        }
183        let names: Vec<&str> = source.unknown_keys().collect();
184        assert_eq!(names.len(), MAX_UNKNOWN_KEYS);
185        assert_eq!(names[0], "k0001");
186        assert!(source.more_unknown_keys());
187    }
188}
rich_data/jsonl.rs

rich_data/
jsonl.rs
