json.rs (1804B)
1 use anyhow::Result; 2 use polars::prelude::*; 3 use std::io::Write as IoWrite; 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 json_format = match format { 10 Format::Ndjson => JsonFormat::JsonLines, 11 _ => JsonFormat::Json, 12 }; 13 14 match path { 15 Some(p) => { 16 let file = std::fs::File::create(p)?; 17 JsonWriter::new(file) 18 .with_json_format(json_format) 19 .finish(df)?; 20 } 21 None => { 22 let mut buf = Vec::new(); 23 JsonWriter::new(&mut buf) 24 .with_json_format(json_format) 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_json_roundtrip() { 39 let s = Series::new("x".into(), &[1i64, 2]); 40 let mut df = DataFrame::new(vec![s.into_column()]).unwrap(); 41 42 let f = NamedTempFile::with_suffix(".json").unwrap(); 43 write(&mut df, Some(f.path()), Format::Json).unwrap(); 44 45 let result = crate::readers::json::read(f.path(), Format::Json, &crate::reader::ReadOptions::default()).unwrap(); 46 assert_eq!(result.height(), 2); 47 } 48 49 #[test] 50 fn write_ndjson_roundtrip() { 51 let s = Series::new("x".into(), &[1i64, 2]); 52 let mut df = DataFrame::new(vec![s.into_column()]).unwrap(); 53 54 let f = NamedTempFile::with_suffix(".ndjson").unwrap(); 55 write(&mut df, Some(f.path()), Format::Ndjson).unwrap(); 56 57 let result = crate::readers::json::read(f.path(), Format::Ndjson, &crate::reader::ReadOptions::default()).unwrap(); 58 assert_eq!(result.height(), 2); 59 } 60 }