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

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};
@@ -984,18 +984,15 @@ pub struct DataSet {
     instruments: HashMap<String, Instrument>,
     calendar: TradingCalendar,
     market_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
-    market_index: HashMap<NaiveDate, HashMap<String, Arc<DailyMarketSnapshot>>>,
     factor_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
-    factor_index: HashMap<NaiveDate, HashMap<String, Arc<DailyFactorSnapshot>>>,
     factor_text_by_date: BTreeMap<NaiveDate, Vec<FactorTextValue>>,
     factor_text_index: HashMap<(NaiveDate, String, String), FactorTextValue>,
     candidate_by_date: BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>,
-    candidate_index: HashMap<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>,
     corporate_actions_by_date: BTreeMap<NaiveDate, Vec<CorporateAction>>,
     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,
@@ -1247,25 +1244,11 @@ impl DataSet {
             .collect::<HashMap<_, _>>();
 
         let market = market.into_iter().map(Arc::new).collect::<Vec<_>>();
-        let market_by_date = group_arc_by_date(&market, |item| item.date);
-        let mut market_index =
-            HashMap::<NaiveDate, HashMap<String, Arc<DailyMarketSnapshot>>>::new();
-        for item in market {
-            market_index
-                .entry(item.date)
-                .or_default()
-                .insert(item.symbol.clone(), item);
-        }
+        let mut market_by_date = group_arc_by_date(&market, |item| item.date);
+        sort_arc_groups_by_symbol(&mut market_by_date, |item| item.symbol.as_str());
 
-        let factor_by_date = group_arc_by_date(&factors, |item| item.date);
-        let mut factor_index =
-            HashMap::<NaiveDate, HashMap<String, Arc<DailyFactorSnapshot>>>::new();
-        for item in factors {
-            factor_index
-                .entry(item.date)
-                .or_default()
-                .insert(item.symbol.clone(), item);
-        }
+        let mut factor_by_date = group_arc_by_date(&factors, |item| item.date);
+        sort_arc_groups_by_symbol(&mut factor_by_date, |item| item.symbol.as_str());
         let factor_texts = factor_texts
             .into_iter()
             .filter_map(|mut item| {
@@ -1283,15 +1266,8 @@ impl DataSet {
             .map(|item| ((item.date, item.symbol.clone(), item.field.clone()), item))
             .collect::<HashMap<_, _>>();
 
-        let candidate_by_date = group_arc_by_date(&candidates, |item| item.date);
-        let mut candidate_index =
-            HashMap::<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>::new();
-        for item in candidates {
-            candidate_index
-                .entry(item.date)
-                .or_default()
-                .insert(item.symbol.clone(), item);
-        }
+        let mut candidate_by_date = group_arc_by_date(&candidates, |item| item.date);
+        sort_arc_groups_by_symbol(&mut candidate_by_date, |item| item.symbol.as_str());
         let corporate_actions_by_date = group_by_date(corporate_actions, |item| item.date);
         let execution_quotes_by_date = build_execution_quote_index(execution_quotes);
         let order_book_depth_index = build_order_book_depth_index(order_book_depth);
@@ -1300,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);
@@ -1309,18 +1284,15 @@ impl DataSet {
             instruments,
             calendar,
             market_by_date,
-            market_index,
             factor_by_date,
-            factor_index,
             factor_text_by_date,
             factor_text_index,
             candidate_by_date,
-            candidate_index,
             corporate_actions_by_date,
             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,
@@ -1366,24 +1338,54 @@ impl DataSet {
     }
 
     pub fn market(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
-        self.market_index
+        self.market_by_date
             .get(&date)
-            .and_then(|rows| rows.get(symbol))
-            .map(Arc::as_ref)
+            .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_index
+        self.factor_by_date
             .get(&date)
-            .and_then(|rows| rows.get(symbol))
-            .map(Arc::as_ref)
+            .and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
     }
 
     pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> {
-        self.candidate_index
+        self.candidate_by_date
             .get(&date)
-            .and_then(|rows| rows.get(symbol))
-            .map(Arc::as_ref)
+            .and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
     }
 
     pub fn benchmark(&self, date: NaiveDate) -> Option<&BenchmarkSnapshot> {
@@ -1615,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()
     }
@@ -2173,22 +2174,18 @@ 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;
         }
         let previous_date = *series.dates.get(end - 1)?;
-        self.market_index
-            .get(&previous_date)
-            .and_then(|rows| rows.get(symbol))
-            .map(Arc::as_ref)
+        self.market(previous_date, symbol)
     }
 
     pub fn factor_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyFactorSnapshot> {
@@ -2255,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()
     }
@@ -2269,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()
     }
@@ -2311,24 +2306,12 @@ impl DataSet {
         let start = days.len().saturating_sub(count);
         days[start..]
             .iter()
-            .map(|day| {
-                evaluator(
-                    self.candidate_index
-                        .get(day)
-                        .and_then(|rows| rows.get(symbol))
-                        .map(Arc::as_ref),
-                    self.market_index
-                        .get(day)
-                        .and_then(|rows| rows.get(symbol))
-                        .map(Arc::as_ref),
-                )
-            })
+            .map(|day| evaluator(self.candidate(*day, symbol), self.market(*day, symbol)))
             .collect()
     }
 
     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))
     }
 
