Re-implement infer_json_schema to make use of TapeDecoder, removing

the need to parse rows into `serde_json::Value`s first.

Author: Rafferty97

arrow-json/src/reader/schema.rs

@@ -22,8 +22,8 @@ use arrow_schema::{ArrowError, Schema};
 use bumpalo::Bump;
 use serde_json::Value;
 
-use self::infer::{EMPTY_OBJECT_TY, infer_json_type};
-use super::ValueIter;
+use super::tape::TapeDecoder;
+use infer::{ANY_TY, EMPTY_OBJECT_TY, InferredType, TapeValue, infer_json_type};
 
 mod infer;
 
@@ -100,12 +100,29 @@ pub fn infer_json_schema_from_seekable<R: BufRead + Seek>(
 /// file.seek(SeekFrom::Start(0)).unwrap();
 /// ```
 pub fn infer_json_schema<R: BufRead>(
-    reader: R,
+    mut reader: R,
     max_read_records: Option<usize>,
 ) -> Result<(Schema, usize), ArrowError> {
-    let mut values = ValueIter::new(reader, max_read_records);
-    let schema = infer_json_schema_from_iterator(&mut values)?;
-    Ok((schema, values.record_count()))
+    let arena = Bump::new();
+    let mut decoder = SchemaDecoder::new(max_read_records, &arena);
+
+    loop {
+        let buf = reader.fill_buf()?;
+        let read = buf.len();
+
+        if read == 0 {
+            break;
+        }
+
+        let decoded = decoder.decode(buf)?;
+        reader.consume(decoded);
+
+        if decoded != read {
+            break;
+        }
+    }
+
+    decoder.finish()
 }
 
 /// Infer the fields of a JSON file by reading all items from the JSON Value Iterator.
@@ -136,6 +153,60 @@ where
         .into_schema()
 }
 
