csv.rs (1824B)
1 use anyhow::Result; 2 use polars::prelude::*; 3 use std::io::Write; 4 use std::path::Path; 5 6 use crate::format::Format; 7 8 pub fn write(df: &mut DataFrame, path: Option<&Path>, format: Format) -> Result<()> { 9 let separator = match format { 10 Format::Tsv => b'\t', 11 _ => b',', 12 }; 13 14 match path { 15 Some(p) => { 16 let file = std::fs::File::create(p)?; 17 CsvWriter::new(file) 18 .with_separator(separator) 19 .finish(df)?; 20 } 21 None => { 22 let mut buf = Vec::new(); 23 CsvWriter::new(&mut buf) 24 .with_separator(separator) 25 .finish(df)?; 26 std::io::stdout().write_all(&buf)?; 27 } 28 } 29 Ok(()) 30 } 31 32 #[cfg(test)] 33 mod tests { 34 use super::*; 35 use tempfile::NamedTempFile; 36 37 #[test] 38 fn write_csv_roundtrip() { 39 let s1 = Series::new("name".into(), &["Alice", "Bob"]); 40 let s2 = Series::new("value".into(), &[100i64, 200]); 41 let mut df = DataFrame::new(vec![s1.into_column(), s2.into_column()]).unwrap(); 42 43 let f = NamedTempFile::with_suffix(".csv").unwrap(); 44 write(&mut df, Some(f.path()), Format::Csv).unwrap(); 45 46 let result = crate::readers::csv::read(f.path(), &crate::reader::ReadOptions::default()).unwrap(); 47 assert_eq!(result.height(), 2); 48 assert_eq!(result.get_column_names(), df.get_column_names()); 49 } 50 51 #[test] 52 fn write_tsv_uses_tab() { 53 let s = Series::new("x".into(), &[1i64]); 54 let mut df = DataFrame::new(vec![s.into_column()]).unwrap(); 55 56 let f = NamedTempFile::with_suffix(".tsv").unwrap(); 57 write(&mut df, Some(f.path()), Format::Tsv).unwrap(); 58 59 let content = std::fs::read_to_string(f.path()).unwrap(); 60 assert!(!content.contains(',')); 61 } 62 }