arrow.rs (1064B)
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 arrow requires -o PATH"))?; 7 let file = std::fs::File::create(path)?; 8 IpcWriter::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_arrow_roundtrip() { 19 let s = Series::new("x".into(), &[1i64, 2, 3]); 20 let mut df = DataFrame::new(vec![s.into_column()]).unwrap(); 21 22 let f = NamedTempFile::with_suffix(".arrow").unwrap(); 23 write(&mut df, Some(f.path())).unwrap(); 24 25 let result = crate::readers::arrow::read(f.path(), &crate::reader::ReadOptions::default()).unwrap(); 26 assert_eq!(result.height(), 3); 27 } 28 29 #[test] 30 fn write_arrow_no_path_errors() { 31 let s = Series::new("x".into(), &[1i64]); 32 let mut df = DataFrame::new(vec![s.into_column()]).unwrap(); 33 assert!(write(&mut df, None).is_err()); 34 } 35 }