懒加载日线序列缓存降低回测内存

crates/fidc-core/src/data.rs

@@ -1,7 +1,7 @@
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::fs;
 use std::path::Path;
-use std::sync::{Arc, OnceLock};
+use std::sync::{Arc, OnceLock, RwLock};
 
 use chrono::{NaiveDate, NaiveDateTime};
 use serde::{Deserialize, Serialize};
@@ -992,7 +992,7 @@ pub struct DataSet {
     execution_quotes_by_date: HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>,
     order_book_depth_index: HashMap<(NaiveDate, String), Vec<IntradayOrderBookDepthLevel>>,
     benchmark_by_date: BTreeMap<NaiveDate, BenchmarkSnapshot>,
-    market_series_by_symbol: HashMap<String, SymbolPriceSeries>,
+    market_series_by_symbol: Arc<RwLock<HashMap<String, Arc<SymbolPriceSeries>>>>,
     benchmark_series_cache: BenchmarkPriceSeries,
     eligible_universe_by_date: Arc<OnceLock<BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>>>>,
     benchmark_code: String,
@@ -1276,7 +1276,6 @@ impl DataSet {
             .into_iter()
             .map(|item| (item.date, item))
             .collect::<BTreeMap<_, _>>();
-        let market_series_by_symbol = build_market_series(&market_by_date);
         let benchmark_series_cache =
             BenchmarkPriceSeries::new(&benchmark_by_date.values().cloned().collect::<Vec<_>>());
         let futures_params_by_symbol = build_futures_params_index(futures_params);
@@ -1293,7 +1292,7 @@ impl DataSet {
             execution_quotes_by_date,
             order_book_depth_index,
             benchmark_by_date,
-            market_series_by_symbol,
+            market_series_by_symbol: Arc::new(RwLock::new(HashMap::new())),
             benchmark_series_cache,
             eligible_universe_by_date: Arc::new(OnceLock::new()),
             benchmark_code,
@@ -1344,6 +1343,39 @@ impl DataSet {
             .and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
     }
 
+    fn market_series(&self, symbol: &str) -> Option<Arc<SymbolPriceSeries>> {
+        if let Some(series) = self
+            .market_series_by_symbol
+            .read()
+            .expect("market series cache lock poisoned")
+            .get(symbol)
+            .cloned()
+        {
+            return Some(series);
+        }
+
+        let rows = self
+            .market_by_date
+            .values()
+            .filter_map(|day_rows| find_arc_by_symbol(day_rows, symbol, |row| row.symbol.as_str()))
+            .collect::<Vec<_>>();
+        if rows.is_empty() {
+            return None;
+        }
+
+        let series = Arc::new(SymbolPriceSeries::new(symbol.to_string(), rows));
+        let mut cache = self
+            .market_series_by_symbol
+            .write()
+            .expect("market series cache lock poisoned");
+        Some(
+            cache
+                .entry(symbol.to_string())
+                .or_insert_with(|| Arc::clone(&series))
+                .clone(),
+        )
+    }
+
     pub fn factor(&self, date: NaiveDate, symbol: &str) -> Option<&DailyFactorSnapshot> {
         self.factor_by_date
             .get(&date)
@@ -1585,8 +1617,7 @@ impl DataSet {
         bar_count: usize,
         include_now: bool,
     ) -> Vec<DailyMarketSnapshot> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .map(|series| series.trailing_snapshots(date, bar_count, include_now))
             .unwrap_or_default()
     }
@@ -2143,13 +2174,12 @@ impl DataSet {
         symbol: &str,
         field: PriceField,
     ) -> Option<f64> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .and_then(|series| series.price_on_or_before(date, field))
     }
 
     pub fn market_before(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
-        let series = self.market_series_by_symbol.get(symbol)?;
+        let series = self.market_series(symbol)?;
         let end = series.previous_completed_end_index(date)?;
         if end == 0 {
             return None;
@@ -2222,8 +2252,7 @@ impl DataSet {
     }
 
     pub fn market_closes_up_to(&self, date: NaiveDate, symbol: &str, lookback: usize) -> Vec<f64> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .map(|series| series.trailing_values(date, lookback, PriceField::Close))
             .unwrap_or_default()
     }
@@ -2236,8 +2265,7 @@ impl DataSet {
         field: &str,
         include_now: bool,
     ) -> Vec<f64> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .map(|series| series.trailing_numeric_values(date, bar_count, field, include_now))
             .unwrap_or_default()
     }
@@ -2283,8 +2311,7 @@ impl DataSet {
     }
 
     pub fn market_decision_close(&self, date: NaiveDate, symbol: &str) -> Option<f64> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .and_then(|series| series.decision_price_on_or_before(date))
     }
 
@@ -2294,8 +2321,7 @@ impl DataSet {
         symbol: &str,
         lookback: usize,
     ) -> Option<f64> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .and_then(|series| series.decision_close_moving_average(date, lookback))
     }
 