@@ -2338,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))
     }
 
@@ -2349,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))
     }
 
@@ -2437,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)
@@ -2468,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" => {
@@ -2496,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),
@@ -2549,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))
     }
 
@@ -2622,8 +2594,8 @@ impl DataSet {
             .get_or_init(|| {
                 build_eligible_universe(
                     &self.factor_by_date,
-                    &self.candidate_index,
-                    &self.market_index,
+                    &self.candidate_by_date,
+                    &self.market_by_date,
                     &self.instruments,
                 )
             })
@@ -3523,6 +3495,24 @@ where
     grouped
 }
 
+fn sort_arc_groups_by_symbol<T, F>(groups: &mut BTreeMap<NaiveDate, Vec<Arc<T>>>, symbol_of: F)
+where
+    F: Fn(&T) -> &str + Copy,
+{
+    for rows in groups.values_mut() {
+        rows.sort_by(|left, right| symbol_of(left.as_ref()).cmp(symbol_of(right.as_ref())));
+    }
+}
+
+fn find_arc_by_symbol<'a, T, F>(rows: &'a [Arc<T>], symbol: &str, symbol_of: F) -> Option<&'a T>
+where
+    F: Fn(&T) -> &str,
+{
+    rows.binary_search_by(|row| symbol_of(row.as_ref()).cmp(symbol))
+        .ok()
+        .map(|index| rows[index].as_ref())
+}
+
 fn collect_benchmark_code(benchmarks: &[BenchmarkSnapshot]) -> Result<String, DataSetError> {
     let mut codes = benchmarks
         .iter()
@@ -3591,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>> {
@@ -3672,8 +3640,8 @@ fn build_order_book_depth_index(
 
 fn build_eligible_universe(
     factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
-    candidate_index: &HashMap<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>,
-    market_index: &HashMap<NaiveDate, HashMap<String, Arc<DailyMarketSnapshot>>>,
+    candidate_by_date: &BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>,
+    market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
     instruments: &HashMap<String, Instrument>,
 ) -> BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>> {
     let mut per_date = BTreeMap::<NaiveDate, Vec<EligibleUniverseSnapshot>>::new();
@@ -3684,19 +3652,16 @@ fn build_eligible_universe(
             if factor.market_cap_bn <= 0.0 || !factor.market_cap_bn.is_finite() {
                 continue;
             }
-            let Some(candidate) = candidate_index
-                .get(date)
-                .and_then(|rows| rows.get(&factor.symbol))
-            else {
+            let Some(candidate) = candidate_by_date.get(date).and_then(|rows| {
+                find_arc_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str())
+            }) else {
                 continue;
             };
-            let Some(market) = market_index
-                .get(date)
-                .and_then(|rows| rows.get(&factor.symbol))
-            else {
+            let Some(market) = market_by_date.get(date).and_then(|rows| {
+                find_arc_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str())
+            }) else {
                 continue;
             };
-            let market = market.as_ref();
             if ChinaAShareRiskControl::selection_rejection_reason(
                 *date,
                 candidate,