gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Download 360 DBN files (36.3 MB) using Rust databento client

Browse Source 
- Created data/examples/download_ml_training_data.rs using reqwest + Databento HTTP API
- Downloaded 90 days × 4 symbols (ES.FUT, NQ.FUT, ZN.FUT, 6E.FUT)
- Files saved to test_data/real/databento/ml_training/
- Total: 360 files, 15 MB compressed DBN format
- Used existing Rust pattern from download_nq_fut.rs
- API key loaded from .env file
- 100% success rate (360/360 files)
- Ready for ML training benchmarks

Next: Create simplified training benchmark for RTX 3050 Ti GPU measurements
This commit is contained in:
jgrusewski
2025-10-13 13:30:02 +02:00
parent 08821565d6
commit e8a68ee39f
539 changed files with 1070448 additions and 489 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
data/src/lib.rs
View File

@@ -144,6 +144,7 @@ pub mod error;
pub mod features; // Feature engineering for ML models
pub mod parquet_persistence; // Parquet market data persistence for replay
pub mod providers; // Data providers (Databento, Benzinga)
pub mod replay; // Real market data replay infrastructure for tests/benchmarks
pub mod storage;
pub mod training_pipeline; // Training data pipeline for ML models
pub mod types;

146
data/src/parquet_persistence.rs
View File

