MeiliSearch/meilisearch-lib/src/document_formats.rs

105 lines
3.2 KiB
Rust
Raw Normal View History

2021-09-28 22:58:48 +02:00
use std::fmt;
use std::io::{self, BufRead, BufReader, BufWriter, Cursor, Read, Seek, Write};
2021-09-28 11:59:55 +02:00
use meilisearch_error::{internal_error, Code, ErrorCode};
2021-09-28 11:59:55 +02:00
use milli::documents::DocumentBatchBuilder;
type Result<T> = std::result::Result<T, DocumentFormatError>;
#[derive(Debug)]
pub enum PayloadType {
2021-09-29 10:17:52 +02:00
Ndjson,
Json,
2021-09-28 22:58:48 +02:00
Csv,
2021-09-28 11:59:55 +02:00
}
impl fmt::Display for PayloadType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
2021-09-29 10:17:52 +02:00
PayloadType::Ndjson => write!(f, "ndjson"),
PayloadType::Json => write!(f, "json"),
2021-09-28 22:58:48 +02:00
PayloadType::Csv => write!(f, "csv"),
2021-09-28 11:59:55 +02:00
}
}
}
#[derive(thiserror::Error, Debug)]
pub enum DocumentFormatError {
2021-11-03 14:25:49 +01:00
#[error("An internal error has occurred. `{0}`.")]
2021-09-28 11:59:55 +02:00
Internal(Box<dyn std::error::Error + Send + Sync + 'static>),
2021-10-26 19:36:48 +02:00
#[error("The `{1}` payload provided is malformed. `{0}`.")]
2021-09-28 22:22:59 +02:00
MalformedPayload(
Box<dyn std::error::Error + Send + Sync + 'static>,
PayloadType,
),
2021-09-28 11:59:55 +02:00
}
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),
}
}
}
2021-09-30 10:35:24 +02:00
impl ErrorCode for DocumentFormatError {
fn error_code(&self) -> Code {
match self {
DocumentFormatError::Internal(_) => Code::Internal,
DocumentFormatError::MalformedPayload(_, _) => Code::MalformedPayload,
}
}
}
internal_error!(DocumentFormatError: io::Error);
2021-09-28 11:59:55 +02:00
/// reads csv from input and write an obkv batch to writer.
pub fn read_csv(input: impl Read, writer: impl Write + Seek) -> Result<usize> {
let writer = BufWriter::new(writer);
2021-10-26 19:36:48 +02:00
let builder =
DocumentBatchBuilder::from_csv(input, writer).map_err(|e| (PayloadType::Csv, e))?;
let count = builder.finish().map_err(|e| (PayloadType::Csv, e))?;
2021-09-28 22:58:48 +02:00
Ok(count)
2021-09-28 22:58:48 +02:00
}
/// reads jsonl from input and write an obkv batch to writer.
pub fn read_ndjson(input: impl Read, writer: impl Write + Seek) -> Result<usize> {
let mut reader = BufReader::new(input);
let writer = BufWriter::new(writer);
let mut builder = DocumentBatchBuilder::new(writer).map_err(|e| (PayloadType::Ndjson, e))?;
let mut buf = String::new();
2021-09-28 11:59:55 +02:00
while reader.read_line(&mut buf)? > 0 {
// skip empty lines
if buf == "\n" {
buf.clear();
continue;
}
builder
.extend_from_json(Cursor::new(&buf.as_bytes()))
.map_err(|e| (PayloadType::Ndjson, e))?;
buf.clear();
2021-09-28 11:59:55 +02:00
}
let count = builder.finish().map_err(|e| (PayloadType::Ndjson, e))?;
2021-09-28 11:59:55 +02:00
Ok(count)
2021-09-28 11:59:55 +02:00
}
/// reads json from input and write an obkv batch to writer.
pub fn read_json(input: impl Read, writer: impl Write + Seek) -> Result<usize> {
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))?;
let count = builder.finish().map_err(|e| (PayloadType::Json, e))?;
Ok(count)
}