瘦身回测数据集按日索引内存

crates/fidc-core/src/data.rs

@@ -984,13 +984,10 @@ 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>>,
@@ -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);
@@ -1309,13 +1285,10 @@ 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,
@@ -1366,24 +1339,21 @@ 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()))
     }
 
     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> {
@@ -2185,10 +2155,7 @@ impl DataSet {
             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> {
@@ -2311,18 +2278,7 @@ 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()
     }
 
@@ -2622,8 +2578,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 +3479,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()
@@ -3672,8 +3646,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 +3658,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,