@@ -4,10 +4,10 @@
//! and trade replay capabilities in the Foxhunt HFT system.
use anyhow::{Context, Result};
use arrow::array::{Float64Array, StringArray, TimestampNanosecondArray, UInt64Array};
use arrow::array::{Array, Float64Array, StringArray, TimestampNanosecondArray, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use arrow::record_batch::RecordBatch;
use parquet::arrow::ArrowWriter;
use parquet::arrow::{arrow_reader::ParquetRecordBatchReaderBuilder, ArrowWriter};
use parquet::file::properties::{EnabledStatistics, WriterProperties};
use serde::{Deserialize, Serialize};
use std::fs::File;
@@ -15,7 +15,7 @@ use std::path::Path;
use std::sync::Arc;
use tokio::sync::{mpsc, RwLock};
use tokio::time::{Duration, Instant};
use tracing::{debug, error, info, warn};
use tracing::{debug, error, info};
// Import the Parquet-specific market data event from trading_engine
pub use trading_engine::types::metrics::ParquetMarketDataEvent as MarketDataEvent;
@@ -197,6 +197,9 @@ impl ParquetMarketDataWriter {
Field::new("quantity", DataType::Float64, true),
Field::new("sequence", DataType::UInt64, false),
Field::new("latency_ns", DataType::UInt64, true),
Field::new("open", DataType::Float64, true),
Field::new("high", DataType::Float64, true),
Field::new("low", DataType::Float64, true),
]));
// Convert events to Arrow arrays
@@ -260,6 +263,9 @@ impl ParquetMarketDataWriter {
let mut quantities = Vec::with_capacity(len);
let mut sequences = Vec::with_capacity(len);
let mut latencies = Vec::with_capacity(len);
let mut opens = Vec::with_capacity(len);
let mut highs = Vec::with_capacity(len);
let mut lows = Vec::with_capacity(len);
for event in events {
timestamps.push(Some(event.timestamp_ns as i64));
@@ -271,6 +277,9 @@ impl ParquetMarketDataWriter {
quantities.push(event.quantity);
sequences.push(event.sequence);
latencies.push(event.latency_ns);
opens.push(event.open);
highs.push(event.high);
lows.push(event.low);
}
// Create Arrow arrays matching ParquetMarketDataEvent schema
@@ -282,6 +291,9 @@ impl ParquetMarketDataWriter {
let quantity_array = Float64Array::from(quantities);
let sequence_array = UInt64Array::from(sequences);
let latency_array = UInt64Array::from(latencies);
let open_array = Float64Array::from(opens);
let high_array = Float64Array::from(highs);
let low_array = Float64Array::from(lows);
// Create record batch
RecordBatch::try_new(
@@ -295,6 +307,9 @@ impl ParquetMarketDataWriter {
Arc::new(quantity_array),
Arc::new(sequence_array),
Arc::new(latency_array),
Arc::new(open_array),
Arc::new(high_array),
Arc::new(low_array),
],
)
.context("Failed to create Arrow RecordBatch")
@@ -349,10 +364,124 @@ impl ParquetMarketDataReader {
pub async fn read_file(&self, filename: &str) -> Result<Vec<MarketDataEvent>> {
let filepath = Path::new(&self.base_path).join(filename);
// This would be implemented using parquet::arrow::async_reader
// For now, return placeholder
warn!("Parquet reader not fully implemented yet: {:?}", filepath);
Ok(Vec::new())
debug!("Reading Parquet file: {:?}", filepath);
// Open the Parquet file
let file = File::open(&filepath)
.with_context(|| format!("Failed to open Parquet file: {:?}", filepath))?;
// Create Parquet reader
let builder = ParquetRecordBatchReaderBuilder::try_new(file)
.with_context(|| format!("Failed to create Parquet reader for: {:?}", filepath))?;
let reader = builder.build()
.context("Failed to build Parquet record batch reader")?;
let mut events = Vec::new();
// Read all record batches
for batch_result in reader {
let batch = batch_result.context("Failed to read record batch from Parquet")?;
// Extract columns
let timestamps = batch.column(0).as_any().downcast_ref::<TimestampNanosecondArray>()
.context("Failed to cast timestamp column")?;
let symbols = batch.column(1).as_any().downcast_ref::<StringArray>()
.context("Failed to cast symbol column")?;
let venues = batch.column(2).as_any().downcast_ref::<StringArray>()
.context("Failed to cast venue column")?;
let event_types = batch.column(3).as_any().downcast_ref::<StringArray>()
.context("Failed to cast event_type column")?;
let prices = batch.column(4).as_any().downcast_ref::<Float64Array>()
.context("Failed to cast price column")?;
let quantities = batch.column(5).as_any().downcast_ref::<Float64Array>()
.context("Failed to cast quantity column")?;
let sequences = batch.column(6).as_any().downcast_ref::<UInt64Array>()
.context("Failed to cast sequence column")?;
let latencies = batch.column(7).as_any().downcast_ref::<UInt64Array>()
.context("Failed to cast latency column")?;
// Handle optional OHLC columns (not present in all files)
let opens = if batch.num_columns() > 8 {
Some(batch.column(8).as_any().downcast_ref::<Float64Array>()
.context("Failed to cast open column")?)
} else {
None
};
let highs = if batch.num_columns() > 9 {
Some(batch.column(9).as_any().downcast_ref::<Float64Array>()
.context("Failed to cast high column")?)
} else {
None
};
let lows = if batch.num_columns() > 10 {
Some(batch.column(10).as_any().downcast_ref::<Float64Array>()
.context("Failed to cast low column")?)
} else {
None
};
// Convert rows to MarketDataEvent structs
for i in 0..batch.num_rows() {
let timestamp_ns = if timestamps.is_null(i) {
0
} else {
timestamps.value(i) as u64
};
let symbol = if symbols.is_null(i) {
String::new()
} else {
symbols.value(i).to_string()
};
let venue = if venues.is_null(i) {
String::new()
} else {
venues.value(i).to_string()
};
let event_type_str = if event_types.is_null(i) {
"Trade"
} else {
event_types.value(i)
};
// Parse event type string back to enum
let event_type = match event_type_str {
"Trade" => trading_engine::types::metrics::MarketDataEventType::Trade,
"Quote" => trading_engine::types::metrics::MarketDataEventType::Quote,
"OrderBookUpdate" => trading_engine::types::metrics::MarketDataEventType::OrderBookUpdate,
_ => trading_engine::types::metrics::MarketDataEventType::Trade,
};
let price = if prices.is_null(i) { None } else { Some(prices.value(i)) };
let quantity = if quantities.is_null(i) { None } else { Some(quantities.value(i)) };
let sequence = sequences.value(i);
let latency_ns = if latencies.is_null(i) { None } else { Some(latencies.value(i)) };
let open = opens.and_then(|arr| if arr.is_null(i) { None } else { Some(arr.value(i)) });
let high = highs.and_then(|arr| if arr.is_null(i) { None } else { Some(arr.value(i)) });
let low = lows.and_then(|arr| if arr.is_null(i) { None } else { Some(arr.value(i)) });
events.push(MarketDataEvent {
timestamp_ns,
symbol,
venue,
event_type,
price,
quantity,
sequence,
latency_ns,
open,
high,
low,
});
}
}
info!("Successfully read {} events from {:?}", events.len(), filepath);
Ok(events)
}
}
@@ -393,6 +522,9 @@ mod tests {
quantity: Some(0.1),
sequence: 1,
latency_ns: Some(1000),
open: None,
high: None,
low: None,
};
let result = writer.record(event);

456
data/src/providers/databento/dbn_to_parquet_converter.rs Normal file
View File

@@ -0,0 +1,456 @@
//! Production-Ready DBN to Parquet Converter
//!
//! High-performance converter that transforms Databento binary format (DBN) files containing
//! OHLCV market data into Parquet format for backtesting and analytics. Maintains <1μs per
//! event processing target through streaming architecture and zero-copy operations.
//!
//! ## Features
//!
//! - **Streaming Processing**: Memory-efficient streaming for large files (O(batch_size) memory)
//! - **Zero-Copy Where Possible**: Leverages DbnParser's zero-copy deserialization
//! - **Comprehensive Error Handling**: Detailed error messages with recovery strategies
//! - **Performance Metrics**: Tracks throughput, latency, and conversion statistics
//! - **Production-Ready**: Proper logging, error recovery, and resource management
//!
//! ## Usage
//!
//! ```no_run
//! use data::providers::databento::DbnToParquetConverter;
//! use data::parquet_persistence::ParquetConfig;
//!
//! #[tokio::main]
//! async fn main() -> anyhow::Result<()> {
//! let config = ParquetConfig::default();
//! let mut converter = DbnToParquetConverter::new(config).await?;
//!
//! let report = converter.convert_file("ES.FUT_ohlcv-1m_2024-01-02.dbn").await?;
//! println!("Converted {} events in {:?}", report.events_processed, report.duration);
//!
//! Ok(())
//! }
//! ```
use crate::error::{DataError, Result};
use crate::parquet_persistence::{ParquetConfig, ParquetMarketDataWriter};
use crate::providers::databento::dbn_parser::{DbnParser, ProcessedMessage};
use anyhow::{Context, Result as AnyhowResult};
use common::Price;
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::time::{Duration, Instant};
use tokio::fs::File;
use tokio::io::AsyncReadExt;
use trading_engine::types::metrics::{MarketDataEventType, ParquetMarketDataEvent};
use tracing::{debug, error, info};
/// DBN to Parquet converter with streaming support
///
/// Converts Databento binary format (DBN) files to Parquet format using the existing
/// DbnParser and ParquetMarketDataWriter infrastructure. Maintains <1μs per event
/// processing target through efficient streaming and batching.
pub struct DbnToParquetConverter {
/// DBN parser for reading binary format
parser: DbnParser,
/// Parquet writer for output
writer: ParquetMarketDataWriter,
/// Conversion metrics tracker
metrics: ConversionMetrics,
/// Batch size for streaming processing
batch_size: usize,
}
impl DbnToParquetConverter {
/// Create new converter with specified Parquet configuration
///
/// # Arguments
/// * `config` - Parquet writer configuration (base path, compression, etc.)
///
/// # Returns
/// * `Result<Self>` - Configured converter or error
///
/// # Errors
/// Returns error if ParquetMarketDataWriter creation fails
pub async fn new(config: ParquetConfig) -> AnyhowResult<Self> {
let parser = DbnParser::new()
.map_err(|e| anyhow::anyhow!("Failed to create DBN parser: {}", e))?;
let writer = ParquetMarketDataWriter::new(config)
.await
.context("Failed to create Parquet writer")?;
Ok(Self {
parser,
writer,
metrics: ConversionMetrics::default(),
batch_size: 10000, // Process in 10k event batches
})
}
/// Convert a DBN file to Parquet format
///
/// Reads the DBN file in streaming fashion, converts OHLCV messages to Parquet events,
/// and writes them via the ParquetMarketDataWriter. Maintains <1μs per event target.
///
/// # Arguments
/// * `dbn_path` - Path to the DBN file to convert
///
/// # Returns
/// * `ConversionReport` - Statistics about the conversion process
///
/// # Errors
/// Returns error if file I/O fails, DBN parsing fails, or Parquet writing fails
pub async fn convert_file<P: AsRef<Path>>(
&mut self,
dbn_path: P,
) -> AnyhowResult<ConversionReport> {
let path = dbn_path.as_ref();
info!("Starting DBN to Parquet conversion: {:?}", path);
let start_time = Instant::now();
self.metrics = ConversionMetrics::default();
// Read entire file (for simplicity - could be optimized with mmap for huge files)
let mut file = File::open(path)
.await
.with_context(|| format!("Failed to open DBN file: {:?}", path))?;
let mut buffer = Vec::new();
file.read_to_end(&mut buffer)
.await
.with_context(|| format!("Failed to read DBN file: {:?}", path))?;
debug!("Read {} bytes from DBN file", buffer.len());
// Skip DBN file header/metadata section
// DBN files have: magic (4) + dataset name (variable) + metadata + data records
// We need to find the start of data records by scanning for valid message headers
let data_offset = self.find_data_start(&buffer)?;
debug!("Found data section at offset {}", data_offset);
// Parse DBN messages in batches (from data section only)
let messages = self.parser.parse_batch(&buffer[data_offset..])
.map_err(|e| anyhow::anyhow!("DBN parsing failed: {}", e))?;
info!("Parsed {} messages from DBN file", messages.len());
// Convert and write in batches
let mut batch = Vec::with_capacity(self.batch_size);
for message in messages {
match self.convert_message(message) {
Ok(Some(event)) => {
batch.push(event);
self.metrics.events_converted += 1;
// Write batch when full
if batch.len() >= self.batch_size {
self.write_batch(&mut batch).await?;
}
}
Ok(None) => {
// Non-OHLCV message, skip
self.metrics.events_skipped += 1;
}
Err(e) => {
error!("Failed to convert message: {}", e);
self.metrics.events_failed += 1;
// Continue processing other messages
}
}
}
// Write remaining batch
if !batch.is_empty() {
self.write_batch(&mut batch).await?;
}
let duration = start_time.elapsed();
let report = ConversionReport {
events_processed: self.metrics.events_converted,
events_skipped: self.metrics.events_skipped,
events_failed: self.metrics.events_failed,
duration,
throughput_events_per_sec: (self.metrics.events_converted as f64
/ duration.as_secs_f64()) as u64,
};
info!("Conversion complete: {} events in {:?} ({} events/sec)",
report.events_processed, report.duration, report.throughput_events_per_sec);
Ok(report)
}
/// Convert a single DBN message to Parquet event
///
/// Only converts OHLCV messages - other message types return None.
/// Maps DBN OHLCV fields to ParquetMarketDataEvent with proper field placement.
///
/// # Field Mapping
/// - `open` → `open`
/// - `high` → `high`
/// - `low` → `low`
/// - `close` → `price` (most important price for analytics)
/// - `volume` → `quantity`
/// - `ts_event` → `timestamp_ns`
/// - `symbol` → `symbol`
/// - `venue` → "DATABENTO" (default, DBN doesn't always provide venue)
///
/// # Arguments
/// * `message` - Parsed DBN message
///
/// # Returns
/// * `Ok(Some(event))` - Successfully converted OHLCV event
/// * `Ok(None)` - Non-OHLCV message (skipped)
/// * `Err(e)` - Conversion error
fn convert_message(&self, message: ProcessedMessage) -> Result<Option<ParquetMarketDataEvent>> {
match message {
ProcessedMessage::Ohlcv {
symbol,
timestamp,
open,
high,
low,
close,
volume,
} => {
// Convert hardware timestamp to nanoseconds
let timestamp_ns = timestamp.as_nanos();
// Convert Price to f64 (Price is a newtype around Decimal)
let open_f64 = price_to_f64(open)?;
let high_f64 = price_to_f64(high)?;
let low_f64 = price_to_f64(low)?;
let close_f64 = price_to_f64(close)?;
// Convert volume Decimal to f64
let volume_f64 = decimal_to_f64(volume)?;
Ok(Some(ParquetMarketDataEvent {
timestamp_ns,
symbol,
venue: "DATABENTO".to_string(), // Default venue for DBN files
event_type: MarketDataEventType::Ohlcv,
price: Some(close_f64), // Close price in standard price field
quantity: Some(volume_f64), // Volume in quantity field
sequence: 0, // OHLCV bars don't have sequence numbers
latency_ns: None, // No latency measurement for historical data
open: Some(open_f64),
high: Some(high_f64),
low: Some(low_f64),
}))
}
// Non-OHLCV messages are skipped
_ => Ok(None),
}
}
/// Write a batch of events to Parquet
///
/// Writes accumulated events and clears the batch buffer.
///
/// # Arguments
/// * `batch` - Mutable reference to batch vector (will be cleared)
///
/// # Errors
/// Returns error if Parquet writer fails
async fn write_batch(&self, batch: &mut Vec<ParquetMarketDataEvent>) -> AnyhowResult<()> {
for event in batch.drain(..) {
self.writer.record(event)
.context("Failed to record event to Parquet writer")?;
}
Ok(())
}
/// Find the start of data records in a DBN file
///
/// DBN files contain:
/// 1. Magic bytes "DBN" + version (4 bytes)
/// 2. Dataset name + metadata (variable length)
/// 3. Data records (what we want)
///
/// This method scans the file to find where OHLCV data records begin.
/// We specifically look for 0x42 (OHLCV Bar) record type since this converter
/// is designed for OHLCV files.
///
/// # Arguments
/// * `buffer` - Complete file contents
///
/// # Returns
/// Offset where OHLCV data records begin
///
/// # Errors
/// Returns error if no valid OHLCV records found
fn find_data_start(&self, buffer: &[u8]) -> AnyhowResult<usize> {
// Skip first 4 bytes (magic + version)
if buffer.len() < 4 || &buffer[0..3] != b"DBN" {
return Err(anyhow::anyhow!("Invalid DBN file: missing magic bytes"));
}
// OHLCV message structure (from dbn_parser.rs DbnOhlcvMessage):
// - header (16 bytes): length(2) + rtype(1) + publisher_id(1) + instrument_id(4) + ts_event(8)
// - open (8 bytes)
// - high (8 bytes)
// - low (8 bytes)
// - close (8 bytes)
// - volume (8 bytes)
// Total: 56 bytes
const OHLCV_RECORD_SIZE: usize = 56;
const OHLCV_RTYPE: u8 = 0x42; // 'B' for Bar
let mut offset = 4; // Start after magic bytes
while offset + OHLCV_RECORD_SIZE <= buffer.len() {
// Read length as little-endian u16
let length = u16::from_le_bytes([buffer[offset], buffer[offset + 1]]);
let rtype = buffer[offset + 2];
// Check if this is an OHLCV record
if rtype == OHLCV_RTYPE && length == OHLCV_RECORD_SIZE as u16 {
// Validate that the next few records are also OHLCV
// (to avoid false positives from metadata)
let mut valid_count = 0;
let mut check_offset = offset;
for _ in 0..3 {
if check_offset + OHLCV_RECORD_SIZE > buffer.len() {
break;
}
let check_length = u16::from_le_bytes([
buffer[check_offset],
buffer[check_offset + 1]
]);
let check_rtype = buffer[check_offset + 2];
if check_rtype == OHLCV_RTYPE && check_length == OHLCV_RECORD_SIZE as u16 {
valid_count += 1;
check_offset += OHLCV_RECORD_SIZE;
} else {
break;
}
}
// If we found at least 2 consecutive OHLCV records, we're in the data section
if valid_count >= 2 {
return Ok(offset);
}
}
// Move to next byte
offset += 1;
}
Err(anyhow::anyhow!("Could not find OHLCV data records in DBN file"))
}
}
/// Convert Price to f64
fn price_to_f64(price: Price) -> Result<f64> {
// Price has a to_f64() method
Ok(price.to_f64())
}
/// Convert Decimal to f64
fn decimal_to_f64(decimal: Decimal) -> Result<f64> {
decimal.to_string().parse::<f64>()
.map_err(|e| DataError::Conversion(format!("Failed to convert Decimal to f64: {}", e)))
}
/// Conversion metrics tracker
#[derive(Debug, Clone, Default)]
struct ConversionMetrics {
/// Total events successfully converted
events_converted: u64,
/// Events skipped (non-OHLCV messages)
events_skipped: u64,
/// Events that failed to convert
events_failed: u64,
}
/// Conversion report with statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversionReport {
/// Total OHLCV events successfully processed
pub events_processed: u64,
/// Non-OHLCV events skipped
pub events_skipped: u64,
/// Events that failed conversion
pub events_failed: u64,
/// Total conversion duration
pub duration: Duration,
/// Throughput in events per second
pub throughput_events_per_sec: u64,
}
impl ConversionReport {
/// Check if conversion was successful (no failures)
pub fn is_success(&self) -> bool {
self.events_failed == 0
}
/// Get success rate as percentage
pub fn success_rate(&self) -> f64 {
let total = self.events_processed + self.events_failed;
if total == 0 {
100.0
} else {
(self.events_processed as f64 / total as f64) * 100.0
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[tokio::test]
async fn test_converter_creation() {
let temp_dir = tempdir().unwrap();
let config = ParquetConfig {
base_path: temp_dir.path().to_string_lossy().to_string(),
..Default::default()
};
let converter = DbnToParquetConverter::new(config).await;
assert!(converter.is_ok());
}
#[test]
fn test_conversion_report_success_rate() {
let report = ConversionReport {
events_processed: 95,
events_skipped: 5,
events_failed: 5,
duration: Duration::from_secs(1),
throughput_events_per_sec: 95,
};
assert_eq!(report.success_rate(), 95.0);
assert!(!report.is_success());
}
#[test]
fn test_conversion_report_perfect_success() {
let report = ConversionReport {
events_processed: 100,
events_skipped: 0,
events_failed: 0,
duration: Duration::from_secs(1),
throughput_events_per_sec: 100,
};
assert_eq!(report.success_rate(), 100.0);
assert!(report.is_success());
}
}

4
data/src/providers/databento/mod.rs
View File

@@ -70,11 +70,15 @@
// Module declarations
pub mod client;
pub mod dbn_parser;
pub mod dbn_to_parquet_converter;
pub mod parser;
pub mod stream;
pub mod types;
pub mod websocket_client;
// Re-export converter for convenience
pub use dbn_to_parquet_converter::{DbnToParquetConverter, ConversionReport};
// Import all major components
// DO NOT RE-EXPORT - Use explicit imports at usage sites
// pub use crate::providers::databento::{

478
data/src/replay/market_data_streamer.rs Normal file
View File

@@ -0,0 +1,478 @@
//! Real-time streaming of historical market data
//!
//! Provides time-accurate replay of historical market data for backtesting
//! and simulation. Maintains the original timing relationships between events
//! while allowing speed multipliers for faster simulation.
//!
//! # Features
//!
//! - **Time-accurate replay**: Preserves original inter-event delays
//! - **Speed control**: 1x, 10x, 100x or custom speed multipliers
//! - **High throughput**: Async channel-based streaming
//! - **Backpressure handling**: Bounded channels prevent memory exhaustion
//!
//! # Examples
//!
//! ## Real-time replay (1x speed)
//! ```no_run
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("market_data.parquet");
//! let events = loader.load_all().await?;
//!
//! let streamer = MarketDataStreamer::new(events);
//! let mut rx = streamer.stream().await;
//!
//! while let Some(event) = rx.recv().await {
//! println!("Event at {}: {:?}", event.timestamp_ns, event.event_type);
//! }
//! # Ok(())
//! # }
//! ```
//!
//! ## Fast replay (10x speed)
//! ```no_run
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("market_data.parquet");
//! let events = loader.load_all().await?;
//!
//! let streamer = MarketDataStreamer::new(events)
//! .with_speed(10.0) // 10x faster
//! .with_buffer_size(10000); // Larger buffer for high throughput
//!
//! let mut rx = streamer.stream().await;
//!
//! while let Some(event) = rx.recv().await {
//! // Process events at 10x speed
//! }
//! # Ok(())
//! # }
//! ```
//!
//! ## Instant replay (no delays)
//! ```no_run
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("market_data.parquet");
//! let events = loader.load_all().await?;
//!
//! // Speed = f64::INFINITY means no delays
//! let streamer = MarketDataStreamer::new(events)
//! .with_speed(f64::INFINITY);
//!
//! let mut rx = streamer.stream().await;
//! // Events streamed as fast as possible
//! # Ok(())
//! # }
//! ```
use std::time::Duration;
use tokio::sync::mpsc;
use trading_engine::types::metrics::ParquetMarketDataEvent;
/// Market data streamer for time-accurate replay
///
/// Streams historical market data events while preserving the original
/// timing relationships. Supports speed multipliers for faster simulation
/// and configurable channel buffer sizes for backpressure control.
#[derive(Debug)]
pub struct MarketDataStreamer {
events: Vec<ParquetMarketDataEvent>,
replay_speed: f64,
buffer_size: usize,
}
impl MarketDataStreamer {
/// Create new streamer with events
///
/// Default configuration:
/// - Speed: 1.0 (real-time)
/// - Buffer size: 1000 events
///
/// # Arguments
/// * `events` - Vector of market data events to stream (must be sorted by timestamp)
///
/// # Examples
/// ```
/// use data::replay::MarketDataStreamer;
///
/// let events = vec![]; // Your events here
/// let streamer = MarketDataStreamer::new(events);
/// ```
pub fn new(events: Vec<ParquetMarketDataEvent>) -> Self {
Self {
events,
replay_speed: 1.0,
buffer_size: 1000,
}
}
/// Set replay speed multiplier
///
/// - 1.0 = real-time (original timing)
/// - 2.0 = 2x faster
/// - 0.5 = half speed (slow motion)
/// - f64::INFINITY = instant (no delays)
///
/// # Arguments
/// * `speed` - Speed multiplier (must be > 0.0)
///
/// # Panics
/// Panics if speed <= 0.0 or is NaN
///
/// # Examples
/// ```
/// use data::replay::MarketDataStreamer;
///
/// let events = vec![];
/// let streamer = MarketDataStreamer::new(events)
/// .with_speed(10.0); // 10x faster
/// ```
pub fn with_speed(mut self, speed: f64) -> Self {
assert!(speed > 0.0 && !speed.is_nan(), "Speed must be positive");
self.replay_speed = speed;
self
}
/// Set channel buffer size
///
/// Controls backpressure behavior:
/// - Small buffer (100-1000): More memory efficient, may slow producer
/// - Large buffer (10000+): Higher throughput, more memory usage
///
/// # Arguments
/// * `size` - Buffer size in events
///
/// # Examples
/// ```
/// use data::replay::MarketDataStreamer;
///
/// let events = vec![];
/// let streamer = MarketDataStreamer::new(events)
/// .with_buffer_size(10000); // Large buffer for high throughput
/// ```
pub fn with_buffer_size(mut self, size: usize) -> Self {
self.buffer_size = size;
self
}
/// Stream events with time-accurate delays
///
/// Spawns a background task that sends events through an async channel
/// with delays calculated from timestamp differences and replay speed.
///
/// # Returns
/// Receiver side of the event channel
///
/// # Timing Behavior
/// - First event: Sent immediately
/// - Subsequent events: Delayed by (timestamp_diff / replay_speed)
/// - Speed = f64::INFINITY: No delays (instant replay)
///
/// # Backpressure
/// If receiver is slow, producer will block when channel buffer fills.
/// This prevents unbounded memory growth.
///
/// # Examples
/// ```no_run
/// # use data::replay::MarketDataStreamer;
/// # async fn example() -> anyhow::Result<()> {
/// let events = vec![]; // Your events
/// let streamer = MarketDataStreamer::new(events);
/// let mut rx = streamer.stream().await;
///
/// while let Some(event) = rx.recv().await {
/// println!("Event: {:?}", event);
/// }
/// # Ok(())
/// # }
/// ```
pub async fn stream(self) -> mpsc::Receiver<ParquetMarketDataEvent> {
let (tx, rx) = mpsc::channel(self.buffer_size);
tokio::spawn(async move {
let mut prev_timestamp = None;
for event in self.events {
// Calculate delay based on timestamp difference
if let Some(prev_ts) = prev_timestamp {
let delay_ns = event.timestamp_ns.saturating_sub(prev_ts);
// Only sleep if replay_speed is finite
if self.replay_speed.is_finite() {
let delay_secs = (delay_ns as f64 / 1_000_000_000.0) / self.replay_speed;
// Only sleep for meaningful delays (> 1 microsecond)
if delay_secs > 0.000_001 {
tokio::time::sleep(Duration::from_secs_f64(delay_secs)).await;
}
}
// If replay_speed is infinite, no sleep (instant replay)
}
prev_timestamp = Some(event.timestamp_ns);
// Send event to channel
if tx.send(event).await.is_err() {
// Receiver dropped, stop streaming
break;
}
}
// Channel will close when tx is dropped here
});
rx
}
/// Get event count
///
/// Returns the number of events that will be streamed.
///
/// # Examples
/// ```
/// use data::replay::MarketDataStreamer;
///
/// let events = vec![]; // Your events
/// let streamer = MarketDataStreamer::new(events);
/// println!("Will stream {} events", streamer.event_count());
/// ```
pub fn event_count(&self) -> usize {
self.events.len()
}
/// Get time span covered by events
///
/// Returns the time difference between first and last event in nanoseconds.
/// Returns None if there are fewer than 2 events.
///
/// # Examples
/// ```
/// use data::replay::MarketDataStreamer;
///
/// let events = vec![]; // Your events (sorted by timestamp)
/// let streamer = MarketDataStreamer::new(events);
///
/// if let Some(span_ns) = streamer.time_span_ns() {
/// println!("Events span {} seconds", span_ns as f64 / 1e9);
/// }
/// ```
pub fn time_span_ns(&self) -> Option<u64> {
if self.events.len() < 2 {
return None;
}
let first = self.events.first()?.timestamp_ns;
let last = self.events.last()?.timestamp_ns;
Some(last.saturating_sub(first))
}
/// Estimate replay duration
///
/// Calculates how long the replay will take at current speed.
/// Returns None if there are fewer than 2 events or speed is infinite.
///
/// # Returns
/// Estimated replay duration in seconds
///
/// # Examples
/// ```
/// use data::replay::MarketDataStreamer;
///
/// let events = vec![]; // Your events
/// let streamer = MarketDataStreamer::new(events)
/// .with_speed(10.0);
///
/// if let Some(duration_secs) = streamer.estimate_replay_duration_secs() {
/// println!("Replay will take ~{:.2} seconds", duration_secs);
/// }
/// ```
pub fn estimate_replay_duration_secs(&self) -> Option<f64> {
if !self.replay_speed.is_finite() {
return None; // Infinite speed = instant replay
}
let span_ns = self.time_span_ns()?;
Some((span_ns as f64 / 1_000_000_000.0) / self.replay_speed)
}
}
#[cfg(test)]
mod tests {
use super::*;
use trading_engine::types::metrics::MarketDataEventType;
fn create_test_event(timestamp_ns: u64, symbol: &str) -> ParquetMarketDataEvent {
ParquetMarketDataEvent {
timestamp_ns,
symbol: symbol.to_string(),
venue: "test_venue".to_string(),
event_type: MarketDataEventType::Trade,
price: Some(100.0),
quantity: Some(10.0),
sequence: 0,
latency_ns: None,
open: None,
high: None,
low: None,
}
}
#[test]
fn test_market_data_streamer_creation() {
let events = vec![];
let streamer = MarketDataStreamer::new(events);
assert_eq!(streamer.replay_speed, 1.0);
assert_eq!(streamer.buffer_size, 1000);
}
#[test]
fn test_with_speed() {
let events = vec![];
let streamer = MarketDataStreamer::new(events).with_speed(2.0);
assert_eq!(streamer.replay_speed, 2.0);
}
#[test]
fn test_with_buffer_size() {
let events = vec![];
let streamer = MarketDataStreamer::new(events).with_buffer_size(5000);
assert_eq!(streamer.buffer_size, 5000);
}
#[test]
#[should_panic(expected = "Speed must be positive")]
fn test_invalid_speed_zero() {
let events = vec![];
let _ = MarketDataStreamer::new(events).with_speed(0.0);
}
#[test]
#[should_panic(expected = "Speed must be positive")]
fn test_invalid_speed_negative() {
let events = vec![];
let _ = MarketDataStreamer::new(events).with_speed(-1.0);
}
#[test]
fn test_event_count() {
let events = vec![
create_test_event(1000, "BTC/USD"),
create_test_event(2000, "BTC/USD"),
create_test_event(3000, "BTC/USD"),
];
let streamer = MarketDataStreamer::new(events);
assert_eq!(streamer.event_count(), 3);
}
#[test]
fn test_time_span_ns() {
let events = vec![
create_test_event(1000, "BTC/USD"),
create_test_event(2000, "BTC/USD"),
create_test_event(5000, "BTC/USD"),
];
let streamer = MarketDataStreamer::new(events);
assert_eq!(streamer.time_span_ns(), Some(4000));
}
#[test]
fn test_time_span_ns_insufficient_events() {
let events = vec![create_test_event(1000, "BTC/USD")];
let streamer = MarketDataStreamer::new(events);
assert_eq!(streamer.time_span_ns(), None);
}
#[test]
fn test_estimate_replay_duration() {
let events = vec![
create_test_event(0, "BTC/USD"),
create_test_event(10_000_000_000, "BTC/USD"), // 10 seconds later
];
let streamer = MarketDataStreamer::new(events).with_speed(2.0);
let duration = streamer.estimate_replay_duration_secs();
assert_eq!(duration, Some(5.0)); // 10 seconds / 2.0 speed = 5 seconds
}
#[test]
fn test_estimate_replay_duration_infinite_speed() {
let events = vec![
create_test_event(0, "BTC/USD"),
create_test_event(10_000_000_000, "BTC/USD"),
];
let streamer = MarketDataStreamer::new(events).with_speed(f64::INFINITY);
assert_eq!(streamer.estimate_replay_duration_secs(), None);
}
#[tokio::test]
async fn test_stream_basic() {
let events = vec![
create_test_event(1000, "BTC/USD"),
create_test_event(2000, "ETH/USD"),
create_test_event(3000, "SOL/USD"),
];
let streamer = MarketDataStreamer::new(events).with_speed(f64::INFINITY); // Instant replay
let mut rx = streamer.stream().await;
let mut received = vec![];
while let Some(event) = rx.recv().await {
received.push(event);
}
assert_eq!(received.len(), 3);
assert_eq!(received[0].symbol, "BTC/USD");
assert_eq!(received[1].symbol, "ETH/USD");
assert_eq!(received[2].symbol, "SOL/USD");
}
#[tokio::test]
async fn test_stream_respects_timing() {
let events = vec![
create_test_event(0, "BTC/USD"),
create_test_event(100_000_000, "BTC/USD"), // 100ms later
];
let streamer = MarketDataStreamer::new(events).with_speed(1.0);
let start = tokio::time::Instant::now();
let mut rx = streamer.stream().await;
let mut count = 0;
while let Some(_event) = rx.recv().await {
count += 1;
}
let elapsed = start.elapsed();
assert_eq!(count, 2);
// Should take at least 90ms (allowing some slack for timing)
assert!(elapsed.as_millis() >= 90);
}
#[tokio::test]
async fn test_stream_early_drop() {
let events = vec![
create_test_event(1000, "BTC/USD"),
create_test_event(2000, "ETH/USD"),
create_test_event(3000, "SOL/USD"),
];
let streamer = MarketDataStreamer::new(events).with_speed(f64::INFINITY);
let mut rx = streamer.stream().await;
// Receive only first event, then drop receiver
let first = rx.recv().await;
assert!(first.is_some());
assert_eq!(first.unwrap().symbol, "BTC/USD");
drop(rx); // Producer should detect and stop
// Test passes if no panic/hang occurs
}
}

241
data/src/replay/mod.rs Normal file
View File

@@ -0,0 +1,241 @@
//! Real market data replay infrastructure
//!
//! This module provides comprehensive infrastructure for loading and replaying
//! real market data from Parquet files. It's designed for use across multiple
//! contexts: tests, benchmarks, backtesting, and ML training.
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────┐
//! │ Parquet Files │
//! │ (Historical Data) │
//! └──────────┬──────────┘
//! │
//! ▼
//! ┌─────────────────────┐
//! │ ParquetDataLoader │ ← Load data from Parquet
//! │ - Bulk loading │
//! │ - Streaming │
//! └──────────┬──────────┘
//! │
//! ▼
//! ┌─────────────────────┐
//! │ MarketDataStreamer │ ← Time-accurate replay
//! │ - Speed control │
//! │ - Timing accuracy │
//! └──────────┬──────────┘
//! │
//! ▼
//! ┌─────────────────────┐
//! │ Consumers │
//! │ - Tests │
//! │ - Benchmarks │
//! │ - Backtesting │
//! │ - ML Training │
//! └─────────────────────┘
//! ```
//!
//! # Core Components
//!
//! ## ParquetDataLoader
//!
//! Efficient Parquet file loading with two modes:
//! - **Bulk loading**: Load entire file into memory for random access
//! - **Streaming**: Process events batch-by-batch for memory efficiency
//!
//! ## MarketDataStreamer
//!
//! Time-accurate replay of historical data with:
//! - Preserved timing relationships between events
//! - Configurable speed multipliers (1x, 10x, 100x, etc.)
//! - Backpressure-safe async channels
//!
//! # Usage Examples
//!
//! ## Basic: Load and Stream Real Data
//!
//! ```no_run
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! # async fn example() -> anyhow::Result<()> {
//! // Load market data from Parquet file
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let events = loader.load_all().await?;
//!
//! println!("Loaded {} events", events.len());
//!
//! // Stream events at real-time speed
//! let streamer = MarketDataStreamer::new(events);
//! let mut rx = streamer.stream().await;
//!
//! while let Some(event) = rx.recv().await {
//! println!("Event: {:?} at {}", event.event_type, event.timestamp_ns);
//! }
//! # Ok(())
//! # }
//! ```
//!
//! ## Fast Backtesting: 10x Speed
//!
//! ```no_run
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let events = loader.load_all().await?;
//!
//! // 10x faster replay for quick backtesting
//! let streamer = MarketDataStreamer::new(events)
//! .with_speed(10.0)
//! .with_buffer_size(10000);
//!
//! let mut rx = streamer.stream().await;
//!
//! // Process events 10x faster than real-time
//! while let Some(event) = rx.recv().await {
//! // Your backtest logic here
//! }
//! # Ok(())
//! # }
//! ```
//!
//! ## ML Training: Instant Replay (No Delays)
//!
//! ```no_run
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let events = loader.load_all().await?;
//!
//! // Instant replay (no timing delays)
//! let streamer = MarketDataStreamer::new(events)
//! .with_speed(f64::INFINITY);
//!
//! let mut rx = streamer.stream().await;
//!
//! // Process events as fast as possible for ML training
//! while let Some(event) = rx.recv().await {
//! // Feature extraction, model training, etc.
//! }
//! # Ok(())
//! # }
//! ```
//!
//! ## Memory-Efficient Streaming
//!
//! ```no_run
//! use data::replay::ParquetDataLoader;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("large_dataset.parquet");
//! let mut stream = loader.load_stream().await?;
//!
//! // Process batches without loading entire file into memory
//! while let Some(batch) = stream.next().await {
//! let events = batch?;
//! println!("Processing batch of {} events", events.len());
//! // Process batch...
//! }
//! # Ok(())
//! # }
//! ```
//!
//! ## Integration with Benchmarks
//!
//! ```no_run
//! use criterion::{criterion_group, criterion_main, Criterion};
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! async fn load_test_data() -> Vec<trading_engine::types::metrics::ParquetMarketDataEvent> {
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! loader.load_all().await.expect("Failed to load test data")
//! }
//!
//! fn benchmark_strategy(c: &mut Criterion) {
//! let rt = tokio::runtime::Runtime::new().unwrap();
//! let events = rt.block_on(load_test_data());
//!
//! c.bench_function("strategy_with_real_data", |b| {
//! b.to_async(&rt).iter(|| async {
//! let streamer = MarketDataStreamer::new(events.clone())
//! .with_speed(f64::INFINITY); // Instant for benchmarking
//! let mut rx = streamer.stream().await;
//!
//! while let Some(event) = rx.recv().await {
//! // Your strategy logic
//! }
//! });
//! });
//! }
//!
//! criterion_group!(benches, benchmark_strategy);
//! criterion_main!(benches);
//! ```
//!
//! ## Integration with Tests
//!
//! ```no_run
//! #[cfg(test)]
//! mod tests {
//! use data::replay::{ParquetDataLoader, MarketDataStreamer};
//!
//! #[tokio::test]
//! async fn test_strategy_with_real_data() {
//! // Load real market data
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let events = loader.load_all().await.expect("Failed to load test data");
//!
//! // Stream at high speed for tests
//! let streamer = MarketDataStreamer::new(events)
//! .with_speed(100.0); // 100x faster
//!
//! let mut rx = streamer.stream().await;
//!
//! let mut processed_count = 0;
//! while let Some(event) = rx.recv().await {
//! // Test your strategy logic
//! processed_count += 1;
//! }
//!
//! assert!(processed_count > 0, "Should process events");
//! }
//! }
//! ```
//!
//! # Performance Characteristics
//!
//! ## ParquetDataLoader
//! - **Bulk loading**: O(n) time, O(n) memory
//! - **Streaming**: O(1) memory per batch, amortized O(n) time
//! - **Throughput**: ~1-5M events/sec depending on schema complexity
//!
//! ## MarketDataStreamer
//! - **Latency**: < 1μs per event (excluding intentional delays)
//! - **Throughput**: Limited by channel buffer and consumer speed
//! - **Memory**: O(buffer_size) for channel, O(1) for streamer
//!
//! # File Format
//!
//! Expected Parquet schema (all fields from ParquetMarketDataEvent):
//! - `timestamp_ns`: Timestamp (nanoseconds) - REQUIRED
//! - `symbol`: String - REQUIRED
//! - `venue`: String - REQUIRED
//! - `event_type`: String (Trade/Quote/OrderBookUpdate/StatusUpdate/Ohlcv) - REQUIRED
//! - `sequence`: UInt64 - REQUIRED
//! - `price`: Float64 - OPTIONAL
//! - `quantity`: Float64 - OPTIONAL
//! - `latency_ns`: UInt64 - OPTIONAL
//! - `open`: Float64 - OPTIONAL (OHLCV only)
//! - `high`: Float64 - OPTIONAL (OHLCV only)
//! - `low`: Float64 - OPTIONAL (OHLCV only)
pub mod market_data_streamer;
pub mod parquet_loader;
pub use market_data_streamer::MarketDataStreamer;
pub use parquet_loader::{ParquetDataLoader, ParquetEventStream};
// Re-export the event type for convenience
pub use trading_engine::types::metrics::ParquetMarketDataEvent;

334
data/src/replay/parquet_loader.rs Normal file
View File

@@ -0,0 +1,334 @@
//! Parquet data loader for real market data replay
//!
//! Provides efficient loading of real market data from Parquet files
//! for use in tests, benchmarks, backtesting, and ML training.
//!
//! # Architecture
//!
//! The loader provides two main APIs:
//! - **Bulk Loading**: Load entire Parquet file into memory (`load_all`)
//! - **Streaming**: Load events in batches for memory efficiency (`load_stream`)
//!
//! # Examples
//!
//! ## Load all events from Parquet file
//! ```no_run
//! use data::replay::ParquetDataLoader;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let events = loader.load_all().await?;
//! println!("Loaded {} events", events.len());
//! # Ok(())
//! # }
//! ```
//!
//! ## Stream events for memory efficiency
//! ```no_run
//! use data::replay::ParquetDataLoader;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let mut stream = loader.load_stream().await?;
//!
//! while let Some(batch) = stream.next().await {
//! let batch = batch?;
//! println!("Processing batch of {} events", batch.len());
//! }
//! # Ok(())
//! # }
//! ```
use anyhow::{Context, Result};
use arrow::array::{
Array, Float64Array, StringArray, as_primitive_array,
};
use arrow::datatypes::{TimestampNanosecondType, UInt64Type};
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::fs::File;
use std::path::Path;
use trading_engine::types::metrics::{MarketDataEventType, ParquetMarketDataEvent};
/// Parquet data loader for real market data
///
/// Provides efficient loading capabilities for Parquet files containing
/// historical market data. Supports both bulk loading (all events at once)
/// and streaming (batch-by-batch processing).
pub struct ParquetDataLoader {
file_path: String,
}
impl ParquetDataLoader {
/// Create new loader for a Parquet file
///
/// # Arguments
/// * `file_path` - Path to the Parquet file to load
///
/// # Examples
/// ```
/// use data::replay::ParquetDataLoader;
///
/// let loader = ParquetDataLoader::new("market_data.parquet");
/// ```
pub fn new(file_path: impl AsRef<Path>) -> Self {
Self {
file_path: file_path.as_ref().to_string_lossy().to_string(),
}
}
/// Load all events from Parquet file into memory
///
/// Loads the entire Parquet file into a vector. Use this for smaller files
/// or when you need random access to all events.
///
/// # Returns
/// * `Ok(Vec<ParquetMarketDataEvent>)` - All events from the file
/// * `Err(_)` - If file cannot be read or parsed
///
/// # Examples
/// ```no_run
/// # use data::replay::ParquetDataLoader;
/// # async fn example() -> anyhow::Result<()> {
/// let loader = ParquetDataLoader::new("market_data.parquet");
/// let events = loader.load_all().await?;
/// println!("Loaded {} events", events.len());
/// # Ok(())
/// # }
/// ```
pub async fn load_all(&self) -> Result<Vec<ParquetMarketDataEvent>> {
// Read from file synchronously (Parquet doesn't have async reader)
let file = File::open(&self.file_path)
.with_context(|| format!("Failed to open Parquet file: {}", self.file_path))?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file)
.context("Failed to create Parquet reader")?;
let reader = builder.build().context("Failed to build Parquet reader")?;
let mut events = Vec::new();
for batch_result in reader {
let batch = batch_result.context("Failed to read record batch")?;
events.extend(Self::batch_to_events(&batch)?);
}
Ok(events)
}
/// Load events in streaming fashion (iterator-based)
///
/// Returns a stream of event batches for memory-efficient processing.
/// Each batch contains multiple events from a single RecordBatch.
///
/// # Returns
/// * `Ok(ParquetEventStream)` - Stream of event batches
/// * `Err(_)` - If file cannot be opened
///
/// # Examples
/// ```no_run
/// # use data::replay::ParquetDataLoader;
/// # async fn example() -> anyhow::Result<()> {
/// let loader = ParquetDataLoader::new("market_data.parquet");
/// let mut stream = loader.load_stream().await?;
///
/// while let Some(batch) = stream.next().await {
/// let events = batch?;
/// // Process events...
/// }
/// # Ok(())
/// # }
/// ```
pub async fn load_stream(&self) -> Result<ParquetEventStream> {
let file = File::open(&self.file_path)
.with_context(|| format!("Failed to open Parquet file: {}", self.file_path))?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file)
.context("Failed to create Parquet reader")?;
let reader = builder.build().context("Failed to build Parquet reader")?;
Ok(ParquetEventStream { reader })
}
/// Convert Arrow RecordBatch to events
///
/// Internal helper to convert Parquet's Arrow RecordBatch representation
/// into our typed ParquetMarketDataEvent structures.
///
/// # Arguments
/// * `batch` - Arrow RecordBatch from Parquet reader
///
/// # Returns
/// * `Ok(Vec<ParquetMarketDataEvent>)` - Converted events
/// * `Err(_)` - If batch schema is invalid or data cannot be parsed
fn batch_to_events(batch: &RecordBatch) -> Result<Vec<ParquetMarketDataEvent>> {
let num_rows = batch.num_rows();
let mut events = Vec::with_capacity(num_rows);
// Extract columns - all required fields
let timestamp_ns = as_primitive_array::<TimestampNanosecondType>(
batch
.column_by_name("timestamp_ns")
.context("Missing timestamp_ns column")?,
);
let symbol = batch
.column_by_name("symbol")
.context("Missing symbol column")?
.as_any()
.downcast_ref::<StringArray>()
.context("Invalid symbol column type")?;
let venue = batch
.column_by_name("venue")
.context("Missing venue column")?
.as_any()
.downcast_ref::<StringArray>()
.context("Invalid venue column type")?;
let event_type = batch
.column_by_name("event_type")
.context("Missing event_type column")?
.as_any()
.downcast_ref::<StringArray>()
.context("Invalid event_type column type")?;
let sequence = as_primitive_array::<UInt64Type>(
batch
.column_by_name("sequence")
.context("Missing sequence column")?,
);
// Optional columns
let price = batch
.column_by_name("price")
.and_then(|col| col.as_any().downcast_ref::<Float64Array>());
let quantity = batch
.column_by_name("quantity")
.and_then(|col| col.as_any().downcast_ref::<Float64Array>());
let latency_ns = batch
.column_by_name("latency_ns")
.and_then(|col| as_primitive_array::<UInt64Type>(col).into());
let open = batch
.column_by_name("open")
.and_then(|col| col.as_any().downcast_ref::<Float64Array>());
let high = batch
.column_by_name("high")
.and_then(|col| col.as_any().downcast_ref::<Float64Array>());
let low = batch
.column_by_name("low")
.and_then(|col| col.as_any().downcast_ref::<Float64Array>());
for i in 0..num_rows {
let event_type_str = event_type.value(i);
let parsed_event_type = match event_type_str {
"Trade" => MarketDataEventType::Trade,
"Quote" => MarketDataEventType::Quote,
"OrderBookUpdate" => MarketDataEventType::OrderBookUpdate,
"StatusUpdate" => MarketDataEventType::StatusUpdate,
"Ohlcv" => MarketDataEventType::Ohlcv,
_ => {
return Err(anyhow::anyhow!("Unknown event type: {}", event_type_str));
}
};
events.push(ParquetMarketDataEvent {
timestamp_ns: timestamp_ns.value(i) as u64,
symbol: symbol.value(i).to_string(),
venue: venue.value(i).to_string(),
event_type: parsed_event_type,
price: price.and_then(|arr| {
if arr.is_null(i) {
None
} else {
Some(arr.value(i))
}
}),
quantity: quantity.and_then(|arr| {
if arr.is_null(i) {
None
} else {
Some(arr.value(i))
}
}),
sequence: sequence.value(i),
latency_ns: latency_ns.and_then(|arr| {
if arr.is_null(i) {
None
} else {
Some(arr.value(i))
}
}),
open: open.and_then(|arr| {
if arr.is_null(i) {
None
} else {
Some(arr.value(i))
}
}),
high: high.and_then(|arr| {
if arr.is_null(i) {
None
} else {
Some(arr.value(i))
}
}),
low: low.and_then(|arr| {
if arr.is_null(i) {
None
} else {
Some(arr.value(i))
}
}),
});
}
Ok(events)
}
}
/// Streaming iterator over Parquet events
///
/// Provides batch-by-batch streaming of events for memory-efficient processing.
/// Each call to `next()` returns a batch of events from the Parquet file.
pub struct ParquetEventStream {
reader: parquet::arrow::arrow_reader::ParquetRecordBatchReader,
}
impl ParquetEventStream {
/// Get next batch of events
///
/// Returns `None` when all batches have been read.
///
/// # Returns
/// * `Some(Ok(Vec<ParquetMarketDataEvent>))` - Next batch of events
/// * `Some(Err(_))` - Error reading or parsing batch
/// * `None` - End of stream
pub async fn next(&mut self) -> Option<Result<Vec<ParquetMarketDataEvent>>> {
// Note: Parquet reader is synchronous, but we provide async API for consistency
match self.reader.next() {
Some(Ok(batch)) => Some(ParquetDataLoader::batch_to_events(&batch)),
Some(Err(e)) => Some(Err(anyhow::anyhow!("Failed to read batch: {}", e))),
None => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_parquet_loader_creation() {
let loader = ParquetDataLoader::new("test.parquet");
assert!(loader.file_path.contains("test.parquet"));
}
#[tokio::test]
async fn test_parquet_loader_nonexistent_file() {
let loader = ParquetDataLoader::new("nonexistent.parquet");
let result = loader.load_all().await;
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("Failed to open Parquet file"));
}
}