+struct SchemaDecoder<'a> {
+    decoder: TapeDecoder,
+    max_read_records: Option<usize>,
+    record_count: usize,
+    schema: InferredType<'a>,
+    arena: &'a Bump,
+}
+
+impl<'a> SchemaDecoder<'a> {
+    pub fn new(max_read_records: Option<usize>, arena: &'a Bump) -> Self {
+        Self {
+            decoder: TapeDecoder::new(1024, 8),
+            max_read_records,
+            record_count: 0,
+            schema: ANY_TY,
+            arena,
+        }
+    }
+
+    pub fn decode(&mut self, buf: &[u8]) -> Result<usize, ArrowError> {
+        let read = self.decoder.decode(buf)?;
+        if read != buf.len() {
+            self.infer_batch()?;
+        }
+        Ok(read)
+    }
+
+    pub fn finish(mut self) -> Result<(Schema, usize), ArrowError> {
+        self.infer_batch()?;
+        Ok((self.schema.into_schema()?, self.record_count))
+    }
+
+    fn infer_batch(&mut self) -> Result<(), ArrowError> {
+        let tape = self.decoder.finish()?;
+
+        let remaining_records = self
+            .max_read_records
+            .map_or(usize::MAX, |max| max - self.record_count);
+
+        let records = tape
+            .iter_rows()
+            .map(|idx| TapeValue::new(&tape, idx))
+            .take(remaining_records);
+
+        for record in records {
+            self.schema = infer_json_type(record, self.schema, self.arena)?;
+            self.record_count += 1;
+        }
+
+        self.decoder.clear();
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -306,7 +377,7 @@ mod tests {
         let re = infer_json_schema_from_seekable(Cursor::new(b"}"), None);
         assert_eq!(
             re.err().unwrap().to_string(),
-            "Json error: Not valid JSON: expected value at line 1 column 1",
+            "Json error: Encountered unexpected '}' whilst parsing value"
         );
     }
 
@@ -320,14 +391,14 @@ mod tests {
         let (inferred_schema, _) =
             infer_json_schema_from_seekable(Cursor::new(data), None).expect("infer");
         let schema = Schema::new(vec![
-            Field::new("an", list_type_of(DataType::Null), true),
             Field::new("in", DataType::Int64, true),
-            Field::new("n", DataType::Null, true),
-            Field::new("na", list_type_of(DataType::Null), true),
-            Field::new("nas", list_type_of(DataType::Utf8), true),
             Field::new("ni", DataType::Int64, true),
             Field::new("ns", DataType::Utf8, true),
             Field::new("sn", DataType::Utf8, true),
+            Field::new("n", DataType::Null, true),
+            Field::new("an", list_type_of(DataType::Null), true),
+            Field::new("na", list_type_of(DataType::Null), true),
+            Field::new("nas", list_type_of(DataType::Utf8), true),
         ]);
         assert_eq!(inferred_schema, schema);
     }

arrow-json/src/reader/schema/infer.rs

@@ -20,6 +20,8 @@ use std::collections::HashMap;
 use arrow_schema::{ArrowError, DataType, Field, Fields, Schema};
 use bumpalo::Bump;
 
+use crate::reader::tape::{Tape, TapeElement};
+
 #[derive(Clone, Copy, Debug)]
 pub struct InferredType<'t>(&'t TyKind<'t>);
 
@@ -257,6 +259,54 @@ pub enum JsonType {
     Object,
 }
 
+#[derive(Copy, Clone, Debug)]
+pub struct TapeValue<'a> {
+    tape: &'a Tape<'a>,
+    idx: u32,
+}
+
+impl<'a> TapeValue<'a> {
+    pub fn new(tape: &'a Tape<'a>, idx: u32) -> Self {
+        Self { tape, idx }
+    }
+}
+
+impl<'a> JsonValue<'a> for TapeValue<'a> {
+    fn get(&self) -> JsonType {
+        match self.tape.get(self.idx) {
+            TapeElement::Null => JsonType::Null,
+            TapeElement::False => JsonType::Bool,
+            TapeElement::True => JsonType::Bool,
+            TapeElement::I64(_) | TapeElement::I32(_) => JsonType::Int64,
+            TapeElement::F64(_) | TapeElement::F32(_) => JsonType::Float64,
+            TapeElement::Number(s) => {
+                if self.tape.get_string(s).parse::<i64>().is_ok() {
+                    JsonType::Int64
+                } else {
+                    JsonType::Float64
+                }
+            }
+            TapeElement::String(_) => JsonType::String,
+            TapeElement::StartList(_) => JsonType::Array,
+            TapeElement::EndList(_) => unreachable!(),
+            TapeElement::StartObject(_) => JsonType::Object,
+            TapeElement::EndObject(_) => unreachable!(),
+        }
+    }
+
+    fn elements(&self) -> impl Iterator<Item = Self> {
+        self.tape
+            .iter_elements(self.idx)
+            .map(move |idx| Self { idx, ..*self })
+    }
+
+    fn fields(&self) -> impl Iterator<Item = (&'a str, Self)> {
+        self.tape
+            .iter_fields(self.idx)
+            .map(move |(key, idx)| (key, Self { idx, ..*self }))
+    }
+}
+
 impl<'a> JsonValue<'a> for &'a serde_json::Value {
     fn get(&self) -> JsonType {
         use serde_json::Value;

arrow-json/src/reader/tape.rs

@@ -142,6 +142,55 @@ impl<'a> Tape<'a> {
         self.num_rows
     }
 
+    /// Iterates over the rows of the tape
+    pub fn iter_rows(&self) -> impl Iterator<Item = u32> {
+        self.iter_values(0, self.elements.len() as u32)
+    }
+
+    /// Iterates over the elements of the array starting at `idx`
+    pub fn iter_elements(&self, idx: u32) -> impl Iterator<Item = u32> {
+        let end = match self.get(idx) {
+            TapeElement::StartList(end) => end,
+            elem => panic!("Expected the start of a list, found {:?}", elem),
+        };
+
+        self.iter_values(idx, end)
+    }
+
+    /// Iterates over the fields of the objected starting at `idx`
+    pub fn iter_fields(&self, idx: u32) -> impl Iterator<Item = (&'a str, u32)> {
+        let end = match self.get(idx) {
+            TapeElement::StartObject(end) => end,
+            elem => panic!("Expected the start of an object, found {:?}", elem),
+        };
+
+        let mut idx = idx + 1;
+
+        std::iter::from_fn(move || {
+            (idx < end).then(|| {
+                let key = match self.get(idx) {
+                    TapeElement::String(s) => self.get_string(s),
+                    elem => panic!("Expected a string, found {:?}", elem),
+                };
+                let value_idx = idx + 1;
+                idx = self.next(value_idx, "field value").unwrap();
+                (key, value_idx)
+            })
+        })
+    }
+
+    fn iter_values(&self, start: u32, end: u32) -> impl Iterator<Item = u32> {
+        let mut idx = start + 1;
+
+        std::iter::from_fn(move || {
+            (idx < end).then(|| {
+                let value_id = idx;
+                idx = self.next(idx, "value").unwrap();
+                value_id
+            })
+        })
+    }
+
     /// Serialize the tape element at index `idx` to `out` returning the next field index
     fn serialize(&self, out: &mut String, idx: u32) -> u32 {
         match self.get(idx) {