unnest_ndjson/
lib.rs

1//! Convert a large json document into smaller, easier to process documents, quickly.
2//!
3//! Call [unnest_to_ndjson] on your stream, and receive a much nicer stream, or some
4//! callbacks.
5
6use std::convert::TryFrom;
7use std::io;
8use std::io::Read;
9
10use iowrap::Ignore;
11use memchr::memchr;
12
13mod sink;
14mod source;
15
16pub use crate::sink::{MiniWrite, Sinker};
17use source::Source;
18
19/// Control what information is retained for individual result documents
20#[derive(Copy, Clone, Eq, PartialEq)]
21#[non_exhaustive]
22pub enum HeaderStyle {
23    /// No information is retained.
24    None,
25    /// The path to the child document is retained.
26    ///
27    /// `{"a": {"H": 6}, "b": {"H": 7}}` would become,
28    /// with the default formatter and a target of `1`,
29    /// `{"key":["a"],"value":{"H":6}}` and
30    /// `{"key":["b"],"value":{"H":6}}`
31    PathArray,
32}
33
34struct Loc {
35    depth: isize,
36    path: Vec<Vec<u8>>,
37    header_style: HeaderStyle,
38}
39
40impl Loc {
41    fn at_target(&self) -> bool {
42        0 == self.depth
43    }
44
45    fn collecting_keys(&self) -> bool {
46        self.depth <= 0
47    }
48
49    fn producing_regular_output(&self) -> bool {
50        self.depth > 0
51    }
52
53    fn shallower_than_target(&self) -> bool {
54        self.depth < 0
55    }
56
57    fn write_suffix(&self, into: &mut impl Sinker) -> io::Result<()> {
58        into.observe_end(self.header_style)
59    }
60
61    fn compute_header(&self) -> bool {
62        match self.header_style {
63            HeaderStyle::None => false,
64            HeaderStyle::PathArray => true,
65        }
66    }
67}
68
69/// Consume a large JSON document from a `Read`, and write sub documents to a destination.
70///
71/// The typical destination is just a `Write` implementation, like a [std::fs::File],
72/// or a [Vec]. Alternatively, you can use the [Sinker] interface to get access to fragments
73/// of documents.
74///
75/// Configure the level of un-nesting with the `target` parameter. `1` will remove one level
76/// of nesting, such as converting `[{"a":5}, {"a":6}]` into `{"a":5}` and `{"a":6}`.
77///
78/// `header_style` controls how much context to retain. See [HeaderStyle].
79pub fn unnest_to_ndjson<R: Read>(
80    from: R,
81    mut to: impl Sinker,
82    target: usize,
83    header_style: HeaderStyle,
84) -> io::Result<()> {
85    let mut iter = Source::new(from);
86    let depth = -isize::try_from(target).map_err(|_| io::ErrorKind::InvalidData)?;
87    let mut loc = Loc {
88        depth,
89        path: Vec::with_capacity(target),
90        header_style,
91    };
92    loop {
93        match drop_whitespace(&mut iter) {
94            Err(ref e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
95            Err(e) => Err(e)?,
96            Ok(()) => (),
97        }
98        handle_one(&mut iter, &mut to, &mut loc)?;
99    }
100    Ok(())
101}
102
103fn drop_whitespace<R: Read>(from: &mut Source<R>) -> io::Result<()> {
104    loop {
105        match from.buf().iter().position(|&b| !b.is_ascii_whitespace()) {
106            Some(end) => {
107                from.consume(end);
108                return Ok(());
109            }
110            None => {
111                from.all_useless();
112                from.fill()?;
113            }
114        }
115    }
116}
117
118fn handle_one<R: Read>(
119    from: &mut Source<R>,
120    into: &mut impl Sinker,
121    loc: &mut Loc,
122) -> io::Result<()> {
123    if loc.compute_header() && loc.at_target() {
124        into.observe_new_item(&loc.path, loc.header_style)?;
125    }
126    match from.next()? {
127        b'{' => handle_object(from, into, loc)?,
128        b'[' => handle_array(from, into, loc)?,
129        c => {
130            if loc.compute_header() && loc.shallower_than_target() {
131                into.observe_new_item(&loc.path, loc.header_style)?;
132            }
133            if b'"' == c {
134                parse_string(from, into)?;
135            } else {
136                scan_primitive(c, from, into)?
137            }
138            if loc.shallower_than_target() {
139                loc.write_suffix(into)?;
140            }
141        }
142    }
143    if loc.at_target() {
144        loc.write_suffix(into)?;
145    }
146    Ok(())
147}
148
149fn handle_object<R: Read>(
150    from: &mut Source<R>,
151    into: &mut impl Sinker,
152    loc: &mut Loc,
153) -> io::Result<()> {
154    loc.depth += 1;
155
156    if loc.producing_regular_output() {
157        into.write_all(b"{")?;
158    }
159    loop {
160        drop_whitespace(from)?;
161        let s = from.next()?;
162        match s {
163            b',' => continue,
164            b'"' => (),
165            b'}' => break,
166            _ => return Err(io::ErrorKind::InvalidData.into()),
167        }
168        if loc.producing_regular_output() {
169            parse_string(from, into)?;
170        } else {
171            assert!(loc.collecting_keys());
172            if loc.compute_header() {
173                let mut key = Vec::with_capacity(32);
174                parse_string(from, &mut key)?;
175                loc.path.push(key);
176            } else {
177                parse_string(from, &mut Ignore {})?;
178            }
179        }
180        drop_whitespace(from)?;
181        let colon = from.next()?;
182        if b':' != colon {
183            return Err(io::ErrorKind::InvalidData.into());
184        }
185        if loc.producing_regular_output() {
186            into.write_all(b":")?;
187        }
188        drop_whitespace(from)?;
189        handle_one(from, into, loc)?;
190        drop_whitespace(from)?;
191
192        if loc.compute_header() && loc.collecting_keys() {
193            let _ = loc.path.pop().unwrap();
194        }
195
196        let delim = from.next()?;
197        match delim {
198            b'}' => break,
199            b',' => (),
200            _ => return Err(io::ErrorKind::InvalidData.into()),
201        }
202        if loc.producing_regular_output() {
203            into.write_all(b",")?;
204        }
205    }
206    if loc.producing_regular_output() {
207        into.write_all(b"}")?;
208    }
209
210    loc.depth -= 1;
211
212    Ok(())
213}
214
215fn handle_array<R: Read>(
216    from: &mut Source<R>,
217    into: &mut impl Sinker,
218    loc: &mut Loc,
219) -> io::Result<()> {
220    loc.depth += 1;
221
222    if loc.producing_regular_output() {
223        into.write_all(b"[")?;
224    }
225
226    for idx in 0usize.. {
227        drop_whitespace(from)?;
228        if let Ok(b']') = from.peek() {
229            let _infallible = from.next()?;
230            break;
231        }
232
233        if loc.compute_header() && loc.collecting_keys() {
234            loc.path.push(format!("{}", idx).into_bytes());
235        }
236        handle_one(from, into, loc)?;
237        if loc.compute_header() && loc.collecting_keys() {
238            let _ = loc.path.pop().unwrap();
239        }
240
241        drop_whitespace(from)?;
242
243        let delim = from.next()?;
244        match delim {
245            b']' => break,
246            b',' => (),
247            _ => return Err(io::ErrorKind::InvalidData.into()),
248        }
249        if loc.producing_regular_output() {
250            into.write_all(b",")?;
251        }
252    }
253    if loc.producing_regular_output() {
254        into.write_all(b"]")?;
255    }
256
257    loc.depth -= 1;
258
259    Ok(())
260}
261
262fn scan_primitive<R: Read, W: sink::MiniWrite>(
263    start: u8,
264    from: &mut Source<R>,
265    into: &mut W,
266) -> io::Result<()> {
267    into.write_all(&[start])?;
268    while let Ok(b) = from.peek() {
269        if b.is_ascii_whitespace()
270            || b',' == b
271            || b']' == b
272            || b'}' == b
273            || b':' == b
274            || b.is_ascii_control()
275        {
276            break;
277        }
278        // infalliable, as we just peeked it
279        let b = from.next()?;
280        into.write_all(&[b])?;
281    }
282
283    Ok(())
284}
285
286fn parse_string<R: Read, W: sink::MiniWrite>(from: &mut Source<R>, into: &mut W) -> io::Result<()> {
287    into.write_all(b"\"")?;
288    loop {
289        let buf = from.buf();
290        let quote = memchr(b'"', buf).unwrap_or(buf.len());
291        let escape = memchr(b'\\', buf).unwrap_or(buf.len());
292        let safe = quote.min(escape);
293        into.write_all(&buf[..safe])?;
294        from.consume(safe);
295        let b = from.next()?;
296        match b {
297            b'"' => break,
298            b'\r' | b'\n' => return Err(io::ErrorKind::InvalidData.into()),
299            b'\\' => {
300                let e = from.next()?;
301                match e {
302                    b'"' | b'/' | b'\\' | b'b' | b'f' | b'r' | b'n' | b't' => {
303                        into.write_all(&[b'\\', e])?;
304                    }
305                    b'u' => {
306                        into.write_all(&[b'\\', b'u'])?;
307                        for _ in 0..4 {
308                            let h: u8 = from.next()?;
309                            if !h.is_ascii_hexdigit() {
310                                return Err(io::ErrorKind::InvalidData.into());
311                            }
312                            into.write_all(&[h])?
313                        }
314                    }
315                    _ => return Err(io::ErrorKind::InvalidData.into()),
316                }
317            }
318            o => into.write_all(&[o])?,
319        }
320    }
321    into.write_all(b"\"")?;
322    Ok(())
323}
324
325#[cfg(test)]
326mod tests {
327    use std::io;
328
329    use super::parse_string;
330    use super::Source;
331
332    fn ps(buf: &str) -> io::Result<String> {
333        let mut v = Vec::with_capacity(buf.len());
334        let mut buf = Source::new(io::Cursor::new(buf.as_bytes()));
335        // remove leading quote, as scan_one does
336        buf.next()?;
337        parse_string(&mut buf, &mut v)?;
338        Ok(String::from_utf8(v).unwrap())
339    }
340
341    #[test]
342    fn string() -> io::Result<()> {
343        assert_eq!(r#""hello world""#, ps(r#""hello world""#)?);
344        Ok(())
345    }
346}
unnest_ndjson/lib.rs

unnest_ndjson/
lib.rs
