feat: implement reader for parquet type files

Author: bohblue2

hftbacktest/Cargo.toml

@@ -22,7 +22,7 @@ rust-version = "1.82"
 
 [features]
 default = ["backtest", "live"]
-backtest = ["zip", "uuid", "nom", "hftbacktest-derive"]
+backtest = ["zip", "uuid", "nom", "hftbacktest-derive", "rayon", "parquet", "arrow-array"]
 live = ["chrono", "tokio", "futures-util", "iceoryx2", "rand"]
 unstable_fuse = []
 
@@ -42,6 +42,10 @@ nom = { version = "7.1.3", optional = true }
 iceoryx2 = { version = "0.4.1", optional = true, features = ["logger_tracing"] }
 hftbacktest-derive = { path = "../hftbacktest-derive", optional = true, version = "0.2.0" }
 
+parquet = { version = "53.3.0", optional = true }
+arrow-array = { version = "53.3.0", optional = true }
+rayon = { version = "1.10.0", optional = true }
+
 [dev-dependencies]
 tracing-subscriber = { version = "0.3.18", features = [] }
 clap = { version = "4.5.4", features = ["derive"] }

hftbacktest/src/backtest/data/mod.rs

@@ -1,4 +1,5 @@
 mod npy;
+mod parquet;
 mod reader;
 
 use std::{

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

@@ -0,0 +1,72 @@
+use std::fs::File;
+use std::sync::{Arc, Mutex}; 
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+use parquet::file::reader::{FileReader, SerializedFileReader};
+use tracing::{debug, info};
+use crate::backtest::data::{Data, DataPtr, POD};
+use crate::prelude::Event;
+
+use super::NpyDTyped;
+use arrow_array::{UInt64Array, Int64Array, Float64Array};
+use rayon::prelude::*;
+
+
+pub fn read_parquet_file<D: NpyDTyped + Clone>(filepath: &str) -> std::io::Result<Data<D>> {
+    let batch_size = 1024 * 1024;
+    let events_capacity = 150_000_000;
+
+    let file = File::open(filepath)?;
+    let builder = ParquetRecordBatchReaderBuilder::try_new(file)
+        .unwrap()
+        .with_batch_size(batch_size);
+    let reader = builder.build().unwrap();
+    let events = Arc::new(Mutex::new(Vec::with_capacity(events_capacity))); 
+
+    // If we use parallel loading here, we need to re-sort by exch_ts in order.
+    // This is because exch_ts and local_ts are sorted in chronological order.
+    reader.into_iter().par_bridge().for_each(|maybe_batch| { 
+        let batch = maybe_batch.unwrap();
+
+        let ev_col = batch.column(0).as_any().downcast_ref::<UInt64Array>().unwrap();
+        let exch_ts_col = batch.column(1).as_any().downcast_ref::<Int64Array>().unwrap();
+        let local_ts_col = batch.column(2).as_any().downcast_ref::<Int64Array>().unwrap();
+        let px_col = batch.column(3).as_any().downcast_ref::<Float64Array>().unwrap();
+        let qty_col = batch.column(4).as_any().downcast_ref::<Float64Array>().unwrap();
+        let order_id_col = batch.column(5).as_any().downcast_ref::<UInt64Array>().unwrap();
+        let ival_col = batch.column(6).as_any().downcast_ref::<Int64Array>().unwrap();
+        let fval_col = batch.column(7).as_any().downcast_ref::<Float64Array>().unwrap();
+
+        let mut local_events: Vec<Event> = Vec::with_capacity(batch.num_rows());
+        for row in 0..batch.num_rows() {
+            local_events.push(Event {
+                ev: ev_col.value(row),
+                exch_ts: exch_ts_col.value(row),
+                local_ts: local_ts_col.value(row),
+                px: px_col.value(row),
+                qty: qty_col.value(row),
+                order_id: order_id_col.value(row),
+                ival: ival_col.value(row),
+                fval: fval_col.value(row),
+            });
+        }
+        debug!("Read {} events", local_events.len());
+        let mut events = events.lock().unwrap(); 
+        events.extend(local_events); 
+    });
+
+    let mut events = events.lock().unwrap();
+    events.par_sort_by_key(|event| event.exch_ts); 
+    let data_ptr = DataPtr::new(events.len() * std::mem::size_of::<D>());
+    
+    // Copy events to DataPtr
+    unsafe {
+        std::ptr::copy_nonoverlapping(
+            events.as_ptr() as *const u8,
+            data_ptr.ptr as *mut u8,
+            events.len() * std::mem::size_of::<D>()
+        );
+    }
+
+    let data = unsafe { Data::from_data_ptr(data_ptr, 0) };
+    Ok(data)
+}

hftbacktest/src/backtest/data/reader.rs

@@ -24,6 +24,8 @@ use crate::{
     types::Event,
 };
 
+use super::parquet;
+
 /// Data source for the [`Reader`].
 #[derive(Clone, Debug)]
 pub enum DataSource<D>
@@ -432,6 +434,30 @@ where
                         }
                     }
                 });
+            } else if key.ends_with(".parquet") {
+                let tx = self.tx.clone();
+                let filepath = key.to_string();
+                let preprocessor = self.preprocessor.clone();
+
+                let _ = thread::spawn(move || {
+                    let load_data = |filepath: &str| {
+                        let mut data = parquet::read_parquet_file::<D>(filepath)?;
+                        if let Some(preprocessor) = &preprocessor {
+                            preprocessor.preprocess(&mut data)?;
+                        }
+                        Ok(data)
+                    };
+                    // SendError occurs only if Reader is already destroyed. Since no data is needed
+                    // once the Reader is destroyed, SendError is safely suppressed.
+                    match load_data(&filepath) {
+                        Ok(data) => {
+                            let _ = tx.send(LoadDataResult::ok(filepath, data));
+                        }
+                        Err(err) => {
+                            let _ = tx.send(LoadDataResult::err(filepath, err));
+                        }
+                    }
+                });
             } else {
                 return Err(BacktestError::DataError(IoError::new(
                     ErrorKind::InvalidData,