kibble 0.1.0

//! Read any supported data file (jsonl or parquet) as a stream of JSON records.
//! Both formats reduce to `Vec<serde_json::Value>` so downstream schema detection
//! is format-agnostic.

use serde_json::Value;
use std::path::Path;

pub enum RecordFormat {
    Jsonl,
    Parquet,
    Unsupported(String),
}

pub fn detect_format(path: &Path) -> RecordFormat {
    match path.extension().and_then(|e| e.to_str()).map(|s| s.to_lowercase()).as_deref() {
        Some("jsonl") | Some("json") => RecordFormat::Jsonl,
        Some("parquet") => RecordFormat::Parquet,
        Some(other) => RecordFormat::Unsupported(other.to_string()),
        None => RecordFormat::Unsupported(String::new()),
    }
}

/// Read up to `limit` records (None = all) from `path`.
pub fn read_records(path: &Path, limit: Option<usize>) -> std::io::Result<Vec<Value>> {
    match detect_format(path) {
        RecordFormat::Jsonl => read_jsonl(path, limit),
        RecordFormat::Parquet => read_parquet(path, limit),
        RecordFormat::Unsupported(ext) => Err(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            format!("unsupported record format: .{ext} ({})", path.display()),
        )),
    }
}

fn read_jsonl(path: &Path, limit: Option<usize>) -> std::io::Result<Vec<Value>> {
    let text = std::fs::read_to_string(path)?;
    let mut out = Vec::new();
    for line in text.lines() {
        if line.trim().is_empty() {
            continue;
        }
        let v: Value = serde_json::from_str(line)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, format!("{}: {e}", path.display())))?;
        out.push(v);
        if let Some(n) = limit {
            if out.len() >= n {
                break;
            }
        }
    }
    Ok(out)
}

fn read_parquet(path: &Path, limit: Option<usize>) -> std::io::Result<Vec<Value>> {
    use parquet::file::reader::{FileReader, SerializedFileReader};
    let file = std::fs::File::open(path)?;
    let reader = SerializedFileReader::new(file)
        .map_err(|e| std::io::Error::other(format!("parquet open {}: {e}", path.display())))?;
    let iter = reader
        .get_row_iter(None)
        .map_err(|e| std::io::Error::other(format!("parquet rows {}: {e}", path.display())))?;
    let mut out = Vec::new();
    for row in iter {
        let row = row.map_err(|e| std::io::Error::other(format!("parquet row {}: {e}", path.display())))?;
        out.push(row.to_json_value());
        if let Some(n) = limit {
            if out.len() >= n {
                break;
            }
        }
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::Path;

    fn tmp(name: &str) -> std::path::PathBuf {
        let d = std::env::temp_dir().join(format!("kibble_records_{}", std::process::id()));
        let _ = std::fs::create_dir_all(&d);
        d.join(name)
    }

    #[test]
    fn detect_format_by_extension() {
        assert!(matches!(detect_format(Path::new("a.jsonl")), RecordFormat::Jsonl));
        assert!(matches!(detect_format(Path::new("a.json")), RecordFormat::Jsonl));
        assert!(matches!(detect_format(Path::new("a.parquet")), RecordFormat::Parquet));
        assert!(matches!(detect_format(Path::new("a.csv")), RecordFormat::Unsupported(_)));
        assert!(matches!(detect_format(Path::new("noext")), RecordFormat::Unsupported(_)));
    }

    #[test]
    fn reads_jsonl_records_skipping_blank_lines() {
        let p = tmp("d.jsonl");
        std::fs::write(&p, "{\"a\":1}\n\n{\"a\":2}\n").unwrap();
        let recs = read_records(&p, None).unwrap();
        assert_eq!(recs.len(), 2);
        assert_eq!(recs[0]["a"], 1);
        assert_eq!(recs[1]["a"], 2);
    }

    #[test]
    fn jsonl_limit_is_respected() {
        let p = tmp("lim.jsonl");
        std::fs::write(&p, "{\"a\":1}\n{\"a\":2}\n{\"a\":3}\n").unwrap();
        let recs = read_records(&p, Some(2)).unwrap();
        assert_eq!(recs.len(), 2);
    }

    #[test]
    fn unsupported_format_errors() {
        let p = tmp("x.csv");
        std::fs::write(&p, "a,b\n1,2\n").unwrap();
        assert!(read_records(&p, None).is_err());
    }

    fn write_test_parquet(path: &Path) {
        use parquet::data_type::{ByteArray, ByteArrayType};
        use parquet::file::properties::WriterProperties;
        use parquet::file::writer::SerializedFileWriter;
        use parquet::schema::parser::parse_message_type;
        use std::sync::Arc;
        let schema = parse_message_type(
            "message s { REQUIRED BYTE_ARRAY instruction (UTF8); REQUIRED BYTE_ARRAY output (UTF8); }",
        )
        .unwrap();
        let props = Arc::new(WriterProperties::builder().build());
        let file = std::fs::File::create(path).unwrap();
        let mut w = SerializedFileWriter::new(file, Arc::new(schema), props).unwrap();
        let mut rg = w.next_row_group().unwrap();
        let mut c = rg.next_column().unwrap().unwrap();
        c.typed::<ByteArrayType>()
            .write_batch(&[ByteArray::from("q1"), ByteArray::from("q2")], None, None)
            .unwrap();
        c.close().unwrap();
        let mut c = rg.next_column().unwrap().unwrap();
        c.typed::<ByteArrayType>()
            .write_batch(&[ByteArray::from("a1"), ByteArray::from("a2")], None, None)
            .unwrap();
        c.close().unwrap();
        rg.close().unwrap();
        w.close().unwrap();
    }

    #[test]
    fn reads_parquet_records_as_json_objects() {
        let p = tmp("d.parquet");
        write_test_parquet(&p);
        let recs = read_records(&p, None).unwrap();
        assert_eq!(recs.len(), 2);
        assert_eq!(recs[0]["instruction"], "q1");
        assert_eq!(recs[0]["output"], "a1");
        assert_eq!(recs[1]["instruction"], "q2");
    }
}