arrow_json/reader/
tape.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::reader::serializer::TapeSerializer;
19use arrow_schema::ArrowError;
20use memchr::memchr2;
21use serde_core::Serialize;
22use std::fmt::Write;
23
24/// We decode JSON to a flattened tape representation,
25/// allowing for efficient traversal of the JSON data
26///
27/// This approach is inspired by [simdjson]
28///
29/// Uses `u32` for offsets to ensure `TapeElement` is 64-bits. A future
30/// iteration may increase this to a custom `u56` type.
31///
32/// Numbers take more than one form, and matches must handle all of them. Parsing JSON
33/// text always yields [`Self::Number`], holding the value textually (read via
34/// [`Tape::get_string`]). Serializing Rust values (see
35/// [`Decoder::serialize`](super::Decoder::serialize)) yields [`Self::I32`],
36/// [`Self::I64`], [`Self::F32`] or [`Self::F64`] — 64-bit values spanning two
37/// elements, high bits first — or [`Self::Number`] for integers exceeding `i64`.
38///
39/// [simdjson]: https://github.com/simdjson/simdjson/blob/master/doc/tape.md
40#[derive(Debug, Copy, Clone, PartialEq, Eq)]
41#[non_exhaustive]
42pub enum TapeElement {
43    /// The start of an object, i.e. `{`
44    ///
45    /// Contains the offset of the corresponding [`Self::EndObject`]
46    StartObject(u32),
47    /// The end of an object, i.e. `}`
48    ///
49    /// Contains the offset of the corresponding [`Self::StartObject`]
50    EndObject(u32),
51    /// The start of a list , i.e. `[`
52    ///
53    /// Contains the offset of the corresponding [`Self::EndList`]
54    StartList(u32),
55    /// The end of a list , i.e. `]`
56    ///
57    /// Contains the offset of the corresponding [`Self::StartList`]
58    EndList(u32),
59    /// A string value
60    ///
61    /// Contains the offset into the [`Tape`] string data
62    String(u32),
63    /// A numeric value
64    ///
65    /// Contains the offset into the [`Tape`] string data
66    Number(u32),
67
68    /// The high bits of a i64
69    ///
70    /// Followed by [`Self::I32`] containing the low bits
71    I64(i32),
72
73    /// A 32-bit signed integer
74    ///
75    /// May be preceded by [`Self::I64`] containing high bits
76    I32(i32),
77
78    /// The high bits of a 64-bit float
79    ///
80    /// Followed by [`Self::F32`] containing the low bits
81    F64(u32),
82
83    /// A 32-bit float or the low-bits of a 64-bit float if preceded by [`Self::F64`]
84    F32(u32),
85
86    /// A true literal
87    True,
88    /// A false literal
89    False,
90    /// A null literal
91    Null,
92}
93
94/// A decoded JSON tape
95///
96/// String and numeric data is stored alongside an array of [`TapeElement`]
97///
98/// The first element is always [`TapeElement::Null`]
99///
100/// This approach to decoding JSON is inspired by [simdjson]
101///
102/// String data is copied into the tape with escapes resolved, so [`Tape::get_string`]
103/// borrows from the tape, not the input. A `Tape` is read, never constructed.
104///
105/// [simdjson]: https://github.com/simdjson/simdjson/blob/master/doc/tape.md
106#[derive(Debug)]
107pub struct Tape<'a> {
108    elements: &'a [TapeElement],
109    strings: &'a str,
110    string_offsets: &'a [usize],
111    num_rows: usize,
112}
113
114impl<'a> Tape<'a> {
115    /// Returns the string for the given string index
116    #[inline]
117    pub fn get_string(&self, idx: u32) -> &'a str {
118        let end_offset = self.string_offsets[idx as usize + 1];
119        let start_offset = self.string_offsets[idx as usize];
120        // SAFETY:
121        // Verified offsets
122        unsafe { self.strings.get_unchecked(start_offset..end_offset) }
123    }
124
125    /// Returns the tape element at `idx`
126    pub fn get(&self, idx: u32) -> TapeElement {
127        self.elements[idx as usize]
128    }
129
130    /// Returns the index of the next field at the same level as `cur_idx`
131    ///
132    /// Return an error if `cur_idx` is not the start of a field
133    pub fn next(&self, cur_idx: u32, expected: &str) -> Result<u32, ArrowError> {
134        match self.get(cur_idx) {
135            TapeElement::String(_)
136            | TapeElement::Number(_)
137            | TapeElement::True
138            | TapeElement::False
139            | TapeElement::Null
140            | TapeElement::I32(_)
141            | TapeElement::F32(_) => Ok(cur_idx + 1),
142            TapeElement::I64(_) | TapeElement::F64(_) => Ok(cur_idx + 2),
143            TapeElement::StartList(end_idx) => Ok(end_idx + 1),
144            TapeElement::StartObject(end_idx) => Ok(end_idx + 1),
145            TapeElement::EndObject(_) | TapeElement::EndList(_) => {
146                Err(self.error(cur_idx, expected))
147            }
148        }
149    }
150
151    /// Returns the number of rows
152    pub fn num_rows(&self) -> usize {
153        self.num_rows
154    }
155
156    /// Serialize the tape element at index `idx` to `out` returning the next field index
157    fn serialize(&self, out: &mut String, idx: u32) -> u32 {
158        match self.get(idx) {
159            TapeElement::StartObject(end) => {
160                out.push('{');
161                let mut cur_idx = idx + 1;
162                while cur_idx < end {
163                    cur_idx = self.serialize(out, cur_idx);
164                    out.push_str(": ");
165                    cur_idx = self.serialize(out, cur_idx);
166                }
167                out.push('}');
168                return end + 1;
169            }
170            TapeElement::EndObject(_) => out.push('}'),
171            TapeElement::StartList(end) => {
172                out.push('[');
173                let mut cur_idx = idx + 1;
174                while cur_idx < end {
175                    cur_idx = self.serialize(out, cur_idx);
176                    if cur_idx < end {
177                        out.push_str(", ");
178                    }
179                }
180                out.push(']');
181                return end + 1;
182            }
183            TapeElement::EndList(_) => out.push(']'),
184            TapeElement::String(s) => {
185                out.push('"');
186                out.push_str(self.get_string(s));
187                out.push('"')
188            }
189            TapeElement::Number(n) => out.push_str(self.get_string(n)),
190            TapeElement::True => out.push_str("true"),
191            TapeElement::False => out.push_str("false"),
192            TapeElement::Null => out.push_str("null"),
193            TapeElement::I64(high) => match self.get(idx + 1) {
194                TapeElement::I32(low) => {
195                    let val = ((high as i64) << 32) | (low as u32) as i64;
196                    let _ = write!(out, "{val}");
197                    return idx + 2;
198                }
199                _ => unreachable!(),
200            },
201            TapeElement::I32(val) => {
202                let _ = write!(out, "{val}");
203            }
204            TapeElement::F64(high) => match self.get(idx + 1) {
205                TapeElement::F32(low) => {
206                    let val = f64::from_bits(((high as u64) << 32) | low as u64);
207                    let _ = write!(out, "{val}");
208                    return idx + 2;
209                }
210                _ => unreachable!(),
211            },
212            TapeElement::F32(val) => {
213                let _ = write!(out, "{}", f32::from_bits(val));
214            }
215        }
216        idx + 1
217    }
218
219    /// Returns an error reading index `idx`
220    pub fn error(&self, idx: u32, expected: &str) -> ArrowError {
221        let mut out = String::with_capacity(64);
222        self.serialize(&mut out, idx);
223        ArrowError::JsonError(format!("expected {expected} got {out}"))
224    }
225}
226
227/// States based on <https://www.json.org/json-en.html>
228#[derive(Debug, Copy, Clone)]
229enum DecoderState {
230    /// Decoding a top-level list, where each element is
231    /// treated as an individual row
232    ///
233    /// This can only appear as the first element on the stack,
234    /// and it is valid to flush a batch when this is the only
235    /// state on the stack
236    TopLevelList,
237    /// Decoding an object
238    ///
239    /// Contains index of start [`TapeElement::StartObject`]
240    Object(u32),
241    /// Decoding a list
242    ///
243    /// Contains index of start [`TapeElement::StartList`]
244    List(u32),
245    String,
246    Value,
247    Number,
248    Colon,
249    Escape,
250    /// A unicode escape sequence,
251    ///
252    /// Consists of a `(low surrogate, high surrogate, decoded length)`
253    Unicode(u16, u16, u8),
254    /// A boolean or null literal
255    ///
256    /// Consists of `(literal, decoded length)`
257    Literal(Literal, u8),
258}
259
260impl DecoderState {
261    fn as_str(&self) -> &'static str {
262        match self {
263            DecoderState::TopLevelList => "top-level list",
264            DecoderState::Object(_) => "object",
265            DecoderState::List(_) => "list",
266            DecoderState::String => "string",
267            DecoderState::Value => "value",
268            DecoderState::Number => "number",
269            DecoderState::Colon => "colon",
270            DecoderState::Escape => "escape",
271            DecoderState::Unicode(_, _, _) => "unicode literal",
272            DecoderState::Literal(d, _) => d.as_str(),
273        }
274    }
275}
276
277#[derive(Debug, Copy, Clone)]
278enum Literal {
279    Null,
280    True,
281    False,
282}
283
284impl Literal {
285    fn element(&self) -> TapeElement {
286        match self {
287            Literal::Null => TapeElement::Null,
288            Literal::True => TapeElement::True,
289            Literal::False => TapeElement::False,
290        }
291    }
292
293    fn as_str(&self) -> &'static str {
294        match self {
295            Literal::Null => "null",
296            Literal::True => "true",
297            Literal::False => "false",
298        }
299    }
300
301    fn bytes(&self) -> &'static [u8] {
302        self.as_str().as_bytes()
303    }
304}
305
306/// Evaluates to the next element in the iterator or breaks the current loop
307macro_rules! next {
308    ($next:ident) => {
309        match $next.next() {
310            Some(b) => b,
311            None => break,
312        }
313    };
314}
315
316/// Implements a state machine for decoding JSON to a tape
317#[derive(Debug)]
318pub struct TapeDecoder {
319    /// The decoded elements
320    elements: Vec<TapeElement>,
321
322    /// The number of rows decoded, including any in progress if `!stack.is_empty()`
323    cur_row: usize,
324
325    /// Number of rows to read per batch
326    batch_size: usize,
327
328    /// Whether to flatten top-level arrays into the stream of JSON rows,
329    /// meaning that each of their elements will be treated as an individual row
330    ///
331    /// When `false` (the default), the entire top-level array will be treated as one row
332    flatten_top_level_arrays: bool,
333
334    /// A buffer of parsed string data
335    ///
336    /// Note: if part way through a record, i.e. `stack` is not empty,
337    /// this may contain truncated UTF-8 data
338    bytes: Vec<u8>,
339
340    /// Offsets into `data`
341    offsets: Vec<usize>,
342
343    /// A stack of [`DecoderState`]
344    stack: Vec<DecoderState>,
345}
346
347/// Configuration for a `TapeDecoder`.
348#[derive(Clone, Copy, Debug)]
349pub struct TapeDecoderOptions {
350    /// The batch size
351    pub batch_size: usize,
352    /// The estimated number of fields in each row
353    pub num_fields: usize,
354    /// Whether to flatten top-level arrays
355    pub flatten_top_level_arrays: bool,
356}
357
358impl TapeDecoder {
359    /// Create a new [`TapeDecoder`] with the provided options
360    pub fn new(options: TapeDecoderOptions) -> Self {
361        let TapeDecoderOptions {
362            batch_size,
363            num_fields,
364            flatten_top_level_arrays,
365        } = options;
366
367        let tokens_per_row = 2 + num_fields * 2;
368        let mut offsets = Vec::with_capacity(batch_size * (num_fields * 2) + 1);
369        offsets.push(0);
370
371        let mut elements = Vec::with_capacity(batch_size * tokens_per_row);
372        elements.push(TapeElement::Null);
373
374        Self {
375            offsets,
376            elements,
377            batch_size,
378            flatten_top_level_arrays,
379            cur_row: 0,
380            bytes: Vec::with_capacity(num_fields * 2 * 8),
381            stack: Vec::with_capacity(10),
382        }
383    }
384
385    pub fn decode(&mut self, buf: &[u8]) -> Result<usize, ArrowError> {
386        let mut iter = BufIter::new(buf);
387
388        while !iter.is_empty() {
389            let state = match self.stack.last_mut() {
390                Some(l) => l,
391                None => {
392                    iter.skip_whitespace();
393                    if self.cur_row >= self.batch_size {
394                        break;
395                    }
396
397                    match iter.peek() {
398                        Some(b'[') if self.flatten_top_level_arrays => {
399                            // Consume the `[` without writing it
400                            iter.next();
401                            self.stack.push(DecoderState::TopLevelList);
402                        }
403                        Some(_) => {
404                            // Start of row
405                            self.cur_row += 1;
406                            self.stack.push(DecoderState::Value);
407                        }
408                        None => break,
409                    }
410
411                    // There is now a top-most state to process
412                    self.stack.last_mut().unwrap()
413                }
414            };
415
416            match state {
417                // Decoding a top-level list
418                DecoderState::TopLevelList => {
419                    iter.advance_until(|b| !json_whitespace(b) && b != b',');
420                    if self.cur_row >= self.batch_size {
421                        break;
422                    }
423
424                    match iter.peek() {
425                        Some(b']') => {
426                            // Consume the `]` without writing it
427                            iter.next();
428                            self.stack.pop();
429                        }
430                        Some(_) => {
431                            // Start of row
432                            self.cur_row += 1;
433                            self.stack.push(DecoderState::Value);
434                        }
435                        None => break,
436                    }
437                }
438                // Decoding an object
439                DecoderState::Object(start_idx) => {
440                    iter.advance_until(|b| !json_whitespace(b) && b != b',');
441                    match next!(iter) {
442                        b'"' => {
443                            self.stack.push(DecoderState::Value);
444                            self.stack.push(DecoderState::Colon);
445                            self.stack.push(DecoderState::String);
446                        }
447                        b'}' => {
448                            let start_idx = *start_idx;
449                            let end_idx = self.elements.len() as u32;
450                            self.elements[start_idx as usize] = TapeElement::StartObject(end_idx);
451                            self.elements.push(TapeElement::EndObject(start_idx));
452                            self.stack.pop();
453                        }
454                        b => return Err(err(b, "parsing object")),
455                    }
456                }
457                // Decoding a list
458                DecoderState::List(start_idx) => {
459                    iter.advance_until(|b| !json_whitespace(b) && b != b',');
460                    match iter.peek() {
461                        Some(b']') => {
462                            iter.next();
463                            let start_idx = *start_idx;
464                            let end_idx = self.elements.len() as u32;
465                            self.elements[start_idx as usize] = TapeElement::StartList(end_idx);
466                            self.elements.push(TapeElement::EndList(start_idx));
467                            self.stack.pop();
468                        }
469                        Some(_) => self.stack.push(DecoderState::Value),
470                        None => break,
471                    }
472                }
473                // Decoding a string
474                DecoderState::String => {
475                    let s = iter.skip_chrs(b'\\', b'"');
476                    self.bytes.extend_from_slice(s);
477
478                    match next!(iter) {
479                        b'\\' => self.stack.push(DecoderState::Escape),
480                        b'"' => {
481                            let idx = self.offsets.len() - 1;
482                            self.elements.push(TapeElement::String(idx as _));
483                            self.offsets.push(self.bytes.len());
484                            self.stack.pop();
485                        }
486                        b => unreachable!("{}", b),
487                    }
488                }
489                state @ DecoderState::Value => {
490                    iter.skip_whitespace();
491                    *state = match next!(iter) {
492                        b'"' => DecoderState::String,
493                        b @ (b'-' | b'0'..=b'9') => {
494                            self.bytes.push(b);
495                            DecoderState::Number
496                        }
497                        b'n' => DecoderState::Literal(Literal::Null, 1),
498                        b'f' => DecoderState::Literal(Literal::False, 1),
499                        b't' => DecoderState::Literal(Literal::True, 1),
500                        b'[' => {
501                            let idx = self.elements.len() as u32;
502                            self.elements.push(TapeElement::StartList(u32::MAX));
503                            DecoderState::List(idx)
504                        }
505                        b'{' => {
506                            let idx = self.elements.len() as u32;
507                            self.elements.push(TapeElement::StartObject(u32::MAX));
508                            DecoderState::Object(idx)
509                        }
510                        b => return Err(err(b, "parsing value")),
511                    };
512                }
513                DecoderState::Number => {
514                    let s = iter.advance_until(|b| {
515                        !matches!(b, b'0'..=b'9' | b'-' | b'+' | b'.' | b'e' | b'E')
516                    });
517                    self.bytes.extend_from_slice(s);
518
519                    if !iter.is_empty() {
520                        self.stack.pop();
521                        let idx = self.offsets.len() - 1;
522                        self.elements.push(TapeElement::Number(idx as _));
523                        self.offsets.push(self.bytes.len());
524                    }
525                }
526                DecoderState::Colon => {
527                    iter.skip_whitespace();
528                    match next!(iter) {
529                        b':' => self.stack.pop(),
530                        b => return Err(err(b, "parsing colon")),
531                    };
532                }
533                DecoderState::Literal(literal, idx) => {
534                    let bytes = literal.bytes();
535                    let expected = bytes.iter().skip(*idx as usize).copied();
536                    for (expected, b) in expected.zip(&mut iter) {
537                        match b == expected {
538                            true => *idx += 1,
539                            false => return Err(err(b, "parsing literal")),
540                        }
541                    }
542                    if *idx == bytes.len() as u8 {
543                        let element = literal.element();
544                        self.stack.pop();
545                        self.elements.push(element);
546                    }
547                }
548                DecoderState::Escape => {
549                    let v = match next!(iter) {
550                        b'u' => {
551                            self.stack.pop();
552                            self.stack.push(DecoderState::Unicode(0, 0, 0));
553                            continue;
554                        }
555                        b'"' => b'"',
556                        b'\\' => b'\\',
557                        b'/' => b'/',
558                        b'b' => 8,  // BS
559                        b'f' => 12, // FF
560                        b'n' => b'\n',
561                        b'r' => b'\r',
562                        b't' => b'\t',
563                        b => return Err(err(b, "parsing escape sequence")),
564                    };
565
566                    self.stack.pop();
567                    self.bytes.push(v);
568                }
569                // Parse a unicode escape sequence
570                DecoderState::Unicode(high, low, idx) => loop {
571                    match *idx {
572                        0..=3 => *high = (*high << 4) | parse_hex(next!(iter))? as u16,
573                        4 => {
574                            if let Some(c) = char::from_u32(*high as u32) {
575                                write_char(c, &mut self.bytes);
576                                self.stack.pop();
577                                break;
578                            }
579
580                            match next!(iter) {
581                                b'\\' => {}
582                                b => return Err(err(b, "parsing surrogate pair escape")),
583                            }
584                        }
585                        5 => match next!(iter) {
586                            b'u' => {}
587                            b => return Err(err(b, "parsing surrogate pair unicode")),
588                        },
589                        6..=9 => *low = (*low << 4) | parse_hex(next!(iter))? as u16,
590                        _ => {
591                            let c = char_from_surrogate_pair(*low, *high)?;
592                            write_char(c, &mut self.bytes);
593                            self.stack.pop();
594                            break;
595                        }
596                    }
597                    *idx += 1;
598                },
599            }
600        }
601
602        Ok(buf.len() - iter.len())
603    }
604
605    /// Writes any type that implements [`Serialize`] into this [`TapeDecoder`]
606    pub fn serialize<S: Serialize>(&mut self, rows: &[S]) -> Result<(), ArrowError> {
607        if let Some(b) = self.stack.last() {
608            return Err(ArrowError::JsonError(format!(
609                "Cannot serialize to tape containing partial decode state {}",
610                b.as_str()
611            )));
612        }
613
614        let mut serializer =
615            TapeSerializer::new(&mut self.elements, &mut self.bytes, &mut self.offsets);
616
617        rows.iter()
618            .try_for_each(|row| row.serialize(&mut serializer))
619            .map_err(|e| ArrowError::JsonError(e.to_string()))?;
620
621        self.cur_row += rows.len();
622
623        Ok(())
624    }
625
626    /// The number of buffered rows, including the partially decoded row (if any).
627    pub fn num_buffered_rows(&self) -> usize {
628        self.cur_row
629    }
630
631    /// True if the decoder is part way through decoding a row. If so, calling [`Self::finish`]
632    /// would return an error.
633    pub fn has_partial_row(&self) -> bool {
634        !matches!(self.stack.last(), None | Some(DecoderState::TopLevelList))
635    }
636
637    /// Finishes the current [`Tape`]
638    pub fn finish(&self) -> Result<Tape<'_>, ArrowError> {
639        match self.stack.last() {
640            None => {}
641            Some(DecoderState::TopLevelList) => {}
642            Some(state) => {
643                return Err(ArrowError::JsonError(format!(
644                    "Truncated record whilst reading {}",
645                    state.as_str()
646                )));
647            }
648        }
649
650        if self.offsets.len() >= u32::MAX as usize {
651            return Err(ArrowError::JsonError(format!(
652                "Encountered more than {} bytes of string data, consider using a smaller batch size",
653                u32::MAX
654            )));
655        }
656
657        if self.offsets.len() >= u32::MAX as usize {
658            return Err(ArrowError::JsonError(format!(
659                "Encountered more than {} JSON elements, consider using a smaller batch size",
660                u32::MAX
661            )));
662        }
663
664        // Sanity check
665        assert_eq!(
666            self.offsets.last().copied().unwrap_or_default(),
667            self.bytes.len()
668        );
669
670        let strings = simdutf8::basic::from_utf8(&self.bytes)
671            .map_err(|_| ArrowError::JsonError("Encountered non-UTF-8 data".to_string()))?;
672
673        for offset in self.offsets.iter().copied() {
674            if !strings.is_char_boundary(offset) {
675                return Err(ArrowError::JsonError(
676                    "Encountered truncated UTF-8 sequence".to_string(),
677                ));
678            }
679        }
680
681        Ok(Tape {
682            strings,
683            elements: &self.elements,
684            string_offsets: &self.offsets,
685            num_rows: self.cur_row,
686        })
687    }
688
689    /// Clears this [`TapeDecoder`] in preparation to read the next batch
690    pub fn clear(&mut self) {
691        assert!(!self.has_partial_row());
692
693        self.cur_row = 0;
694        self.bytes.clear();
695        self.elements.clear();
696        self.elements.push(TapeElement::Null);
697        self.offsets.clear();
698        self.offsets.push(0);
699    }
700}
701
702/// A wrapper around a slice iterator that provides some helper functionality
703struct BufIter<'a> {
704    buf: &'a [u8],
705    pos: usize,
706}
707
708impl<'a> BufIter<'a> {
709    fn new(buf: &'a [u8]) -> Self {
710        Self { buf, pos: 0 }
711    }
712
713    #[inline]
714    fn as_slice(&self) -> &'a [u8] {
715        &self.buf[self.pos..]
716    }
717
718    #[inline]
719    fn is_empty(&self) -> bool {
720        self.pos >= self.buf.len()
721    }
722
723    fn peek(&self) -> Option<u8> {
724        self.buf.get(self.pos).copied()
725    }
726
727    #[inline]
728    fn advance(&mut self, skip: usize) {
729        self.pos += skip;
730    }
731
732    fn advance_until<F: FnMut(u8) -> bool>(&mut self, f: F) -> &[u8] {
733        let s = self.as_slice();
734        match s.iter().copied().position(f) {
735            Some(x) => {
736                self.advance(x);
737                &s[..x]
738            }
739            None => {
740                self.advance(s.len());
741                s
742            }
743        }
744    }
745
746    fn skip_chrs(&mut self, c1: u8, c2: u8) -> &[u8] {
747        let s = self.as_slice();
748        match memchr2(c1, c2, s) {
749            Some(p) => {
750                self.advance(p);
751                &s[..p]
752            }
753            None => {
754                self.advance(s.len());
755                s
756            }
757        }
758    }
759
760    fn skip_whitespace(&mut self) {
761        self.advance_until(|b| !json_whitespace(b));
762    }
763}
764
765impl Iterator for BufIter<'_> {
766    type Item = u8;
767
768    fn next(&mut self) -> Option<Self::Item> {
769        let b = self.peek();
770        self.pos += 1;
771        b
772    }
773
774    fn size_hint(&self) -> (usize, Option<usize>) {
775        let s = self.buf.len().saturating_sub(self.pos);
776        (s, Some(s))
777    }
778}
779
780impl ExactSizeIterator for BufIter<'_> {}
781
782/// Returns an error for a given byte `b` and context `ctx`
783fn err(b: u8, ctx: &str) -> ArrowError {
784    ArrowError::JsonError(format!(
785        "Encountered unexpected '{}' whilst {ctx}",
786        b as char
787    ))
788}
789
790/// Creates a character from an UTF-16 surrogate pair
791fn char_from_surrogate_pair(low: u16, high: u16) -> Result<char, ArrowError> {
792    match (low, high) {
793        (0xDC00..=0xDFFF, 0xD800..=0xDBFF) => {
794            let n = (((high - 0xD800) as u32) << 10) | ((low - 0xDC00) as u32 + 0x1_0000);
795            char::from_u32(n)
796                .ok_or_else(|| ArrowError::JsonError(format!("Invalid UTF-16 surrogate pair {n}")))
797        }
798        _ => Err(ArrowError::JsonError(format!(
799            "Invalid UTF-16 surrogate pair. High: {high:#04X}, Low: {low:#04X}"
800        ))),
801    }
802}
803
804/// Writes `c` as UTF-8 to `out`
805fn write_char(c: char, out: &mut Vec<u8>) {
806    let mut t = [0; 4];
807    out.extend_from_slice(c.encode_utf8(&mut t).as_bytes());
808}
809
810/// Evaluates to true if `b` is a valid JSON whitespace character
811#[inline]
812fn json_whitespace(b: u8) -> bool {
813    matches!(b, b' ' | b'\n' | b'\r' | b'\t')
814}
815
816/// Parse a hex character to `u8`
817fn parse_hex(b: u8) -> Result<u8, ArrowError> {
818    let digit = char::from(b)
819        .to_digit(16)
820        .ok_or_else(|| err(b, "unicode escape"))?;
821    Ok(digit as u8)
822}
823
824#[cfg(test)]
825mod tests {
826    use super::*;
827
828    /// Helper method to create a `TapeDecoder` with sensible defaults
829    fn tape_decoder() -> TapeDecoder {
830        TapeDecoder::new(TapeDecoderOptions {
831            batch_size: 16,
832            num_fields: 2,
833            flatten_top_level_arrays: false,
834        })
835    }
836
837    #[test]
838    fn test_sizes() {
839        assert_eq!(std::mem::size_of::<DecoderState>(), 8);
840        assert_eq!(std::mem::size_of::<TapeElement>(), 8);
841    }
842
843    #[test]
844    fn test_basic() {
845        let a = r#"
846        {"hello": "world", "foo": 2, "bar": 45}
847
848        {"foo": "bar"}
849
850        {"fiz": null}
851
852        {"a": true, "b": false, "c": null}
853
854        {"a": "", "": "a"}
855
856        {"a": "b", "object": {"nested": "hello", "foo": 23}, "b": {}, "c": {"foo": null }}
857
858        {"a": ["", "foo", ["bar", "c"]], "b": {"1": []}, "c": {"2": [1, 2, 3]} }
859        "#;
860        let mut decoder = tape_decoder();
861        decoder.decode(a.as_bytes()).unwrap();
862        assert!(!decoder.has_partial_row());
863        assert_eq!(decoder.num_buffered_rows(), 7);
864
865        let finished = decoder.finish().unwrap();
866        assert!(!decoder.has_partial_row());
867        assert_eq!(decoder.num_buffered_rows(), 7); // didn't call clear() yet
868        assert_eq!(
869            finished.elements,
870            &[
871                TapeElement::Null,
872                TapeElement::StartObject(8), // {"hello": "world", "foo": 2, "bar": 45}
873                TapeElement::String(0),      // "hello"
874                TapeElement::String(1),      // "world"
875                TapeElement::String(2),      // "foo"
876                TapeElement::Number(3),      // 2
877                TapeElement::String(4),      // "bar"
878                TapeElement::Number(5),      // 45
879                TapeElement::EndObject(1),
880                TapeElement::StartObject(12), // {"foo": "bar"}
881                TapeElement::String(6),       // "foo"
882                TapeElement::String(7),       // "bar"
883                TapeElement::EndObject(9),
884                TapeElement::StartObject(16), // {"fiz": null}
885                TapeElement::String(8),       // "fiz
886                TapeElement::Null,            // null
887                TapeElement::EndObject(13),
888                TapeElement::StartObject(24), // {"a": true, "b": false, "c": null}
889                TapeElement::String(9),       // "a"
890                TapeElement::True,            // true
891                TapeElement::String(10),      // "b"
892                TapeElement::False,           // false
893                TapeElement::String(11),      // "c"
894                TapeElement::Null,            // null
895                TapeElement::EndObject(17),
896                TapeElement::StartObject(30), // {"a": "", "": "a"}
897                TapeElement::String(12),      // "a"
898                TapeElement::String(13),      // ""
899                TapeElement::String(14),      // ""
900                TapeElement::String(15),      // "a"
901                TapeElement::EndObject(25),
902                TapeElement::StartObject(49), // {"a": "b", "object": {"nested": "hello", "foo": 23}, "b": {}, "c": {"foo": null }}
903                TapeElement::String(16),      // "a"
904                TapeElement::String(17),      // "b"
905                TapeElement::String(18),      // "object"
906                TapeElement::StartObject(40), // {"nested": "hello", "foo": 23}
907                TapeElement::String(19),      // "nested"
908                TapeElement::String(20),      // "hello"
909                TapeElement::String(21),      // "foo"
910                TapeElement::Number(22),      // 23
911                TapeElement::EndObject(35),
912                TapeElement::String(23),      // "b"
913                TapeElement::StartObject(43), // {}
914                TapeElement::EndObject(42),
915                TapeElement::String(24),      // "c"
916                TapeElement::StartObject(48), // {"foo": null }
917                TapeElement::String(25),      // "foo"
918                TapeElement::Null,            // null
919                TapeElement::EndObject(45),
920                TapeElement::EndObject(31),
921                TapeElement::StartObject(75), // {"a": ["", "foo", ["bar", "c"]], "b": {"1": []}, "c": {"2": [1, 2, 3]} }
922                TapeElement::String(26),      // "a"
923                TapeElement::StartList(59),   // ["", "foo", ["bar", "c"]]
924                TapeElement::String(27),      // ""
925                TapeElement::String(28),      // "foo"
926                TapeElement::StartList(58),   // ["bar", "c"]
927                TapeElement::String(29),      // "bar"
928                TapeElement::String(30),      // "c"
929                TapeElement::EndList(55),
930                TapeElement::EndList(52),
931                TapeElement::String(31),      // "b"
932                TapeElement::StartObject(65), // {"1": []}
933                TapeElement::String(32),      // "1"
934                TapeElement::StartList(64),   // []
935                TapeElement::EndList(63),
936                TapeElement::EndObject(61),
937                TapeElement::String(33),      // "c"
938                TapeElement::StartObject(74), // {"2": [1, 2, 3]}
939                TapeElement::String(34),      // "2"
940                TapeElement::StartList(73),   // [1, 2, 3]
941                TapeElement::Number(35),      // 1
942                TapeElement::Number(36),      // 2
943                TapeElement::Number(37),      // 3
944                TapeElement::EndList(69),
945                TapeElement::EndObject(67),
946                TapeElement::EndObject(50)
947            ]
948        );
949
950        assert_eq!(
951            finished.strings,
952            "helloworldfoo2bar45foobarfizabcaaabobjectnestedhellofoo23bcfooafoobarcb1c2123"
953        );
954        assert_eq!(
955            &finished.string_offsets,
956            &[
957                0, 5, 10, 13, 14, 17, 19, 22, 25, 28, 29, 30, 31, 32, 32, 32, 33, 34, 35, 41, 47,
958                52, 55, 57, 58, 59, 62, 63, 63, 66, 69, 70, 71, 72, 73, 74, 75, 76, 77
959            ]
960        );
961
962        decoder.clear();
963        assert!(!decoder.has_partial_row());
964        assert_eq!(decoder.num_buffered_rows(), 0);
965    }
966
967    #[test]
968    fn test_invalid() {
969        // Test invalid
970        let mut decoder = tape_decoder();
971        let err = decoder.decode(b"hello").unwrap_err().to_string();
972        assert_eq!(
973            err,
974            "Json error: Encountered unexpected 'h' whilst parsing value"
975        );
976
977        let mut decoder = tape_decoder();
978        let err = decoder.decode(b"{\"hello\": }").unwrap_err().to_string();
979        assert_eq!(
980            err,
981            "Json error: Encountered unexpected '}' whilst parsing value"
982        );
983
984        let mut decoder = tape_decoder();
985        let err = decoder
986            .decode(b"{\"hello\": [ false, tru ]}")
987            .unwrap_err()
988            .to_string();
989        assert_eq!(
990            err,
991            "Json error: Encountered unexpected ' ' whilst parsing literal"
992        );
993
994        let mut decoder = tape_decoder();
995        let err = decoder
996            .decode(b"{\"hello\": \"\\ud8\"}")
997            .unwrap_err()
998            .to_string();
999        assert_eq!(
1000            err,
1001            "Json error: Encountered unexpected '\"' whilst unicode escape"
1002        );
1003
1004        // Missing surrogate pair
1005        let mut decoder = tape_decoder();
1006        let err = decoder
1007            .decode(b"{\"hello\": \"\\ud83d\"}")
1008            .unwrap_err()
1009            .to_string();
1010        assert_eq!(
1011            err,
1012            "Json error: Encountered unexpected '\"' whilst parsing surrogate pair escape"
1013        );
1014
1015        // Test truncation
1016        let mut decoder = tape_decoder();
1017        decoder.decode(b"{\"he").unwrap();
1018        assert!(decoder.has_partial_row());
1019        assert_eq!(decoder.num_buffered_rows(), 1);
1020        let err = decoder.finish().unwrap_err().to_string();
1021        assert_eq!(err, "Json error: Truncated record whilst reading string");
1022
1023        let mut decoder = tape_decoder();
1024        decoder.decode(b"{\"hello\" : ").unwrap();
1025        let err = decoder.finish().unwrap_err().to_string();
1026        assert_eq!(err, "Json error: Truncated record whilst reading value");
1027
1028        let mut decoder = tape_decoder();
1029        decoder.decode(b"{\"hello\" : [").unwrap();
1030        let err = decoder.finish().unwrap_err().to_string();
1031        assert_eq!(err, "Json error: Truncated record whilst reading list");
1032
1033        let mut decoder = tape_decoder();
1034        decoder.decode(b"{\"hello\" : tru").unwrap();
1035        let err = decoder.finish().unwrap_err().to_string();
1036        assert_eq!(err, "Json error: Truncated record whilst reading true");
1037
1038        let mut decoder = tape_decoder();
1039        decoder.decode(b"{\"hello\" : nu").unwrap();
1040        let err = decoder.finish().unwrap_err().to_string();
1041        assert_eq!(err, "Json error: Truncated record whilst reading null");
1042
1043        // Test invalid UTF-8
1044        let mut decoder = tape_decoder();
1045        decoder.decode(b"{\"hello\" : \"world\xFF\"}").unwrap();
1046        let err = decoder.finish().unwrap_err().to_string();
1047        assert_eq!(err, "Json error: Encountered non-UTF-8 data");
1048
1049        let mut decoder = tape_decoder();
1050        decoder.decode(b"{\"\xe2\" : \"\x96\xa1\"}").unwrap();
1051        let err = decoder.finish().unwrap_err().to_string();
1052        assert_eq!(err, "Json error: Encountered truncated UTF-8 sequence");
1053    }
1054
1055    #[test]
1056    fn test_invalid_surrogates() {
1057        let mut decoder = tape_decoder();
1058        let res = decoder.decode(b"{\"test\": \"\\ud800\\ud801\"}");
1059        assert!(res.is_err());
1060
1061        let mut decoder = tape_decoder();
1062        let res = decoder.decode(b"{\"test\": \"\\udc00\\udc01\"}");
1063        assert!(res.is_err());
1064    }
1065
1066    #[test]
1067    fn test_flatten_top_level_arrays() {
1068        let input = r#"
1069        [
1070            {"hello": "world", "foo": 2, "bar": 45},
1071            {"a": true, "b": false, "c": null}
1072        ]
1073        [
1074            {"a": "b", "object": {"nested": "hello", "foo": 23}},
1075            {"a": ["", "foo", ["bar", "c"]]},
1076            {"hello": "world", "foo": 2, "bar": 27}
1077        ]"#;
1078        const TOTAL_ROWS: usize = 5;
1079
1080        // Check that regular decoding returns two rows
1081        let mut decoder = tape_decoder();
1082        decoder.decode(input.as_bytes()).unwrap();
1083        assert!(!decoder.has_partial_row());
1084        assert_eq!(decoder.num_buffered_rows(), 2);
1085
1086        let expected = TestCase::new()
1087            // {"hello": "world", "foo": 2, "bar": 45}
1088            .start_object(6)
1089            .string("hello")
1090            .string("world")
1091            .string("foo")
1092            .number("2")
1093            .string("bar")
1094            .number("45")
1095            .end_object(6)
1096            // {"a": true, "b": false, "c": null}
1097            .start_object(6)
1098            .string("a")
1099            .r#true()
1100            .string("b")
1101            .r#false()
1102            .string("c")
1103            .null()
1104            .end_object(6)
1105            // {"a": "b", "object": {"nested": "hello", "foo": 23}}
1106            .start_object(9)
1107            .string("a")
1108            .string("b")
1109            .string("object")
1110            .start_object(4)
1111            .string("nested")
1112            .string("hello")
1113            .string("foo")
1114            .number("23")
1115            .end_object(4)
1116            .end_object(9)
1117            // {"a": ["", "foo", ["bar", "c"]]}
1118            .start_object(9)
1119            .string("a")
1120            .start_list(6)
1121            .string("")
1122            .string("foo")
1123            .start_list(2)
1124            .string("bar")
1125            .string("c")
1126            .end_list(2)
1127            .end_list(6)
1128            .end_object(9)
1129            // {"hello": "world", "foo": 2, "bar": 27}
1130            .start_object(6)
1131            .string("hello")
1132            .string("world")
1133            .string("foo")
1134            .number("2")
1135            .string("bar")
1136            .number("27")
1137            .end_object(6);
1138
1139        // Check that decoding with `flatten_top_level_arrays` yields rows correctly,
1140        // and respects the configured batch size
1141        for batch_size in [1, 2, 3, 4, 8] {
1142            let mut decoder = TapeDecoder::new(TapeDecoderOptions {
1143                batch_size,
1144                num_fields: 2,
1145                flatten_top_level_arrays: true,
1146            });
1147            decoder.decode(input.as_bytes()).unwrap();
1148            assert!(!decoder.has_partial_row());
1149            assert_eq!(decoder.num_buffered_rows(), batch_size.min(TOTAL_ROWS));
1150
1151            let finished = decoder.finish().unwrap();
1152            assert!(!decoder.has_partial_row());
1153            assert_eq!(decoder.num_buffered_rows(), batch_size.min(TOTAL_ROWS)); // didn't call clear() yet
1154            assert_eq!(
1155                finished.elements,
1156                &expected.elements[..finished.elements.len()]
1157            );
1158            assert_eq!(
1159                finished.strings,
1160                &expected.strings[..finished.strings.len()]
1161            );
1162            assert_eq!(
1163                finished.string_offsets,
1164                &expected.string_offsets[..finished.string_offsets.len()]
1165            );
1166
1167            decoder.clear();
1168            assert!(!decoder.has_partial_row());
1169            assert_eq!(decoder.num_buffered_rows(), 0);
1170        }
1171    }
1172
1173    /// The expected elements, strings and string offsets for a test case
1174    struct TestCase {
1175        elements: Vec<TapeElement>,
1176        strings: String,
1177        string_offsets: Vec<usize>,
1178    }
1179
1180    impl TestCase {
1181        fn new() -> Self {
1182            Self {
1183                elements: vec![TapeElement::Null],
1184                strings: String::new(),
1185                string_offsets: vec![0],
1186            }
1187        }
1188
1189        fn start_object(mut self, len: usize) -> Self {
1190            let end_idx = (self.elements.len() + len + 1) as u32;
1191            self.elements.push(TapeElement::StartObject(end_idx));
1192            self
1193        }
1194
1195        fn end_object(mut self, len: usize) -> Self {
1196            let start_idx = (self.elements.len() - len - 1) as u32;
1197            self.elements.push(TapeElement::EndObject(start_idx));
1198            self
1199        }
1200
1201        fn start_list(mut self, len: usize) -> Self {
1202            let end_idx = (self.elements.len() + len + 1) as u32;
1203            self.elements.push(TapeElement::StartList(end_idx));
1204            self
1205        }
1206
1207        fn end_list(mut self, len: usize) -> Self {
1208            let start_idx = (self.elements.len() - len - 1) as u32;
1209            self.elements.push(TapeElement::EndList(start_idx));
1210            self
1211        }
1212
1213        fn string(mut self, raw: &str) -> Self {
1214            let idx = (self.string_offsets.len() - 1) as u32;
1215            let start = self.strings.len();
1216            let end = start + raw.len();
1217            self.elements.push(TapeElement::String(idx));
1218            self.strings.push_str(raw);
1219            self.string_offsets.push(end);
1220            self
1221        }
1222
1223        fn number(mut self, raw: &str) -> Self {
1224            let idx = (self.string_offsets.len() - 1) as u32;
1225            let start = self.strings.len();
1226            let end = start + raw.len();
1227            self.elements.push(TapeElement::Number(idx));
1228            self.strings.push_str(raw);
1229            self.string_offsets.push(end);
1230            self
1231        }
1232
1233        fn r#true(mut self) -> Self {
1234            self.elements.push(TapeElement::True);
1235            self
1236        }
1237
1238        fn r#false(mut self) -> Self {
1239            self.elements.push(TapeElement::False);
1240            self
1241        }
1242
1243        fn null(mut self) -> Self {
1244            self.elements.push(TapeElement::Null);
1245            self
1246        }
1247    }
1248}
arrow_json/reader/tape.rs

arrow_json/reader/
tape.rs
