use lora_executor::{LoraValue, Row};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Format {
Jsonl,
Json,
Csv,
}
impl Format {
pub fn from_extension(name: &str) -> Option<Self> {
let lower = name.to_ascii_lowercase();
let ext = lower.rsplit('.').next()?;
match ext {
"jsonl" | "ndjson" => Some(Format::Jsonl),
"json" => Some(Format::Json),
"csv" => Some(Format::Csv),
_ => None,
}
}
pub fn content_type(&self) -> &'static str {
match self {
Format::Jsonl => "application/x-ndjson",
Format::Json => "application/json",
Format::Csv => "text/csv",
}
}
pub fn as_str(&self) -> &'static str {
match self {
Format::Jsonl => "jsonl",
Format::Json => "json",
Format::Csv => "csv",
}
}
pub fn parse(name: &str) -> Option<Self> {
match name.to_ascii_lowercase().as_str() {
"jsonl" | "ndjson" => Some(Format::Jsonl),
"json" => Some(Format::Json),
"csv" => Some(Format::Csv),
_ => None,
}
}
}
pub trait RowEncoder {
fn begin(&mut self, columns: &[String]) -> std::io::Result<()>;
fn write_row(&mut self, row: &Row) -> std::io::Result<()>;
fn write_named_row(&mut self, columns: &[(String, LoraValue)]) -> std::io::Result<()>;
fn finish(&mut self) -> std::io::Result<()>;
}
pub trait RowDecoder {
fn header(&mut self) -> std::io::Result<Option<Vec<String>>>;
fn next_row(&mut self) -> std::io::Result<Option<Vec<(String, LoraValue)>>>;
}
pub trait StreamingRowDecoder {
fn feed(&mut self, chunk: &[u8]) -> std::io::Result<()>;
fn drain(&mut self) -> std::io::Result<Vec<Vec<(String, LoraValue)>>>;
fn finish(&mut self) -> std::io::Result<Vec<Vec<(String, LoraValue)>>>;
fn header(&self) -> Option<&[String]>;
fn bytes_fed(&self) -> u64;
fn rows_emitted(&self) -> u64;
fn set_permissive(&mut self, _on: bool) {}
fn take_errors(&mut self) -> Vec<RowParseError> {
Vec::new()
}
}
#[derive(Debug, Clone)]
pub struct RowParseError {
pub row: u64,
pub column: Option<String>,
pub raw_sample: String,
pub message: String,
}
pub const RAW_SAMPLE_MAX_CHARS: usize = 200;
impl RowParseError {
pub fn make_sample(raw: &str) -> String {
let cleaned: String = raw
.chars()
.map(|c| if c.is_control() { '·' } else { c })
.collect();
if cleaned.chars().count() <= RAW_SAMPLE_MAX_CHARS {
cleaned
} else {
let truncated: String = cleaned.chars().take(RAW_SAMPLE_MAX_CHARS).collect();
format!("{truncated}…")
}
}
pub fn make_sample_from_bytes(raw: &[u8]) -> String {
Self::make_sample(&String::from_utf8_lossy(raw))
}
}
impl std::fmt::Display for RowParseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "row {}", self.row)?;
if let Some(col) = &self.column {
write!(f, ", column `{col}`")?;
}
write!(f, ": {}", self.message)?;
if !self.raw_sample.is_empty() {
write!(f, " (raw: `{}`)", self.raw_sample)?;
}
Ok(())
}
}
impl std::error::Error for RowParseError {}
pub fn row_parse_io_error(err: RowParseError) -> std::io::Error {
std::io::Error::new(std::io::ErrorKind::InvalidData, err)
}
pub fn downcast_row_parse_error(err: &std::io::Error) -> Option<&RowParseError> {
err.get_ref()?.downcast_ref::<RowParseError>()
}
pub fn write_all_rows<E, I>(encoder: &mut E, columns: &[String], rows: I) -> std::io::Result<u64>
where
E: RowEncoder,
I: IntoIterator<Item = Row>,
{
encoder.begin(columns)?;
let mut count = 0u64;
for row in rows {
encoder.write_row(&row)?;
count += 1;
}
encoder.finish()?;
Ok(count)
}
pub fn invalid_data<E: Into<Box<dyn std::error::Error + Send + Sync>>>(err: E) -> std::io::Error {
std::io::Error::new(std::io::ErrorKind::InvalidData, err)
}