parquet.rs (1167B)
1 use anyhow::Result; 2 use polars::prelude::*; 3 use std::path::Path; 4 5 pub fn write(df: &mut DataFrame, path: Option<&Path>) -> Result<()> { 6 let path = path.ok_or_else(|| anyhow::anyhow!("--convert parquet requires -o PATH"))?; 7 let file = std::fs::File::create(path)?; 8 ParquetWriter::new(file).finish(df)?; 9 Ok(()) 10 } 11 12 #[cfg(test)] 13 mod tests { 14 use super::*; 15 use tempfile::NamedTempFile; 16 17 #[test] 18 fn write_parquet_roundtrip() { 19 let s1 = Series::new("name".into(), &["Alice", "Bob"]); 20 let s2 = Series::new("value".into(), &[100i64, 200]); 21 let mut df = DataFrame::new(vec![s1.into_column(), s2.into_column()]).unwrap(); 22 23 let f = NamedTempFile::with_suffix(".parquet").unwrap(); 24 write(&mut df, Some(f.path())).unwrap(); 25 26 let result = crate::readers::parquet::read(f.path(), &crate::reader::ReadOptions::default()).unwrap(); 27 assert_eq!(result.height(), 2); 28 } 29 30 #[test] 31 fn write_parquet_no_path_errors() { 32 let s = Series::new("x".into(), &[1i64]); 33 let mut df = DataFrame::new(vec![s.into_column()]).unwrap(); 34 assert!(write(&mut df, None).is_err()); 35 } 36 }