From 6befda18a58060b12fd14d5afe24b02ae2cee41b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 23 Feb 2026 18:33:26 +0100 Subject: [PATCH] feat(ml): add quarterly download binary for futures baseline Adds `ml/examples/download_baseline.rs` that downloads 730 days of Databento OHLCV-1m data in quarterly chunks for 4 CME futures symbols. Features: - Reads universe config from TOML (symbols, date range, dataset) - Splits date range into calendar-quarter chunks (~90 days each) - Resume support: skips existing non-empty files - Uses `get_range_to_file` for streaming writes to .dbn.zst - Dry-run mode with cost estimate ($0.12/symbol/day) - Confirmation prompt (skippable with --yes) - Per-file progress with timing and byte counts - Failure-tolerant: logs errors and continues Co-Authored-By: Claude Opus 4.6 --- ml/examples/download_baseline.rs | 473 +++++++++++++++++++++++++++++++ 1 file changed, 473 insertions(+) create mode 100644 ml/examples/download_baseline.rs diff --git a/ml/examples/download_baseline.rs b/ml/examples/download_baseline.rs new file mode 100644 index 000000000..edb635c36 --- /dev/null +++ b/ml/examples/download_baseline.rs @@ -0,0 +1,473 @@ +//! Download 730 days of Databento OHLCV-1m data in quarterly chunks +//! +//! Downloads futures baseline data for ML model training, split into +//! ~90-day quarterly files for efficient caching and resume support. +//! +//! Configuration is read from a universe TOML file (default: +//! `config/universe-futures-baseline.toml`). +//! +//! Usage: +//! # Dry run (preview config and cost estimate) +//! cargo run -p ml --example download_baseline --release -- --dry-run +//! +//! # Download with confirmation prompt +//! cargo run -p ml --example download_baseline --release +//! +//! # Skip confirmation prompt +//! cargo run -p ml --example download_baseline --release -- --yes + +use anyhow::{Context, Result}; +use chrono::{Datelike, NaiveDate}; +use clap::Parser; +use databento::historical::timeseries::GetRangeToFileParams; +use databento::historical::DateTimeRange; +use databento::HistoricalClient; +use dbn::Schema; +use serde::Deserialize; +use std::env; +use std::fs; +use std::path::PathBuf; +use std::time::Instant; + +// --------------------------------------------------------------------------- +// CLI +// --------------------------------------------------------------------------- + +#[derive(Debug, Parser)] +#[command( + name = "download_baseline", + about = "Download quarterly OHLCV-1m futures data from Databento" +)] +struct Opts { + /// Output directory for downloaded files + #[arg(long, default_value = "data/cache/futures-baseline")] + output_dir: String, + + /// Path to universe configuration TOML + #[arg(long, default_value = "config/universe-futures-baseline.toml")] + universe_config: String, + + /// Preview only, do not download + #[arg(long)] + dry_run: bool, + + /// Skip confirmation prompt + #[arg(long)] + yes: bool, +} + +// --------------------------------------------------------------------------- +// Universe config +// --------------------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +struct UniverseConfig { + universe: UniverseMeta, + symbols: Vec, +} + +#[derive(Debug, Deserialize)] +struct UniverseMeta { + name: String, + #[allow(dead_code)] + description: String, + date_range_start: String, + date_range_end: String, + #[allow(dead_code)] + bar_size: String, + databento_dataset: String, + #[allow(dead_code)] + databento_schema: String, +} + +#[derive(Debug, Deserialize)] +struct SymbolEntry { + symbol: String, + #[allow(dead_code)] + exchange: String, + #[allow(dead_code)] + asset_class: String, + #[allow(dead_code)] + trading_symbol: Option, + #[allow(dead_code)] + min_daily_volume: Option, +} + +// --------------------------------------------------------------------------- +// Quarter representation +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +struct Quarter { + label: String, + start: NaiveDate, + end: NaiveDate, +} + +/// Split a date range into calendar-quarter chunks. +/// +/// Each chunk begins at the later of `start` and the quarter boundary, and +/// ends at the earlier of `end` and the next quarter boundary. +fn generate_quarters(start: NaiveDate, end: NaiveDate) -> Vec { + let mut quarters = Vec::new(); + let mut cursor = start; + + while cursor < end { + let q = (cursor.month() - 1) / 3 + 1; // 1..4 + let year = cursor.year(); + let label = format!("{}-Q{}", year, q); + + // Next quarter boundary + let next_q_start = if q == 4 { + NaiveDate::from_ymd_opt(year + 1, 1, 1) + } else { + NaiveDate::from_ymd_opt(year, q * 3 + 1, 1) + }; + + let chunk_end = match next_q_start { + Some(nq) if nq < end => nq, + _ => end, + }; + + quarters.push(Quarter { + label, + start: cursor, + end: chunk_end, + }); + + cursor = chunk_end; + } + + quarters +} + +// --------------------------------------------------------------------------- +// Conversion helpers +// --------------------------------------------------------------------------- + +/// Convert a `chrono::NaiveDate` (interpreted as midnight UTC) to UNIX +/// nanoseconds for the Databento `DateTimeRange` API. +fn naive_date_to_unix_nanos(date: NaiveDate) -> u64 { + let ts = date + .and_hms_opt(0, 0, 0) + .map(|dt| dt.and_utc().timestamp()) + .unwrap_or(0); + (ts as u64).saturating_mul(1_000_000_000) +} + +/// Build a `DateTimeRange` from two `NaiveDate`s. +fn date_range(start: NaiveDate, end: NaiveDate) -> Result { + let start_ns = naive_date_to_unix_nanos(start); + let end_ns = naive_date_to_unix_nanos(end); + DateTimeRange::try_from((start_ns, end_ns)).map_err(|e| anyhow::anyhow!("{}", e)) +} + +// --------------------------------------------------------------------------- +// Download logic +// --------------------------------------------------------------------------- + +struct DownloadStats { + successful: usize, + failed: usize, + skipped: usize, + total_bytes: u64, +} + +impl DownloadStats { + fn new() -> Self { + Self { + successful: 0, + failed: 0, + skipped: 0, + total_bytes: 0, + } + } +} + +async fn download_quarter( + client: &mut HistoricalClient, + symbol: &str, + quarter: &Quarter, + dataset: &str, + output_dir: &PathBuf, +) -> Result { + let filename = format!("{}_{}.dbn.zst", symbol, quarter.label); + let symbol_dir = output_dir.join(symbol); + fs::create_dir_all(&symbol_dir).context("Failed to create symbol directory")?; + let file_path = symbol_dir.join(&filename); + + // Resume support: skip if file exists and is non-empty + if file_path.exists() { + let meta = fs::metadata(&file_path).context("Failed to read file metadata")?; + if meta.len() > 0 { + return Ok(meta.len()); + } + } + + let dt_range = date_range(quarter.start, quarter.end)?; + + let params = GetRangeToFileParams::builder() + .dataset(dataset.to_string()) + .symbols(vec![symbol.to_string()]) + .schema(Schema::Ohlcv1M) + .date_time_range(dt_range) + .path(file_path.clone()) + .build(); + + let _decoder = client + .timeseries() + .get_range_to_file(¶ms) + .await + .with_context(|| { + format!( + "Failed to download {} for {}", + quarter.label, symbol + ) + })?; + + let meta = fs::metadata(&file_path).context("Failed to read downloaded file metadata")?; + Ok(meta.len()) +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +#[tokio::main] +async fn main() -> Result<()> { + let opts = Opts::parse(); + + // Load .env for API key + dotenv::dotenv().ok(); + + // ------------------------------------------------------------------ + // Parse universe config + // ------------------------------------------------------------------ + let config_text = + fs::read_to_string(&opts.universe_config).with_context(|| { + format!("Failed to read universe config: {}", opts.universe_config) + })?; + let config: UniverseConfig = + toml::from_str(&config_text).context("Failed to parse universe config TOML")?; + + let start_date = NaiveDate::parse_from_str(&config.universe.date_range_start, "%Y-%m-%d") + .context("Failed to parse date_range_start")?; + let end_date = NaiveDate::parse_from_str(&config.universe.date_range_end, "%Y-%m-%d") + .context("Failed to parse date_range_end")?; + + let symbols: Vec<&str> = config.symbols.iter().map(|s| s.symbol.as_str()).collect(); + let quarters = generate_quarters(start_date, end_date); + let total_days = (end_date - start_date).num_days(); + + // ------------------------------------------------------------------ + // Print config summary + // ------------------------------------------------------------------ + println!("================================================================================"); + println!("Futures Baseline Download - Databento (Quarterly Chunks)"); + println!("================================================================================"); + println!(); + println!("Universe: {}", config.universe.name); + println!("Dataset: {}", config.universe.databento_dataset); + println!("Schema: ohlcv-1m"); + println!("Date range: {} to {}", start_date, end_date); + println!("Total days: {}", total_days); + println!( + "Symbols: {} ({})", + symbols.len(), + symbols.join(", ") + ); + println!("Quarters: {}", quarters.len()); + println!("Output: {}", opts.output_dir); + println!(); + + // Print quarter breakdown + println!("Quarter breakdown:"); + for q in &quarters { + let days = (q.end - q.start).num_days(); + println!(" {} : {} to {} ({} days)", q.label, q.start, q.end, days); + } + println!(); + + // Cost estimate ($0.12 per symbol per day) + let total_files = symbols.len() * quarters.len(); + let estimated_cost = total_days as f64 * symbols.len() as f64 * 0.12; + println!("Total files: {}", total_files); + println!("Estimated cost: ${:.2} ({} symbols x {} days x $0.12/sym/day)", + estimated_cost, symbols.len(), total_days); + println!(); + + // ------------------------------------------------------------------ + // Dry run exit + // ------------------------------------------------------------------ + if opts.dry_run { + println!("[DRY RUN] Preview complete. Remove --dry-run to execute downloads."); + return Ok(()); + } + + // ------------------------------------------------------------------ + // Confirmation + // ------------------------------------------------------------------ + if !opts.yes { + println!( + "This will download data and may incur costs (~${:.2}).", + estimated_cost + ); + print!("Proceed? (yes/no): "); + std::io::Write::flush(&mut std::io::stdout())?; + + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + let trimmed = input.trim(); + if !trimmed.eq_ignore_ascii_case("yes") && !trimmed.eq_ignore_ascii_case("y") { + println!("Download cancelled."); + return Ok(()); + } + println!(); + } + + // ------------------------------------------------------------------ + // Initialize client + // ------------------------------------------------------------------ + let api_key = env::var("DATABENTO_API_KEY") + .context("DATABENTO_API_KEY not found in environment or .env file")?; + + let mut client = HistoricalClient::builder() + .key(api_key) + .map_err(|e| anyhow::anyhow!("Invalid API key: {}", e))? + .build() + .map_err(|e| anyhow::anyhow!("Failed to build Databento client: {}", e))?; + + let output_dir = PathBuf::from(&opts.output_dir); + fs::create_dir_all(&output_dir)?; + println!("Databento client initialized. Output: {}", opts.output_dir); + println!(); + + // ------------------------------------------------------------------ + // Download loop + // ------------------------------------------------------------------ + let mut stats = DownloadStats::new(); + let mut failed_items: Vec = Vec::new(); + + for symbol in &symbols { + println!("[{}]", symbol); + + for quarter in &quarters { + let filename = format!("{}_{}.dbn.zst", symbol, quarter.label); + let symbol_dir = output_dir.join(symbol); + let file_path = symbol_dir.join(&filename); + + // Check for existing file before calling download + if file_path.exists() { + match fs::metadata(&file_path) { + Ok(meta) if meta.len() > 0 => { + println!( + " [SKIP] {} already exists ({} bytes)", + filename, + meta.len() + ); + stats.skipped += 1; + stats.total_bytes += meta.len(); + continue; + } + _ => {} + } + } + + let t0 = Instant::now(); + match download_quarter( + &mut client, + symbol, + quarter, + &config.universe.databento_dataset, + &output_dir, + ) + .await + { + Ok(size) => { + let elapsed = t0.elapsed().as_secs_f64(); + println!( + " [OK] {} -- {} bytes, {:.1}s", + filename, size, elapsed + ); + stats.successful += 1; + stats.total_bytes += size; + } + Err(e) => { + println!(" [FAIL] {} -- {}", filename, e); + stats.failed += 1; + failed_items.push(format!("{}/{}", symbol, quarter.label)); + } + } + } + println!(); + } + + // ------------------------------------------------------------------ + // Summary + // ------------------------------------------------------------------ + println!("================================================================================"); + println!("DOWNLOAD SUMMARY"); + println!("================================================================================"); + println!(); + println!("Successful: {}", stats.successful); + println!("Skipped: {}", stats.skipped); + println!("Failed: {}", stats.failed); + println!( + "Total size: {:.1} MB", + stats.total_bytes as f64 / 1_048_576.0 + ); + println!(); + + if !failed_items.is_empty() { + println!("Failed downloads:"); + for item in &failed_items { + println!(" - {}", item); + } + println!(); + println!("Re-run the command to retry failed downloads (existing files are skipped)."); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_generate_quarters_full_range() { + let start = NaiveDate::from_ymd_opt(2024, 3, 1).unwrap_or_default(); + let end = NaiveDate::from_ymd_opt(2024, 12, 31).unwrap_or_default(); + let qs = generate_quarters(start, end); + assert_eq!(qs.len(), 4); // Q1(partial), Q2, Q3, Q4 + assert_eq!(qs.first().map(|q| q.label.as_str()), Some("2024-Q1")); + assert_eq!(qs.last().map(|q| q.label.as_str()), Some("2024-Q4")); + } + + #[test] + fn test_generate_quarters_cross_year() { + let start = NaiveDate::from_ymd_opt(2024, 10, 1).unwrap_or_default(); + let end = NaiveDate::from_ymd_opt(2025, 4, 1).unwrap_or_default(); + let qs = generate_quarters(start, end); + assert_eq!(qs.len(), 3); // 2024-Q4, 2025-Q1, 2025-Q2(partial) + assert_eq!(qs.first().map(|q| q.label.as_str()), Some("2024-Q4")); + assert_eq!(qs.last().map(|q| q.label.as_str()), Some("2025-Q2")); + } + + #[test] + fn test_generate_quarters_single_day() { + let start = NaiveDate::from_ymd_opt(2024, 6, 15).unwrap_or_default(); + let end = NaiveDate::from_ymd_opt(2024, 6, 16).unwrap_or_default(); + let qs = generate_quarters(start, end); + assert_eq!(qs.len(), 1); + assert_eq!(qs.first().map(|q| q.label.as_str()), Some("2024-Q2")); + } + + #[test] + fn test_naive_date_to_unix_nanos() { + let date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap_or_default(); + let nanos = naive_date_to_unix_nanos(date); + // 2024-01-01 00:00:00 UTC = 1704067200 seconds + assert_eq!(nanos, 1_704_067_200_000_000_000); + } +}