arrow.rs (1015B)
1 use anyhow::Result; 2 use polars::prelude::*; 3 use std::path::Path; 4 5 use crate::reader::ReadOptions; 6 7 pub fn read(path: &Path, opts: &ReadOptions) -> Result<DataFrame> { 8 let file = std::fs::File::open(path)?; 9 let mut df = IpcReader::new(file).finish()?; 10 11 if let Some(skip) = opts.skip_rows 12 && skip > 0 && skip < df.height() { 13 df = df.slice(skip as i64, df.height() - skip); 14 } 15 16 Ok(df) 17 } 18 19 #[cfg(test)] 20 mod tests { 21 use super::*; 22 use tempfile::NamedTempFile; 23 24 #[test] 25 fn read_arrow_roundtrip() { 26 let s1 = Series::new("x".into(), &[1i64, 2, 3]); 27 let mut df = DataFrame::new(vec![s1.into_column()]).unwrap(); 28 29 let f = NamedTempFile::with_suffix(".arrow").unwrap(); 30 let file = std::fs::File::create(f.path()).unwrap(); 31 IpcWriter::new(file).finish(&mut df).unwrap(); 32 33 let result = read(f.path(), &ReadOptions::default()).unwrap(); 34 assert_eq!(result.height(), 3); 35 assert_eq!(result.width(), 1); 36 } 37 }