mirror of
https://github.com/meilisearch/MeiliSearch
synced 2025-07-04 04:17:10 +02:00
Merge #402
402: Optimize document transform r=MarinPostma a=MarinPostma This pr optimizes the transform of documents additions in the obkv format. Instead on accepting any serializable objects, we instead treat json and CSV specifically: - For json, we build a serde `Visitor`, that transform the json straight into obkv without intermediate representation. - For csv, we directly write the lines in the obkv, applying other optimization as well. Co-authored-by: marin postma <postma.marin@protonmail.com>
This commit is contained in:
bors[bot]
GitHub
commit
d7943fe225
14 changed files with 706 additions and 818 deletions
|
|
@ -1,285 +0,0 @@
|
|||
use std::io::{Read, Result as IoResult};
|
||||
use std::num::ParseFloatError;
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
enum AllowedType {
|
||||
String,
|
||||
Number,
|
||||
}
|
||||
|
||||
fn parse_csv_header(header: &str) -> (String, AllowedType) {
|
||||
// if there are several separators we only split on the last one.
|
||||
match header.rsplit_once(':') {
|
||||
Some((field_name, field_type)) => match field_type {
|
||||
"string" => (field_name.to_string(), AllowedType::String),
|
||||
"number" => (field_name.to_string(), AllowedType::Number),
|
||||
// we may return an error in this case.
|
||||
_otherwise => (header.to_string(), AllowedType::String),
|
||||
},
|
||||
None => (header.to_string(), AllowedType::String),
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CSVDocumentDeserializer<R>
|
||||
where
|
||||
R: Read,
|
||||
{
|
||||
documents: csv::StringRecordsIntoIter<R>,
|
||||
headers: Vec<(String, AllowedType)>,
|
||||
}
|
||||
|
||||
impl<R: Read> CSVDocumentDeserializer<R> {
|
||||
pub fn from_reader(reader: R) -> IoResult<Self> {
|
||||
let mut records = csv::Reader::from_reader(reader);
|
||||
|
||||
let headers = records.headers()?.into_iter().map(parse_csv_header).collect();
|
||||
|
||||
Ok(Self { documents: records.into_records(), headers })
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: Read> Iterator for CSVDocumentDeserializer<R> {
|
||||
type Item = anyhow::Result<Map<String, Value>>;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
let csv_document = self.documents.next()?;
|
||||
|
||||
match csv_document {
|
||||
Ok(csv_document) => {
|
||||
let mut document = Map::new();
|
||||
|
||||
for ((field_name, field_type), value) in
|
||||
self.headers.iter().zip(csv_document.into_iter())
|
||||
{
|
||||
let parsed_value: Result<Value, ParseFloatError> = match field_type {
|
||||
AllowedType::Number => {
|
||||
value.parse::<f64>().map(Value::from).map_err(Into::into)
|
||||
}
|
||||
AllowedType::String => Ok(Value::String(value.to_string())),
|
||||
};
|
||||
|
||||
match parsed_value {
|
||||
Ok(value) => drop(document.insert(field_name.to_string(), value)),
|
||||
Err(_e) => {
|
||||
return Some(Err(anyhow::anyhow!(
|
||||
"Value '{}' is not a valid number",
|
||||
value
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(Ok(document))
|
||||
}
|
||||
Err(e) => Some(Err(anyhow::anyhow!("Error parsing csv document: {}", e))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn simple_csv_document() {
|
||||
let documents = r#"city,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city": "Boston",
|
||||
"country": "United States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coma_in_field() {
|
||||
let documents = r#"city,country,pop
|
||||
"Boston","United, States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city": "Boston",
|
||||
"country": "United, States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quote_in_field() {
|
||||
let documents = r#"city,country,pop
|
||||
"Boston","United"" States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city": "Boston",
|
||||
"country": "United\" States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn integer_in_field() {
|
||||
let documents = r#"city,country,pop:number
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city": "Boston",
|
||||
"country": "United States",
|
||||
"pop": 4628910.0,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn float_in_field() {
|
||||
let documents = r#"city,country,pop:number
|
||||
"Boston","United States","4628910.01""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city": "Boston",
|
||||
"country": "United States",
|
||||
"pop": 4628910.01,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn several_double_dot_in_header() {
|
||||
let documents = r#"city:love:string,country:state,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city:love": "Boston",
|
||||
"country:state": "United States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ending_by_double_dot_in_header() {
|
||||
let documents = r#"city:,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city:": "Boston",
|
||||
"country": "United States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn starting_by_double_dot_in_header() {
|
||||
let documents = r#":city,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
":city": "Boston",
|
||||
"country": "United States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn starting_by_double_dot_in_header2() {
|
||||
let documents = r#":string,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"": "Boston",
|
||||
"country": "United States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn double_double_dot_in_header() {
|
||||
let documents = r#"city::string,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
"city:": "Boston",
|
||||
"country": "United States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_type_in_header() {
|
||||
let documents = r#"city,country:number,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert!(csv_iter.next().unwrap().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_column_count1() {
|
||||
let documents = r#"city,country,pop
|
||||
"Boston","United States","4628910", "too much""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert!(csv_iter.next().unwrap().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_column_count2() {
|
||||
let documents = r#"city,country,pop
|
||||
"Boston","United States""#;
|
||||
|
||||
let mut csv_iter = CSVDocumentDeserializer::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert!(csv_iter.next().unwrap().is_err());
|
||||
}
|
||||
}
|
||||
|
|
@ -1,10 +1,9 @@
|
|||
mod documents_from_csv;
|
||||
mod update_store;
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
|
||||
use std::fmt::Display;
|
||||
use std::fs::{create_dir_all, File};
|
||||
use std::io::Cursor;
|
||||
use std::io::{BufRead, BufReader, Cursor};
|
||||
use std::net::SocketAddr;
|
||||
use std::num::{NonZeroU32, NonZeroUsize};
|
||||
use std::path::PathBuf;
|
||||
|
|
@ -39,7 +38,6 @@ use warp::http::Response;
|
|||
use warp::Filter;
|
||||
|
||||
use self::update_store::UpdateStore;
|
||||
use crate::documents_from_csv::CSVDocumentDeserializer;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[global_allocator]
|
||||
|
|
@ -1041,11 +1039,11 @@ fn documents_from_jsonl(reader: impl io::Read) -> anyhow::Result<Vec<u8>> {
|
|||
let mut writer = Cursor::new(Vec::new());
|
||||
let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
|
||||
|
||||
let values = serde_json::Deserializer::from_reader(reader)
|
||||
.into_iter::<serde_json::Map<String, serde_json::Value>>();
|
||||
for document in values {
|
||||
let document = document?;
|
||||
documents.add_documents(document)?;
|
||||
let mut buf = String::new();
|
||||
let mut reader = BufReader::new(reader);
|
||||
|
||||
while reader.read_line(&mut buf)? > 0 {
|
||||
documents.extend_from_json(&mut buf.as_bytes())?;
|
||||
}
|
||||
documents.finish()?;
|
||||
|
||||
|
|
@ -1056,8 +1054,7 @@ fn documents_from_json(reader: impl io::Read) -> anyhow::Result<Vec<u8>> {
|
|||
let mut writer = Cursor::new(Vec::new());
|
||||
let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
|
||||
|
||||
let json: serde_json::Value = serde_json::from_reader(reader)?;
|
||||
documents.add_documents(json)?;
|
||||
documents.extend_from_json(reader)?;
|
||||
documents.finish()?;
|
||||
|
||||
Ok(writer.into_inner())
|
||||
|
|
@ -1065,16 +1062,7 @@ fn documents_from_json(reader: impl io::Read) -> anyhow::Result<Vec<u8>> {
|
|||
|
||||
fn documents_from_csv(reader: impl io::Read) -> anyhow::Result<Vec<u8>> {
|
||||
let mut writer = Cursor::new(Vec::new());
|
||||
let mut documents = milli::documents::DocumentBatchBuilder::new(&mut writer)?;
|
||||
|
||||
let iter = CSVDocumentDeserializer::from_reader(reader)?;
|
||||
|
||||
for doc in iter {
|
||||
let doc = doc?;
|
||||
documents.add_documents(doc)?;
|
||||
}
|
||||
|
||||
documents.finish()?;
|
||||
milli::documents::DocumentBatchBuilder::from_csv(reader, &mut writer)?.finish()?;
|
||||
|
||||
Ok(writer.into_inner())
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue