use futures::stream::TryStreamExt;
use once_cell::sync::Lazy;
use std::{
fs::{self, File},
io::{self, prelude::*, BufWriter},
path::{Path, PathBuf},
};
use tfrecord::{Error, ExampleStream, FeatureKind};
static INPUT_TFRECORD_PATH: Lazy<PathBuf> = Lazy::new(|| {
(move || {
let url = include_str!("../tests/tfrecord_link.txt");
let path = DATA_DIR.join("input.tfrecord");
let mut writer = BufWriter::new(File::create(&path)?);
io::copy(&mut ureq::get(url).call()?.into_reader(), &mut writer)?;
writer.flush()?;
anyhow::Ok(path)
})()
.unwrap()
});
static DATA_DIR: Lazy<&Path> = Lazy::new(|| {
let path = Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/test_data"));
fs::create_dir_all(path).unwrap();
path
});
#[async_std::main]
pub async fn main() -> Result<(), Error> {
let stream = ExampleStream::open(&*INPUT_TFRECORD_PATH, Default::default()).await?;
println!("example_no\tfeature_no\tname\ttype\tsize");
stream
.try_fold(0, |example_index, example| {
async move {
for (feature_index, (name, feature)) in example.into_iter().enumerate() {
print!("{}\t{}\t{}\t", example_index, feature_index, name);
use FeatureKind as F;
match feature.into_kinds() {
Some(F::Bytes(list)) => {
println!("bytes\t{}", list.len());
}
Some(F::F32(list)) => {
println!("float\t{}", list.len());
}
Some(F::I64(list)) => {
println!("int64\t{}", list.len());
}
None => {
println!("none");
}
}
}
Ok(example_index + 1)
}
})
.await?;
Ok(())
}