6beedca110
The function was doing a linear scan through 14.5M sorted MBP-10 snapshots for each of 1.12M OHLCV bars, resulting in ~16.2 trillion comparisons. Replaced with partition_point (binary search) for O(log n) per lookup, reducing total comparisons to ~27M — a ~600,000x improvement. This was the root cause of OFI computation taking 30+ minutes during hyperopt data loading on H100. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
555 lines
18 KiB
Rust
555 lines
18 KiB
Rust
//! MBP-10 Data Loader Helper
|
||
//!
|
||
//! Simple integration layer between DBN MBP-10 parser and OFI feature extraction.
|
||
//! Provides synchronous loading wrapper and snapshot window selection for real-time
|
||
//! Order Flow Imbalance (OFI) calculation.
|
||
//!
|
||
//! # Purpose
|
||
//!
|
||
//! - Load MBP-10 snapshots from DBN files
|
||
//! - Provide snapshot windows for OFI calculation
|
||
//! - Abstract async DBN parser complexity
|
||
//!
|
||
//! # Example
|
||
//!
|
||
//! ```no_run
|
||
//! use ml::features::mbp10_loader::load_mbp10_snapshots_sync;
|
||
//! use ml::features::ofi_calculator::OFICalculator;
|
||
//! use std::path::Path;
|
||
//!
|
||
//! let snapshots = load_mbp10_snapshots_sync(Path::new("test_data/ES.FUT.mbp10.dbn"))?;
|
||
//! let mut calculator = OFICalculator::new();
|
||
//!
|
||
//! for snapshot in &snapshots {
|
||
//! let features = calculator.calculate(snapshot)?;
|
||
//! // Use features for model input
|
||
//! }
|
||
//! ```
|
||
|
||
use data::providers::databento::{dbn_parser::DbnParser, mbp10::Mbp10Snapshot};
|
||
use std::path::Path;
|
||
|
||
use crate::ofi_calculator::OFICalculator;
|
||
use crate::trades_loader::load_trades_sync;
|
||
use crate::MLError;
|
||
|
||
/// Load MBP-10 snapshots from DBN file (synchronous wrapper)
|
||
///
|
||
/// This is a blocking wrapper around the async `DbnParser::parse_mbp10_file()`.
|
||
/// Use this for simple synchronous contexts like testing or batch processing.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `file_path` - Path to the DBN file containing MBP-10 data
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// * `Ok(Vec<Mbp10Snapshot>)` - All snapshots from the file
|
||
/// * `Err(MLError)` - If file cannot be read or parsed
|
||
///
|
||
/// # Performance
|
||
///
|
||
/// - Target: <10ms for typical 180-day files
|
||
/// - Memory: ~1KB per snapshot (10 levels × 48 bytes)
|
||
///
|
||
pub fn load_mbp10_snapshots_sync(file_path: &Path) -> Result<Vec<Mbp10Snapshot>, MLError> {
|
||
let parser = DbnParser::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
|
||
})?;
|
||
|
||
// Use tokio runtime for async operation
|
||
let runtime = tokio::runtime::Runtime::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create async runtime: {}", e))
|
||
})?;
|
||
|
||
runtime.block_on(async {
|
||
parser
|
||
.parse_mbp10_file(file_path)
|
||
.await
|
||
.map_err(|e| MLError::InsufficientData(format!("Failed to parse MBP-10 file: {}", e)))
|
||
})
|
||
}
|
||
|
||
/// Compute OFI features from an MBP-10 file in a single streaming pass.
|
||
///
|
||
/// This is ~10x faster than `load_mbp10_snapshots_sync` + separate OFI iteration
|
||
/// because it avoids materializing millions of intermediate `Mbp10Snapshot` clones.
|
||
/// The parser feeds snapshots by reference directly to the OFI calculator.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `file_path` - Path to the .dbn or .dbn.zst MBP-10 file
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// * `Ok(Vec<[f64; 8]>)` - OFI feature arrays (one per aggregated snapshot)
|
||
/// * `Err(MLError)` - If file cannot be read or parsed
|
||
///
|
||
pub fn compute_ofi_from_file(file_path: &Path) -> Result<Vec<[f64; 8]>, MLError> {
|
||
let parser = DbnParser::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
|
||
})?;
|
||
|
||
let mut calculator = OFICalculator::new();
|
||
let mut features = Vec::new();
|
||
|
||
let snapshot_count = parser
|
||
.parse_mbp10_streaming(file_path, 100, |snapshot| {
|
||
if let Ok(f) = calculator.calculate(snapshot) {
|
||
features.push(f.to_array());
|
||
} // Skip failed calculations (e.g. first snapshot with no prev)
|
||
})
|
||
.map_err(|e| MLError::InsufficientData(format!("MBP-10 streaming parse failed: {}", e)))?;
|
||
|
||
tracing::debug!(
|
||
"Computed {} OFI features from {} snapshots in {:?}",
|
||
features.len(),
|
||
snapshot_count,
|
||
file_path.file_name().unwrap_or_default()
|
||
);
|
||
|
||
Ok(features)
|
||
}
|
||
|
||
/// Compute OFI features from MBP-10 file enriched with trade data.
|
||
///
|
||
/// Interleaves trade events (from a `.dbn.zst` trades file) into the MBP-10
|
||
/// streaming pipeline. Before each MBP-10 snapshot is processed, all trades
|
||
/// with timestamps up to that snapshot are fed into the OFI calculator,
|
||
/// populating VPIN, Kyle's Lambda, and trade imbalance with real trade data
|
||
/// instead of leaving them at zero.
|
||
///
|
||
/// Falls back to `compute_ofi_from_file()` if the trade file doesn't exist
|
||
/// or fails to load.
|
||
pub fn compute_ofi_with_trades(
|
||
mbp10_file: &Path,
|
||
trades_file: &Path,
|
||
) -> Result<Vec<[f64; 8]>, MLError> {
|
||
// Load all trades upfront (sorted by timestamp). Trade files are much
|
||
// smaller than MBP-10 (~120 MB vs ~5+ GB), so this is fine in memory.
|
||
let trades = match load_trades_sync(trades_file) {
|
||
Ok(t) => {
|
||
tracing::info!(
|
||
"Loaded {} trades from {} for OFI enrichment",
|
||
t.len(),
|
||
trades_file.file_name().unwrap_or_default().to_string_lossy()
|
||
);
|
||
t
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
"Failed to load trades from {}: {} \u{2014} falling back to MBP-10 only",
|
||
trades_file.display(),
|
||
e
|
||
);
|
||
return compute_ofi_from_file(mbp10_file);
|
||
}
|
||
};
|
||
|
||
let parser = DbnParser::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
|
||
})?;
|
||
|
||
let mut calculator = OFICalculator::new();
|
||
let mut features = Vec::new();
|
||
let mut trade_cursor: usize = 0;
|
||
|
||
let snapshot_count = parser
|
||
.parse_mbp10_streaming(mbp10_file, 100, |snapshot| {
|
||
// Feed all trades with timestamp <= this snapshot's timestamp.
|
||
// This ensures VPIN/Kyle's Lambda/trade_imbalance are populated
|
||
// before the OFI calculator processes the snapshot.
|
||
while trade_cursor < trades.len()
|
||
&& trades[trade_cursor].timestamp <= snapshot.timestamp
|
||
{
|
||
let t = &trades[trade_cursor];
|
||
calculator.feed_trade(t.price, t.volume, t.is_buy);
|
||
trade_cursor += 1;
|
||
}
|
||
|
||
if let Ok(f) = calculator.calculate(snapshot) {
|
||
features.push(f.to_array());
|
||
} // Skip failed calculations (e.g. first snapshot)
|
||
})
|
||
.map_err(|e| MLError::InsufficientData(format!("MBP-10 streaming parse failed: {}", e)))?;
|
||
|
||
tracing::info!(
|
||
"Computed {} OFI features from {} snapshots + {} trades (cursor at {}/{})",
|
||
features.len(),
|
||
snapshot_count,
|
||
trades.len(),
|
||
trade_cursor,
|
||
trades.len()
|
||
);
|
||
|
||
Ok(features)
|
||
}
|
||
|
||
/// Find matching trade file for an MBP-10 file in the trades directory.
|
||
///
|
||
/// Matches by filename stem: `ES.FUT_2024-Q1.dbn.zst` in MBP-10 dir
|
||
/// → looks for `ES.FUT_2024-Q1.dbn.zst` in trades dir (maintaining
|
||
/// the same subdirectory structure).
|
||
fn find_matching_trade_file(mbp10_file: &Path, trades_dir: &Path) -> Option<std::path::PathBuf> {
|
||
let file_name = mbp10_file.file_name()?;
|
||
|
||
// Try exact match: trades_dir/SYMBOL/filename
|
||
// MBP-10 files are typically under mbp10_dir/ES.FUT/ES.FUT_2024-Q1.dbn.zst
|
||
// Trades would be under trades_dir/ES.FUT/ES.FUT_2024-Q1.dbn.zst
|
||
if let Some(parent) = mbp10_file.parent() {
|
||
if let Some(symbol_dir) = parent.file_name() {
|
||
let candidate = trades_dir.join(symbol_dir).join(file_name);
|
||
if candidate.exists() {
|
||
return Some(candidate);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Fallback: flat structure — trades_dir/filename
|
||
let candidate = trades_dir.join(file_name);
|
||
if candidate.exists() {
|
||
return Some(candidate);
|
||
}
|
||
|
||
None
|
||
}
|
||
|
||
/// Load OFI features from all MBP-10 files in a directory using parallel streaming.
|
||
///
|
||
/// Files are processed concurrently with rayon. Each file streams through the DBN parser
|
||
/// and computes OFI features inline — no intermediate `Vec<Mbp10Snapshot>` allocation.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `mbp10_dir` - Directory containing .dbn or .dbn.zst MBP-10 files
|
||
/// * `file_collector` - Function to collect files from the directory (supports both
|
||
/// flat and recursive collection patterns used by DQN/PPO adapters)
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// * `Some(Vec<[f64; 8]>)` - Concatenated OFI features from all files (file order preserved)
|
||
/// * `None` - If directory doesn't exist, has no files, or all files failed
|
||
///
|
||
#[allow(clippy::cognitive_complexity)]
|
||
pub fn load_ofi_features_parallel<F>(
|
||
mbp10_dir: &Path,
|
||
trades_dir: Option<&Path>,
|
||
file_collector: F,
|
||
) -> Option<Vec<[f64; 8]>>
|
||
where
|
||
F: FnOnce(&Path) -> Vec<std::path::PathBuf>,
|
||
{
|
||
use rayon::prelude::*;
|
||
|
||
if !mbp10_dir.exists() {
|
||
tracing::info!(
|
||
"No MBP10 directory at {}, OFI features will be zero-padded",
|
||
mbp10_dir.display()
|
||
);
|
||
return None;
|
||
}
|
||
|
||
let mbp10_files = file_collector(mbp10_dir);
|
||
if mbp10_files.is_empty() {
|
||
tracing::info!("No MBP10 .dbn files found in {}", mbp10_dir.display());
|
||
return None;
|
||
}
|
||
|
||
tracing::info!(
|
||
"Loading OFI features from {} MBP10 files (parallel streaming{})",
|
||
mbp10_files.len(),
|
||
if trades_dir.is_some() { " + trades" } else { "" }
|
||
);
|
||
|
||
let start = std::time::Instant::now();
|
||
|
||
// Process files in parallel — each file gets its own OFI calculator (independent state).
|
||
// If a trades directory is provided, each MBP-10 file looks for a matching trade file
|
||
// to enrich VPIN, Kyle's Lambda, and trade imbalance with real trade data.
|
||
let per_file_results: Vec<_> = mbp10_files
|
||
.par_iter()
|
||
.map(|file| {
|
||
let result = if let Some(t_dir) = trades_dir {
|
||
if let Some(trade_file) = find_matching_trade_file(file, t_dir) {
|
||
compute_ofi_with_trades(file, &trade_file)
|
||
} else {
|
||
tracing::debug!(
|
||
"No matching trade file for {}, using MBP-10 only",
|
||
file.file_name().unwrap_or_default().to_string_lossy()
|
||
);
|
||
compute_ofi_from_file(file)
|
||
}
|
||
} else {
|
||
compute_ofi_from_file(file)
|
||
};
|
||
|
||
match result {
|
||
Ok(features) => {
|
||
tracing::info!(
|
||
" {} \u{2192} {} OFI features",
|
||
file.file_name().unwrap_or_default().to_string_lossy(),
|
||
features.len()
|
||
);
|
||
features
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!("Failed to compute OFI from {}: {}", file.display(), e);
|
||
Vec::new()
|
||
}
|
||
}
|
||
})
|
||
.collect();
|
||
|
||
// Concatenate results in file order
|
||
let total_features: usize = per_file_results.iter().map(|v| v.len()).sum();
|
||
let mut all_ofi = Vec::with_capacity(total_features);
|
||
for features in per_file_results {
|
||
all_ofi.extend_from_slice(&features);
|
||
}
|
||
|
||
if all_ofi.is_empty() {
|
||
return None;
|
||
}
|
||
|
||
let elapsed = start.elapsed();
|
||
tracing::info!(
|
||
"Computed {} OFI feature vectors from {} files in {:.1}s ({:.0} features/sec)",
|
||
all_ofi.len(),
|
||
mbp10_files.len(),
|
||
elapsed.as_secs_f64(),
|
||
all_ofi.len() as f64 / elapsed.as_secs_f64()
|
||
);
|
||
|
||
Some(all_ofi)
|
||
}
|
||
|
||
/// Get snapshots for a specific timestamp window
|
||
///
|
||
/// Returns a slice of snapshots starting from the first snapshot at or before
|
||
/// the target timestamp. Used for calculating OFI features that require
|
||
/// forward-looking context (e.g., next N snapshots after a given timestamp).
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `snapshots` - All available snapshots (must be sorted by timestamp)
|
||
/// * `target_ts` - Target timestamp (nanoseconds since Unix epoch)
|
||
/// * `window_size` - Number of snapshots to include in window
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Slice of up to `window_size` snapshots starting from the first snapshot
|
||
/// with timestamp >= `target_ts`. Returns empty slice if no snapshots exist
|
||
/// at or after the target timestamp.
|
||
///
|
||
/// # Algorithm
|
||
///
|
||
/// - Simple linear search (TODO: optimize with binary search for large datasets)
|
||
/// - Returns snapshots starting from the first one with timestamp >= `target_ts`
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```ignore
|
||
/// // Snapshots at times: [1000, 2000, 3000, 4000]
|
||
/// // target_ts = 2000, window_size = 2
|
||
/// // Returns: [2000, 3000]
|
||
/// ```
|
||
///
|
||
pub fn get_snapshots_for_timestamp(
|
||
snapshots: &[Mbp10Snapshot],
|
||
target_ts: u64,
|
||
window_size: usize,
|
||
) -> &[Mbp10Snapshot] {
|
||
if snapshots.is_empty() {
|
||
return &[];
|
||
}
|
||
|
||
// Binary search for the first snapshot with timestamp >= target_ts
|
||
// Snapshots are pre-sorted by timestamp (from DBN file ordering)
|
||
let idx = snapshots.partition_point(|s| s.timestamp < target_ts);
|
||
if idx >= snapshots.len() {
|
||
return &[]; // Target timestamp is after all snapshots
|
||
}
|
||
let end_idx = (idx + window_size).min(snapshots.len());
|
||
&snapshots[idx..end_idx]
|
||
}
|
||
|
||
/// Get the most recent N snapshots ending at the given index
|
||
///
|
||
/// Useful for calculating features that require a rolling window of recent data.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `snapshots` - All available snapshots
|
||
/// * `end_idx` - Index of the last snapshot to include (exclusive)
|
||
/// * `window_size` - Number of snapshots to include
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Slice of up to `window_size` snapshots ending at `end_idx`
|
||
///
|
||
pub fn get_recent_snapshots(
|
||
snapshots: &[Mbp10Snapshot],
|
||
end_idx: usize,
|
||
window_size: usize,
|
||
) -> &[Mbp10Snapshot] {
|
||
if snapshots.is_empty() || end_idx == 0 {
|
||
return &[];
|
||
}
|
||
|
||
let actual_end = end_idx.min(snapshots.len());
|
||
let start_idx = actual_end.saturating_sub(window_size);
|
||
|
||
&snapshots[start_idx..actual_end]
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use data::providers::databento::mbp10::BidAskPair;
|
||
|
||
fn create_test_snapshot(timestamp: u64, bid_px: i64) -> Mbp10Snapshot {
|
||
let levels = vec![BidAskPair {
|
||
bid_px,
|
||
bid_sz: 100,
|
||
bid_ct: 5,
|
||
ask_px: bid_px + 1000000000, // 1 tick higher
|
||
ask_sz: 120,
|
||
ask_ct: 6,
|
||
}];
|
||
|
||
Mbp10Snapshot::new("ES.FUT".to_owned(), timestamp, levels, 0, 0)
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_exact_match() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
create_test_snapshot(4000, 150030000000000),
|
||
];
|
||
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2000, 2);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_between() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
];
|
||
|
||
// Target timestamp between 2000 and 3000
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2500, 2);
|
||
assert_eq!(window.len(), 1); // Only snapshot at 3000 remains
|
||
assert_eq!(window[0].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_empty() {
|
||
let snapshots: Vec<Mbp10Snapshot> = vec![];
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2000, 2);
|
||
assert_eq!(window.len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_before_all() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(2000, 150000000000000),
|
||
create_test_snapshot(3000, 150010000000000),
|
||
];
|
||
|
||
// Target is before all snapshots, should return first 2 snapshots
|
||
let window = get_snapshots_for_timestamp(&snapshots, 1000, 2);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_after_all() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
];
|
||
|
||
let window = get_snapshots_for_timestamp(&snapshots, 5000, 2);
|
||
assert_eq!(window.len(), 0); // No snapshots after target
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_window_clipping() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
];
|
||
|
||
// Request 10 snapshots but only 2 available at/after target (2000, 3000)
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2000, 10);
|
||
assert_eq!(window.len(), 2); // 2 snapshots at/after 2000: [2000, 3000]
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_full_window() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
create_test_snapshot(4000, 150030000000000),
|
||
];
|
||
|
||
let window = get_recent_snapshots(&snapshots, 3, 2);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_partial_window() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
];
|
||
|
||
// Request 5 snapshots but only 2 available
|
||
let window = get_recent_snapshots(&snapshots, 2, 5);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 1000);
|
||
assert_eq!(window[1].timestamp, 2000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_empty() {
|
||
let snapshots: Vec<Mbp10Snapshot> = vec![];
|
||
let window = get_recent_snapshots(&snapshots, 0, 2);
|
||
assert_eq!(window.len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_zero_end_idx() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
];
|
||
|
||
let window = get_recent_snapshots(&snapshots, 0, 2);
|
||
assert_eq!(window.len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_one_element() {
|
||
let snapshots = vec![create_test_snapshot(1000, 150000000000000)];
|
||
|
||
let window = get_recent_snapshots(&snapshots, 1, 1);
|
||
assert_eq!(window.len(), 1);
|
||
assert_eq!(window[0].timestamp, 1000);
|
||
}
|
||
}
|