Introduce benchmarks for npy/npz readers

hftbacktest/Cargo.toml

@@ -54,6 +54,11 @@ hftbacktest-derive = { path = "../hftbacktest-derive", optional = true, version
 [dev-dependencies]
 tracing-subscriber = { version = "0.3.18", features = [] }
 clap = { version = "4.5.4", features = ["derive"] }
+criterion = { version = "0.5", features = ["html_reports"] }
+
+[[bench]]
+name = "formats"
+harness = false
 
 [package.metadata.docs.rs]
 all-features = true

hftbacktest/benches/formats.rs

@@ -0,0 +1,62 @@
+use criterion::*;
+use hftbacktest::backtest::data::{read_npy_file, write_npy, write_npz};
+use hftbacktest::{
+    backtest::data::{read_npz_file, Data},
+    types::Event,
+};
+use std::fs::File;
+use std::time::Duration;
+use std::{fs, ops::Index};
+
+fn bench(c: &mut Criterion) {
+    let mut group = c.benchmark_group("format-throughput");
+    let events: Vec<Event> = (0..1000_000)
+        .map(|id| Event {
+            ev: id,
+            exch_ts: 1_000_000,
+            local_ts: 1_000_001,
+            px: 1.0,
+            qty: 1.0,
+            order_id: 1,
+            ival: 100,
+            fval: 100.0,
+        })
+        .collect();
+
+    let mut npy_file = File::create("bench.npy").expect("couldn't create bench.npy");
+    let mut npz_file = File::create("bench.npz").expect("couldn't create bench.npz");
+
+    write_npy(&mut npy_file, &events).expect("failed to generate npy file");
+    write_npz(&mut npz_file, &events).expect("failed to generate npz file");
+
+    group.throughput(Throughput::Elements(events.len() as u64));
+    group.warm_up_time(Duration::from_secs(10));
+    group.bench_function("npz", |b| {
+        b.iter(|| benchmark_npz_file())
+    });
+    group.bench_function("npy", |b| b.iter(|| benchmark_npy_file()));
+    group.finish();
+
+    let _ = fs::remove_file("bench.npy");
+    let _ = fs::remove_file("bench.npz");
+}
+
+#[inline]
+fn read_all(data: Data<Event>) {
+    for index in 0..data.len() {
+        black_box(data.index(index));
+    }
+}
+
+fn benchmark_npz_file() {
+    let data = read_npz_file::<Event>("bench.npz", "data").unwrap();
+    read_all(data);
+}
+
+fn benchmark_npy_file() {
+    let data = read_npy_file::<Event>("bench.npy").unwrap();
+    read_all(data);
+}
+
+criterion_group!(benches, bench);
+criterion_main!(benches);

hftbacktest/src/backtest/data/mod.rs

@@ -10,7 +10,7 @@ use std::{
     slice::SliceIndex,
 };
 
-pub use npy::{read_npy_file, read_npz_file, write_npy, Field, NpyDTyped, NpyHeader};
+pub use npy::{read_npy_file, read_npz_file, write_npy, write_npz, Field, NpyDTyped, NpyHeader};
 pub use reader::{Cache, DataSource, Reader};
 
 use crate::utils::{AlignedArray, CACHE_LINE_SIZE};

hftbacktest/src/backtest/data/npy/mod.rs

@@ -1,8 +1,10 @@
 use std::{
     fs::File,
-    io::{Error, ErrorKind, Read, Write},
+    io::{Error, ErrorKind, Read, Seek, Write},
 };
 
+use zip::{write::FileOptions, ZipWriter};
+
 use crate::backtest::data::{npy::parser::Value, Data, DataPtr, POD};
 
 mod parser;
@@ -255,6 +257,16 @@ pub fn write_npy<W: Write, T: NpyDTyped>(write: &mut W, data: &[T]) -> std::io::
     Ok(())
 }
 
+pub fn write_npz<W: Write + Seek, T: NpyDTyped>(write: &mut W, data: &[T]) -> std::io::Result<()> {
+    let mut archive = ZipWriter::new(write);
+
+    archive.start_file::<_, (), _>("data.npy", FileOptions::default())?;
+    write_npy(&mut archive, data)?;
+    archive.finish()?;
+
+    Ok(())
+}
+
 fn vec_as_bytes<T>(vec: &[T]) -> &[u8] {
     let len = std::mem::size_of_val(vec);
     let ptr = vec.as_ptr() as *const u8;