增加逐日紧凑证券索引

2026-08-24 04:04:03 +08:00
parent 8254ebbb47
commit 589f94e5b2
1 changed files with 126 additions and 9 deletions
@@ -1116,10 +1116,13 @@ pub struct DataSet {
    instruments: HashMap<String, Instrument>,
    calendar: TradingCalendar,
    market_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
+    market_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
    factor_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
+    factor_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
    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_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
    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>>,
@@ -1127,6 +1130,7 @@ pub struct DataSet {
    market_series_by_symbol: Arc<HashMap<String, Arc<SymbolPriceSeries>>>,
    adjusted_close_series_by_symbol: Arc<HashMap<String, Arc<AdjustedCloseSeries>>>,
    benchmark_series_cache: BenchmarkPriceSeries,
+    symbol_id_by_code: Arc<HashMap<String, u32>>,
    eligible_universe_by_date: Arc<OnceLock<BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>>>>,
    benchmark_code: String,
    futures_params_by_symbol: HashMap<String, Vec<FuturesTradingParameter>>,
@@ -1330,6 +1334,24 @@ impl DataSet {

        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 symbol_id_by_code = build_symbol_id_index(
+            &instruments,
+            &market_by_date,
+            &factor_by_date,
+            &candidate_by_date,
+        );
+        let market_symbol_ids_by_date =
+            build_group_symbol_ids(&market_by_date, &symbol_id_by_code, |item| {
+                item.symbol.as_str()
+            });
+        let factor_symbol_ids_by_date =
+            build_group_symbol_ids(&factor_by_date, &symbol_id_by_code, |item| {
+                item.symbol.as_str()
+            });
+        let candidate_symbol_ids_by_date =
+            build_group_symbol_ids(&candidate_by_date, &symbol_id_by_code, |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);
@@ -1346,10 +1368,13 @@ impl DataSet {
            instruments,
            calendar,
            market_by_date,
+            market_symbol_ids_by_date: Arc::new(market_symbol_ids_by_date),
            factor_by_date,
+            factor_symbol_ids_by_date: Arc::new(factor_symbol_ids_by_date),
            factor_text_by_date,
            factor_text_index,
            candidate_by_date,
+            candidate_symbol_ids_by_date: Arc::new(candidate_symbol_ids_by_date),
            corporate_actions_by_date,
            execution_quotes_by_date,
            order_book_depth_index,
@@ -1357,6 +1382,7 @@ impl DataSet {
            market_series_by_symbol: Arc::new(market_series_by_symbol),
            adjusted_close_series_by_symbol: Arc::new(adjusted_close_series_by_symbol),
            benchmark_series_cache,
+            symbol_id_by_code: Arc::new(symbol_id_by_code),
            eligible_universe_by_date: Arc::new(OnceLock::new()),
            benchmark_code,
            futures_params_by_symbol,
@@ -1401,9 +1427,12 @@ impl DataSet {
    }

    pub fn market(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
-        self.market_by_date
-            .get(&date)
-            .and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
+        let symbol_id = *self.symbol_id_by_code.get(symbol)?;
+        find_arc_by_symbol_id(
+            self.market_by_date.get(&date)?,
+            self.market_symbol_ids_by_date.get(&date)?,
+            symbol_id,
+        )
    }

    fn market_series(&self, symbol: &str) -> Option<&SymbolPriceSeries> {
@@ -1417,15 +1446,21 @@ impl DataSet {
    }

    pub fn factor(&self, date: NaiveDate, symbol: &str) -> Option<&DailyFactorSnapshot> {
-        self.factor_by_date
-            .get(&date)
-            .and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
+        let symbol_id = *self.symbol_id_by_code.get(symbol)?;
+        find_arc_by_symbol_id(
+            self.factor_by_date.get(&date)?,
+            self.factor_symbol_ids_by_date.get(&date)?,
+            symbol_id,
+        )
    }

    pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> {
-        self.candidate_by_date
-            .get(&date)
-            .and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
+        let symbol_id = *self.symbol_id_by_code.get(symbol)?;
+        find_arc_by_symbol_id(
+            self.candidate_by_date.get(&date)?,
+            self.candidate_symbol_ids_by_date.get(&date)?,
+            symbol_id,
+        )
    }

    pub fn benchmark(&self, date: NaiveDate) -> Option<&BenchmarkSnapshot> {
@@ -3079,6 +3114,88 @@ where
    }
 }

+fn build_symbol_id_index(
+    instruments: &HashMap<String, Instrument>,
+    market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
+    factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
+    candidate_by_date: &BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>,
+) -> HashMap<String, u32> {
+    let mut symbols = instruments.keys().cloned().collect::<HashSet<_>>();
+    for rows in market_by_date.values() {
+        for row in rows {
+            if !symbols.contains(row.symbol.as_str()) {
+                symbols.insert(row.symbol.clone());
+            }
+        }
+    }
+    for rows in factor_by_date.values() {
+        for row in rows {
+            if !symbols.contains(row.symbol.as_str()) {
+                symbols.insert(row.symbol.clone());
+            }
+        }
+    }
+    for rows in candidate_by_date.values() {
+        for row in rows {
+            if !symbols.contains(row.symbol.as_str()) {
+                symbols.insert(row.symbol.clone());
+            }
+        }
+    }
+    let mut symbols = symbols.into_iter().collect::<Vec<_>>();
+    symbols.sort_unstable();
+    symbols
+        .into_iter()
+        .enumerate()
+        .map(|(index, symbol)| {
+            (
+                symbol,
+                u32::try_from(index).expect("FIDC symbol index exceeds u32 capacity"),
+            )
+        })
+        .collect()
+}
+
+fn build_group_symbol_ids<T, F>(
+    groups: &BTreeMap<NaiveDate, Vec<Arc<T>>>,
+    symbol_id_by_code: &HashMap<String, u32>,
+    symbol_of: F,
+) -> BTreeMap<NaiveDate, Vec<u32>>
+where
+    F: Fn(&T) -> &str + Copy,
+{
+    groups
+        .iter()
+        .map(|(date, rows)| {
+            let symbol_ids = rows
+                .iter()
+                .map(|row| {
+                    *symbol_id_by_code
+                        .get(symbol_of(row.as_ref()))
+                        .expect("snapshot symbol missing from FIDC symbol index")
+                })
+                .collect::<Vec<_>>();
+            debug_assert!(symbol_ids.windows(2).all(|window| window[0] < window[1]));
+            (*date, symbol_ids)
+        })
+        .collect()
+}
+
+fn find_arc_by_symbol_id<'a, T>(
+    rows: &'a [Arc<T>],
+    symbol_ids: &[u32],
+    symbol_id: u32,
+) -> Option<&'a T> {
+    if rows.len() != symbol_ids.len() {
+        return None;
+    }
+    symbol_ids
+        .binary_search(&symbol_id)
+        .ok()
+        .and_then(|index| rows.get(index))
+        .map(Arc::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,