rs_zips2meta2rbat2stream/
lib.rs

1pub use arrow;
2pub use futures;
3
4use std::io;
5
6use io::BufReader;
7
8use std::path::PathBuf;
9
10use std::fs::DirEntry;
11use std::fs::File;
12
13use futures::Stream;
14
15use arrow::record_batch::RecordBatch;
16
17pub trait ZipToBatch: Sync + Send + 'static {
18    type ZipId;
19
20    fn zip2batch(&self, id: Self::ZipId) -> Result<RecordBatch, io::Error>;
21}
22
23pub struct ZipConvFs {
24    pub root: PathBuf,
25}
26
27impl ZipToBatch for ZipConvFs {
28    type ZipId = String;
29
30    fn zip2batch(&self, id: Self::ZipId) -> Result<RecordBatch, io::Error> {
31        let fullpath = self.root.join(&id);
32        let f: File = File::open(fullpath)?;
33        let rdr = BufReader::new(f);
34        let z = zip::ZipArchive::new(rdr)?;
35
36        rs_zip2meta2rbat::sync::zip2record_batch(id, z)
37    }
38}
39
40pub fn ids2stream<I, Z>(ids: I, zconv: Z) -> impl Stream<Item = Result<RecordBatch, io::Error>>
41where
42    Z: ZipToBatch,
43    I: Iterator<Item = Result<Z::ZipId, io::Error>>,
44{
45    async_stream::try_stream! {
46        for rid in ids {
47            let id: Z::ZipId = rid?;
48            let rb: RecordBatch = zconv.zip2batch(id)?;
49            yield rb;
50        }
51    }
52}
53
54pub fn dirent2name(dirent: DirEntry) -> Result<String, io::Error> {
55    let ostr = dirent.file_name();
56    ostr.into_string()
57        .map_err(|_| "invalid file name")
58        .map_err(io::Error::other)
59}
60
61pub fn is_zip_name(name: &str) -> bool {
62    name.ends_with(".zip")
63}
64
65pub fn rname2zname(rname: Result<String, io::Error>) -> Option<Result<String, io::Error>> {
66    match rname {
67        Err(e) => Some(Err(e)),
68        Ok(name) => {
69            let is_zip: bool = is_zip_name(&name);
70            is_zip.then_some(Ok(name))
71        }
72    }
73}
74
75/// Gets the zip filenames from the dir and converts them to a stream.
76pub fn dir2zips2stream(
77    dirname: PathBuf,
78) -> Result<impl Stream<Item = Result<RecordBatch, io::Error>>, io::Error> {
79    let dirents = std::fs::read_dir(&dirname)?;
80    let mapd = dirents.map(|rdir| rdir.and_then(dirent2name));
81    let filtered = mapd.filter_map(rname2zname);
82    let zconv = ZipConvFs { root: dirname };
83    Ok(ids2stream(filtered, zconv))
84}
rs_zips2meta2rbat2stream/lib.rs

rs_zips2meta2rbat2stream/
lib.rs
