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    /// Whether the stream has stopped: at the end, or after an error that
75    /// ends it (a read error, a line longer than [`MAX_LINE`]). A line that
76    /// is not JSON does not stop it; the next record follows.
77    pub fn stopped(&self) -> bool {
78        self.done
79    }
80}
81
82impl<R: BufRead> Iterator for Lines<R> {
83    type Item = Result<Json, DataError>;
84
85    fn next(&mut self) -> Option<Self::Item> {
86        while !self.done {
87            self.line += 1;
88            self.buffer.clear();
89            let read = (&mut self.reader)
90                .take(MAX_LINE as u64 + 1)
91                .read_until(b'\n', &mut self.buffer);
92            match read {
93                Err(error) => {
94                    self.done = true;
95                    return Some(Err(DataError::at(self.line, error.to_string())));
96                }
97                Ok(0) => self.done = true,
98                Ok(_) => {
99                    let mut line = self.buffer.as_slice();
100                    if line.len() > MAX_LINE && !line.ends_with(b"\n") {
101                        self.done = true;
102                        return Some(Err(DataError::at(
103                            self.line,
104                            format!("longer than {MAX_LINE} bytes"),
105                        )));
106                    }
107                    if self.line == 1 {
108                        line = line.strip_prefix("\u{feff}".as_bytes()).unwrap_or(line);
109                    }
110                    if line.iter().all(u8::is_ascii_whitespace) {
111                        continue;
112                    }
113                    return Some(
114                        serde_json::from_slice(line)
115                            .map_err(|e| DataError::at(self.line, format!("not JSON: {e}"))),
116                    );
117                }
118            }
119        }
120        None
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn errors_carry_the_line_number() {
130        let error = read("{\"a\": 1}\n\n{\"a\": \n").unwrap_err();
131        assert_eq!(error.line(), Some(3));
132        assert!(error.message().starts_with("not JSON"));
133    }
134
135    #[test]
136    fn a_stream_samples_ahead_then_yields_rows() {
137        let text = "\u{feff}{\"a\": 1}\r\n{\"a\": 2, \"b\": 3}\n[1, 2]\n";
138        let mut source = source_with_sample(text.as_bytes(), 1).unwrap();
139        assert_eq!(source.columns(), ["a"]);
140        let mut count = 0;
141        while let Some(row) = source.next_row() {
142            row.unwrap();
143            count += 1;
144        }
145        assert_eq!(count, 3);
146        assert_eq!(source.unknown_keys().collect::<Vec<_>>(), ["b", "value"]);
147    }
148
149    /// A record of many keys, at the top level or nested, is read in time
150    /// linear in its keys: each key used to be looked up by scanning every
151    /// column, so 40,000 keys (1.7 MB) took 16 s and a 64 MiB line days.
152    #[test]
153    fn a_record_of_many_keys_reads_in_linear_time() {
154        let keys = 200_000;
155        let object: String = (0..keys)
156            .map(|i| format!("\"k{i}\":{i}"))
157            .collect::<Vec<_>>()
158            .join(",");
159        let text = format!("{{{object}}}\n{{\"nested\":{{{object}}}}}\n{{{object}}}\n");
160        let started = std::time::Instant::now();
161        let mut source = source(text.as_bytes()).unwrap();
162        assert_eq!(source.columns().len(), keys + 1);
163        let mut rows = 0;
164        while let Some(row) = source.next_row() {
165            assert_eq!(row.unwrap().len(), keys + 1);
166            rows += 1;
167        }
168        assert_eq!(rows, 3);
169        // Every key is absent from the nested record, so nullable.
170        assert!(source.schema().unwrap().fields()[0].is_nullable());
171        assert!(
172            started.elapsed() < std::time::Duration::from_secs(30),
173            "took {:?}",
174            started.elapsed()
175        );
176    }
177
178    /// A stream whose every record past the sample has a new key keeps only
179    /// the first `MAX_UNKNOWN_KEYS` names: it used to keep them all, so the
180    /// memory grew with the input and `rich profile` printed every one.
181    #[test]
182    fn keys_after_the_sample_are_named_up_to_a_bound() {
183        let text: String = (0..MAX_UNKNOWN_KEYS + 50)
184            .map(|i| format!("{{\"id\": {i}, \"k{i:04}\": 1}}\n"))
185            .collect();
186        let mut source = source_with_sample(text.as_bytes(), 1).unwrap();
187        while let Some(row) = source.next_row() {
188            row.unwrap();
189        }
190        let names: Vec<&str> = source.unknown_keys().collect();
191        assert_eq!(names.len(), MAX_UNKNOWN_KEYS);
192        assert_eq!(names[0], "k0001");
193        assert!(source.more_unknown_keys());
194    }
195}
rich_data/jsonl.rs

rich_data/
jsonl.rs
