mirror of
https://github.com/meilisearch/MeiliSearch
synced 2025-07-04 20:37:15 +02:00
optimize document transform
fix error types bump milli
This commit is contained in:
marin postma
parent
e87146b0d9
commit
4ac005b094
6 changed files with 60 additions and 372 deletions
|
|
@ -1,10 +1,8 @@
|
|||
use std::fmt;
|
||||
use std::io::{self, Read, Result as IoResult, Seek, Write};
|
||||
use std::io::{self, BufRead, BufReader, BufWriter, Cursor, Read, Seek, Write};
|
||||
|
||||
use csv::{Reader as CsvReader, StringRecordsIntoIter};
|
||||
use meilisearch_error::{Code, ErrorCode};
|
||||
use milli::documents::DocumentBatchBuilder;
|
||||
use serde_json::{Deserializer, Map, Value};
|
||||
|
||||
type Result<T> = std::result::Result<T, DocumentFormatError>;
|
||||
|
||||
|
|
@ -36,6 +34,15 @@ pub enum DocumentFormatError {
|
|||
),
|
||||
}
|
||||
|
||||
impl From<(PayloadType, milli::documents::Error)> for DocumentFormatError {
|
||||
fn from((ty, error): (PayloadType, milli::documents::Error)) -> Self {
|
||||
match error {
|
||||
milli::documents::Error::Io(e) => Self::Internal(Box::new(e)),
|
||||
e => Self::MalformedPayload(Box::new(e), ty),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ErrorCode for DocumentFormatError {
|
||||
fn error_code(&self) -> Code {
|
||||
match self {
|
||||
|
|
@ -45,330 +52,47 @@ impl ErrorCode for DocumentFormatError {
|
|||
}
|
||||
}
|
||||
|
||||
internal_error!(DocumentFormatError: milli::documents::Error, io::Error);
|
||||
|
||||
macro_rules! malformed {
|
||||
($type:path, $e:expr) => {
|
||||
$e.map_err(|e| DocumentFormatError::MalformedPayload(Box::new(e), $type))
|
||||
};
|
||||
}
|
||||
internal_error!(DocumentFormatError: io::Error);
|
||||
|
||||
/// reads csv from input and write an obkv batch to writer.
|
||||
pub fn read_csv(input: impl Read, writer: impl Write + Seek) -> Result<()> {
|
||||
let mut builder = DocumentBatchBuilder::new(writer).unwrap();
|
||||
|
||||
let iter = CsvDocumentIter::from_reader(input)?;
|
||||
for doc in iter {
|
||||
let doc = doc?;
|
||||
builder.add_documents(doc).unwrap();
|
||||
}
|
||||
builder.finish().unwrap();
|
||||
let writer = BufWriter::new(writer);
|
||||
DocumentBatchBuilder::from_csv(input, writer)
|
||||
.map_err(|e| (PayloadType::Csv, e))?
|
||||
.finish()
|
||||
.map_err(|e| (PayloadType::Csv, e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// read jsonl from input and write an obkv batch to writer.
|
||||
/// reads jsonl from input and write an obkv batch to writer.
|
||||
pub fn read_ndjson(input: impl Read, writer: impl Write + Seek) -> Result<()> {
|
||||
let mut builder = DocumentBatchBuilder::new(writer)?;
|
||||
let stream = Deserializer::from_reader(input).into_iter::<Map<String, Value>>();
|
||||
let mut reader = BufReader::new(input);
|
||||
let writer = BufWriter::new(writer);
|
||||
|
||||
for value in stream {
|
||||
let value = malformed!(PayloadType::Ndjson, value)?;
|
||||
builder.add_documents(&value)?;
|
||||
let mut builder = DocumentBatchBuilder::new(writer).map_err(|e| (PayloadType::Ndjson, e))?;
|
||||
let mut buf = String::new();
|
||||
|
||||
while reader.read_line(&mut buf)? > 0 {
|
||||
builder
|
||||
.extend_from_json(Cursor::new(&buf.as_bytes()))
|
||||
.map_err(|e| (PayloadType::Ndjson, e))?;
|
||||
buf.clear();
|
||||
}
|
||||
|
||||
builder.finish()?;
|
||||
builder.finish().map_err(|e| (PayloadType::Ndjson, e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// read json from input and write an obkv batch to writer.
|
||||
/// reads json from input and write an obkv batch to writer.
|
||||
pub fn read_json(input: impl Read, writer: impl Write + Seek) -> Result<()> {
|
||||
let mut builder = DocumentBatchBuilder::new(writer).unwrap();
|
||||
|
||||
let documents: Vec<Map<String, Value>> =
|
||||
malformed!(PayloadType::Json, serde_json::from_reader(input))?;
|
||||
builder.add_documents(documents).unwrap();
|
||||
builder.finish().unwrap();
|
||||
let writer = BufWriter::new(writer);
|
||||
let mut builder = DocumentBatchBuilder::new(writer).map_err(|e| (PayloadType::Json, e))?;
|
||||
builder
|
||||
.extend_from_json(input)
|
||||
.map_err(|e| (PayloadType::Json, e))?;
|
||||
builder.finish().map_err(|e| (PayloadType::Json, e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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),
|
||||
// if the pattern isn't reconized, we keep the whole field.
|
||||
_otherwise => (header.to_string(), AllowedType::String),
|
||||
},
|
||||
None => (header.to_string(), AllowedType::String),
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CsvDocumentIter<R>
|
||||
where
|
||||
R: Read,
|
||||
{
|
||||
documents: StringRecordsIntoIter<R>,
|
||||
headers: Vec<(String, AllowedType)>,
|
||||
}
|
||||
|
||||
impl<R: Read> CsvDocumentIter<R> {
|
||||
pub fn from_reader(reader: R) -> IoResult<Self> {
|
||||
let mut records = CsvReader::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 CsvDocumentIter<R> {
|
||||
type Item = 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 = match field_type {
|
||||
AllowedType::Number => {
|
||||
malformed!(PayloadType::Csv, value.parse::<f64>().map(Value::from))
|
||||
}
|
||||
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(e)),
|
||||
}
|
||||
}
|
||||
|
||||
Some(Ok(document))
|
||||
}
|
||||
Err(e) => Some(Err(DocumentFormatError::MalformedPayload(
|
||||
Box::new(e),
|
||||
PayloadType::Csv,
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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 = CsvDocumentIter::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 = CsvDocumentIter::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 = CsvDocumentIter::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 = CsvDocumentIter::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 = CsvDocumentIter::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_colon_in_header() {
|
||||
let documents = r#"city:love:string,country:state,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CsvDocumentIter::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_colon_in_header() {
|
||||
let documents = r#"city:,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CsvDocumentIter::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_colon_in_header() {
|
||||
let documents = r#":city,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CsvDocumentIter::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(csv_iter.next().unwrap().unwrap()),
|
||||
json!({
|
||||
":city": "Boston",
|
||||
"country": "United States",
|
||||
"pop": "4628910",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[ignore]
|
||||
#[test]
|
||||
fn starting_by_colon_in_header2() {
|
||||
let documents = r#":string,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CsvDocumentIter::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert!(csv_iter.next().unwrap().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn double_colon_in_header() {
|
||||
let documents = r#"city::string,country,pop
|
||||
"Boston","United States","4628910""#;
|
||||
|
||||
let mut csv_iter = CsvDocumentIter::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 = CsvDocumentIter::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 = CsvDocumentIter::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 = CsvDocumentIter::from_reader(documents.as_bytes()).unwrap();
|
||||
|
||||
assert!(csv_iter.next().unwrap().is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue