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

arrow_json/reader/
tape.rs
