use std::collections::BTreeMap; use std::fmt; use std::io::{Cursor, Write}; use byteorder::WriteBytesExt; use serde::de::{DeserializeSeed, MapAccess, SeqAccess, Visitor}; use serde::Deserialize; use serde_json::Value; use super::{ByteCounter, DocumentsBatchIndex, Error}; use crate::FieldId; macro_rules! tri { ($e:expr) => { match $e { Ok(r) => r, Err(e) => return Ok(Err(e.into())), } }; } struct FieldIdResolver<'a>(&'a mut DocumentsBatchIndex); impl<'a, 'de> DeserializeSeed<'de> for FieldIdResolver<'a> { type Value = FieldId; fn deserialize(self, deserializer: D) -> Result where D: serde::Deserializer<'de>, { deserializer.deserialize_str(self) } } impl<'a, 'de> Visitor<'de> for FieldIdResolver<'a> { type Value = FieldId; fn visit_str(self, v: &str) -> Result where E: serde::de::Error, { Ok(self.0.insert(v)) } fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "a string") } } struct ValueDeserializer; impl<'de> DeserializeSeed<'de> for ValueDeserializer { type Value = serde_json::Value; fn deserialize(self, deserializer: D) -> Result where D: serde::Deserializer<'de>, { serde_json::Value::deserialize(deserializer) } } pub struct DocumentVisitor<'a, W> { pub inner: &'a mut ByteCounter, pub index: &'a mut DocumentsBatchIndex, pub obkv_buffer: &'a mut Vec, pub value_buffer: &'a mut Vec, pub values: &'a mut BTreeMap, pub count: &'a mut usize, } impl<'a, 'de, W: Write> Visitor<'de> for &mut DocumentVisitor<'a, W> { /// This Visitor value is nothing, since it write the value to a file. type Value = Result<(), Error>; fn visit_seq(self, mut seq: A) -> Result where A: SeqAccess<'de>, { while let Some(v) = seq.next_element_seed(&mut *self)? { tri!(v) } Ok(Ok(())) } fn visit_map(self, mut map: A) -> Result where A: MapAccess<'de>, { while let Some((key, value)) = map.next_entry_seed(FieldIdResolver(&mut *self.index), ValueDeserializer)? { self.values.insert(key, value); } self.obkv_buffer.clear(); let mut obkv = obkv::KvWriter::new(Cursor::new(&mut *self.obkv_buffer)); for (key, value) in self.values.iter() { self.value_buffer.clear(); // This is guaranteed to work tri!(serde_json::to_writer(Cursor::new(&mut *self.value_buffer), value)); tri!(obkv.insert(*key, &self.value_buffer)); } let reader = tri!(obkv.into_inner()).into_inner(); tri!(self.inner.write_u32::(reader.len() as u32)); tri!(self.inner.write_all(reader)); *self.count += 1; self.values.clear(); Ok(Ok(())) } fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "a documents, or a sequence of documents.") } } impl<'a, 'de, W> DeserializeSeed<'de> for &mut DocumentVisitor<'a, W> where W: Write, { type Value = Result<(), Error>; fn deserialize(self, deserializer: D) -> Result where D: serde::Deserializer<'de>, { deserializer.deserialize_map(self) } }