@@ -2305,8 +2331,7 @@ impl DataSet {
         symbol: &str,
         lookback: usize,
     ) -> Option<f64> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .and_then(|series| series.decision_volume_moving_average(date, lookback))
     }
 
@@ -2393,12 +2418,10 @@ impl DataSet {
         let field = normalize_field(field);
         match field.as_str() {
             "close" | "prev_close" | "stock_close" | "price" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .and_then(|series| series.decision_close_moving_average(date, lookback)),
             "volume" | "stock_volume" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .and_then(|series| series.decision_volume_moving_average(date, lookback)),
             "day_open" | "dayopen" => {
                 self.market_moving_average(date, symbol, lookback, PriceField::DayOpen)
@@ -2424,8 +2447,7 @@ impl DataSet {
                 self.market_moving_average(date, symbol, lookback, PriceField::Close)
             }
             "volume" | "stock_volume" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .and_then(|series| series.current_volume_moving_average(date, lookback))
                 .or_else(|| self.factor_moving_average(date, symbol, "daily_volume", lookback)),
             "day_open" | "dayopen" => {
@@ -2452,28 +2474,23 @@ impl DataSet {
         let field = normalize_field(field);
         match field.as_str() {
             "close" | "prev_close" | "stock_close" | "price" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .and_then(|series| series.decision_prev_close_values(date, lookback))
                 .unwrap_or_default(),
             "volume" | "stock_volume" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .and_then(|series| series.decision_volume_values(date, lookback))
                 .unwrap_or_default(),
             "day_open" | "dayopen" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .map(|series| series.trailing_values(date, lookback, PriceField::DayOpen))
                 .unwrap_or_default(),
             "open" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .map(|series| series.trailing_values(date, lookback, PriceField::Open))
                 .unwrap_or_default(),
             "last" | "last_price" => self
-                .market_series_by_symbol
-                .get(symbol)
+                .market_series(symbol)
                 .map(|series| series.trailing_values(date, lookback, PriceField::Last))
                 .unwrap_or_default(),
             other => self.factor_numeric_values(date, symbol, other, lookback),
@@ -2505,8 +2522,7 @@ impl DataSet {
         lookback: usize,
         field: PriceField,
     ) -> Option<f64> {
-        self.market_series_by_symbol
-            .get(symbol)
+        self.market_series(symbol)
             .and_then(|series| series.moving_average(date, lookback, field))
     }
 
@@ -3565,28 +3581,6 @@ mod optional_date_format {
     }
 }
 
-fn build_market_series(
-    market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
-) -> HashMap<String, SymbolPriceSeries> {
-    let mut grouped = HashMap::<String, Vec<&DailyMarketSnapshot>>::new();
-    for rows in market_by_date.values() {
-        for row in rows {
-            grouped
-                .entry(row.symbol.clone())
-                .or_default()
-                .push(row.as_ref());
-        }
-    }
-
-    grouped
-        .into_iter()
-        .map(|(symbol, rows)| {
-            let series = SymbolPriceSeries::new(symbol.clone(), rows);
-            (symbol, series)
-        })
-        .collect()
-}
-
 fn build_futures_params_index(
     rows: Vec<FuturesTradingParameter>,
 ) -> HashMap<String, Vec<FuturesTradingParameter>> {