优化回测数据集索引查找

crates/fidc-core/src/data.rs

@@ -966,13 +966,13 @@ pub struct DataSet {
     instruments: HashMap<String, Instrument>,
     calendar: TradingCalendar,
     market_by_date: BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>,
-    market_index: HashMap<(NaiveDate, String), DailyMarketSnapshot>,
+    market_index: HashMap<NaiveDate, HashMap<String, DailyMarketSnapshot>>,
     factor_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
-    factor_index: HashMap<(NaiveDate, String), 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, String), 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>>,
@@ -1229,16 +1229,23 @@ impl DataSet {
             .collect::<HashMap<_, _>>();
 
         let market_by_date = group_by_date(market.clone(), |item| item.date);
-        let market_index = market
+        let mut market_index = HashMap::<NaiveDate, HashMap<String, DailyMarketSnapshot>>::new();
-            .into_iter()
+        for item in market {
-            .map(|item| ((item.date, item.symbol.clone()), item))
+            market_index
-            .collect::<HashMap<_, _>>();
+                .entry(item.date)
+                .or_default()
+                .insert(item.symbol.clone(), item);
+        }
 
         let factor_by_date = group_arc_by_date(&factors, |item| item.date);
-        let factor_index = factors
+        let mut factor_index =
-            .into_iter()
+            HashMap::<NaiveDate, HashMap<String, Arc<DailyFactorSnapshot>>>::new();
-            .map(|item| ((item.date, item.symbol.clone()), item))
+        for item in factors {
-            .collect::<HashMap<_, _>>();
+            factor_index
+                .entry(item.date)
+                .or_default()
+                .insert(item.symbol.clone(), item);
+        }
         let factor_texts = factor_texts
             .into_iter()
             .filter_map(|mut item| {
@@ -1257,10 +1264,14 @@ impl DataSet {
             .collect::<HashMap<_, _>>();
 
         let candidate_by_date = group_arc_by_date(&candidates, |item| item.date);
-        let candidate_index = candidates
+        let mut candidate_index =
-            .into_iter()
+            HashMap::<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>::new();
-            .map(|item| ((item.date, item.symbol.clone()), item))
+        for item in candidates {
-            .collect::<HashMap<_, _>>();
+            candidate_index
+                .entry(item.date)
+                .or_default()
+                .insert(item.symbol.clone(), item);
+        }
         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);
@@ -1341,18 +1352,22 @@ impl DataSet {
     }
 
     pub fn market(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
-        self.market_index.get(&(date, symbol.to_string()))
+        self.market_index
+            .get(&date)
+            .and_then(|rows| rows.get(symbol))
     }
 
     pub fn factor(&self, date: NaiveDate, symbol: &str) -> Option<&DailyFactorSnapshot> {
         self.factor_index
-            .get(&(date, symbol.to_string()))
+            .get(&date)
+            .and_then(|rows| rows.get(symbol))
             .map(Arc::as_ref)
     }
 
     pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> {
         self.candidate_index
-            .get(&(date, symbol.to_string()))
+            .get(&date)
+            .and_then(|rows| rows.get(symbol))
             .map(Arc::as_ref)
     }
 
@@ -2154,7 +2169,9 @@ impl DataSet {
             return None;
         }
         let previous_date = *series.dates.get(end - 1)?;
-        self.market_index.get(&(previous_date, symbol.to_string()))
+        self.market_index
+            .get(&previous_date)
+            .and_then(|rows| rows.get(symbol))
     }
 
     pub fn factor_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyFactorSnapshot> {
@@ -2276,9 +2293,10 @@ impl DataSet {
             .map(|day| {
                 evaluator(
                     self.candidate_index
-                        .get(&(*day, symbol.to_string()))
+                        .get(day)
+                        .and_then(|rows| rows.get(symbol))
                         .map(Arc::as_ref),
-                    self.market_index.get(&(*day, symbol.to_string())),
+                    self.market_index.get(day).and_then(|rows| rows.get(symbol)),
                 )
             })
             .collect()
@@ -3615,8 +3633,8 @@ fn build_order_book_depth_index(
 
 fn build_eligible_universe(
     factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
-    candidate_index: &HashMap<(NaiveDate, String), Arc<CandidateEligibility>>,
+    candidate_index: &HashMap<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>,
-    market_index: &HashMap<(NaiveDate, String), DailyMarketSnapshot>,
+    market_index: &HashMap<NaiveDate, HashMap<String, DailyMarketSnapshot>>,
     instruments: &HashMap<String, Instrument>,
 ) -> BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>> {
     let mut per_date = BTreeMap::<NaiveDate, Vec<EligibleUniverseSnapshot>>::new();
@@ -3627,11 +3645,16 @@ fn build_eligible_universe(
             if factor.market_cap_bn <= 0.0 || !factor.market_cap_bn.is_finite() {
                 continue;
             }
-            let key = (*date, factor.symbol.clone());
+            let Some(candidate) = candidate_index
-            let Some(candidate) = candidate_index.get(&key) else {
+                .get(date)
+                .and_then(|rows| rows.get(&factor.symbol))
+            else {
                 continue;
             };
-            let Some(market) = market_index.get(&key) else {
+            let Some(market) = market_index
+                .get(date)
+                .and_then(|rows| rows.get(&factor.symbol))
+            else {
                 continue;
             };
             if ChinaAShareRiskControl::selection_rejection_reason(