加速按日股票快照查找

crates/fidc-core/src/data.rs

@@ -566,6 +566,11 @@ struct SymbolPriceSeries {
     valid_volume_start_by_count: Vec<usize>,
 }
 
+type DenseRowPositionIndex = BTreeMap<NaiveDate, Vec<u32>>;
+
+const MISSING_ROW_POSITION: u32 = u32::MAX;
+const MAX_DENSE_ROW_INDEX_BYTES: usize = 256 * 1024 * 1024;
+
 #[derive(Debug, Clone)]
 struct AdjustedCloseSeries {
     dates: Vec<NaiveDate>,
@@ -1243,12 +1248,15 @@ pub struct DataSet {
     calendar: Arc<TradingCalendar>,
     market_by_date: Arc<BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>>,
     market_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
+    market_row_positions_by_date: Arc<Option<DenseRowPositionIndex>>,
     factor_by_date: Arc<BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>>,
     factor_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
+    factor_row_positions_by_date: Arc<Option<DenseRowPositionIndex>>,
     factor_text_by_date: Arc<BTreeMap<NaiveDate, Vec<FactorTextValue>>>,
     factor_text_index: Arc<HashMap<(NaiveDate, String, String), FactorTextValue>>,
     candidate_by_date: Arc<BTreeMap<NaiveDate, Vec<CandidateEligibility>>>,
     candidate_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
+    candidate_row_positions_by_date: Arc<Option<DenseRowPositionIndex>>,
     corporate_actions_by_date: Arc<BTreeMap<NaiveDate, Vec<CorporateAction>>>,
     execution_quotes_by_date: Arc<HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>>,
     execution_quote_dates: Arc<Vec<NaiveDate>>,
@@ -1660,6 +1668,21 @@ impl DataSet {
             build_group_symbol_ids(&candidate_by_date, &symbol_id_by_code, |item| {
                 item.symbol.as_str()
             });
+        let market_row_positions_by_date = build_dense_row_positions(
+            &market_by_date,
+            &market_symbol_ids_by_date,
+            symbol_id_by_code.len(),
+        );
+        let factor_row_positions_by_date = build_dense_row_positions(
+            &factor_by_date,
+            &factor_symbol_ids_by_date,
+            symbol_id_by_code.len(),
+        );
+        let candidate_row_positions_by_date = build_dense_row_positions(
+            &candidate_by_date,
+            &candidate_symbol_ids_by_date,
+            symbol_id_by_code.len(),
+        );
         let mut market_series_by_symbol_id = vec![None; symbol_id_by_code.len()];
         for (symbol, series) in &market_series_by_symbol {
             if let Some(symbol_id) = symbol_id_by_code.get(symbol).copied() {
@@ -1685,12 +1708,15 @@ impl DataSet {
             calendar: Arc::new(calendar),
             market_by_date: Arc::new(market_by_date),
             market_symbol_ids_by_date: Arc::new(market_symbol_ids_by_date),
+            market_row_positions_by_date: Arc::new(market_row_positions_by_date),
             factor_by_date: Arc::new(factor_by_date),
             factor_symbol_ids_by_date: Arc::new(factor_symbol_ids_by_date),
+            factor_row_positions_by_date: Arc::new(factor_row_positions_by_date),
             factor_text_by_date: Arc::new(factor_text_by_date),
             factor_text_index: Arc::new(factor_text_index),
             candidate_by_date: Arc::new(candidate_by_date),
             candidate_symbol_ids_by_date: Arc::new(candidate_symbol_ids_by_date),
+            candidate_row_positions_by_date: Arc::new(candidate_row_positions_by_date),
             corporate_actions_by_date: Arc::new(corporate_actions_by_date),
             execution_quotes_by_date: Arc::new(execution_quotes_by_date),
             execution_quote_dates: Arc::new(execution_quote_dates),
@@ -1759,8 +1785,13 @@ impl DataSet {
         date: NaiveDate,
         symbol_id: u32,
     ) -> Option<&DailyMarketSnapshot> {
+        let rows = self.market_by_date.get(&date)?;
+        if let Some(index) = dense_row_position(&self.market_row_positions_by_date, date, symbol_id)
+        {
+            return rows.get(index);
+        }
         find_by_symbol_id(
-            self.market_by_date.get(&date)?,
+            rows,
             self.market_symbol_ids_by_date.get(&date)?,
             symbol_id,
         )
@@ -1798,8 +1829,13 @@ impl DataSet {
         date: NaiveDate,
         symbol_id: u32,
     ) -> Option<&DailyFactorSnapshot> {
+        let rows = self.factor_by_date.get(&date)?;
+        if let Some(index) = dense_row_position(&self.factor_row_positions_by_date, date, symbol_id)
+        {
+            return rows.get(index);
+        }
         find_by_symbol_id(
-            self.factor_by_date.get(&date)?,
+            rows,
             self.factor_symbol_ids_by_date.get(&date)?,
             symbol_id,
         )
@@ -1815,8 +1851,14 @@ impl DataSet {
         date: NaiveDate,
         symbol_id: u32,
     ) -> Option<&CandidateEligibility> {
+        let rows = self.candidate_by_date.get(&date)?;
+        if let Some(index) =
+            dense_row_position(&self.candidate_row_positions_by_date, date, symbol_id)
+        {
+            return rows.get(index);
+        }
         find_by_symbol_id(
-            self.candidate_by_date.get(&date)?,
+            rows,
             self.candidate_symbol_ids_by_date.get(&date)?,
             symbol_id,
         )
@@ -1832,45 +1874,50 @@ impl DataSet {
             .market_symbol_ids_by_date
             .get(&date)
             .map(Vec::as_slice);
-        let market_index = market_rows
+        let market_index = dense_row_position(
-            .zip(market_symbol_ids)
+            &self.market_row_positions_by_date,
-            .and_then(|(rows, symbol_ids)| symbol_id_index(rows.len(), symbol_ids, symbol_id));
+            date,
+            symbol_id,
+        )
+        .or_else(|| {
+            market_rows
+                .zip(market_symbol_ids)
+                .and_then(|(rows, symbol_ids)| symbol_id_index(rows.len(), symbol_ids, symbol_id))
+        });
         let market = market_index.and_then(|index| market_rows?.get(index));
 
-        let factor = self
+        let factor = self.factor_by_date.get(&date).and_then(|rows| {
-            .factor_by_date
+            dense_row_position(&self.factor_row_positions_by_date, date, symbol_id)
-            .get(&date)
+                .and_then(|index| rows.get(index))
-            .map(Vec::as_slice)
+                .or_else(|| {
-            .zip(
+                    self.factor_symbol_ids_by_date
-                self.factor_symbol_ids_by_date
+                        .get(&date)
-                    .get(&date)
+                        .and_then(|symbol_ids| {
-                    .map(Vec::as_slice),
+                            find_by_symbol_id_with_preferred_index(
-            )
+                                rows,
-            .and_then(|(rows, symbol_ids)| {
+                                symbol_ids,
-                find_by_symbol_id_with_preferred_index(
+                                symbol_id,
-                    rows,
+                                market_index,
-                    symbol_ids,
+                            )
-                    symbol_id,
+                        })
-                    market_index,
+                })
-                )
+        });
-            });
+        let candidate = self.candidate_by_date.get(&date).and_then(|rows| {
-        let candidate = self
+            dense_row_position(&self.candidate_row_positions_by_date, date, symbol_id)
-            .candidate_by_date
+                .and_then(|index| rows.get(index))
-            .get(&date)
+                .or_else(|| {
-            .map(Vec::as_slice)
+                    self.candidate_symbol_ids_by_date
-            .zip(
+                        .get(&date)
-                self.candidate_symbol_ids_by_date
+                        .and_then(|symbol_ids| {
-                    .get(&date)
+                            find_by_symbol_id_with_preferred_index(
-                    .map(Vec::as_slice),
+                                rows,
-            )
+                                symbol_ids,
-            .and_then(|(rows, symbol_ids)| {
+                                symbol_id,
-                find_by_symbol_id_with_preferred_index(
+                                market_index,
-                    rows,
+                            )
-                    symbol_ids,
+                        })
-                    symbol_id,
+                })
-                    market_index,
+        });
-                )
-            });
 
         SymbolSnapshotRefs {
             market,
@@ -3807,6 +3854,49 @@ where
         .collect()
 }
 
+fn build_dense_row_positions<T>(
+    groups: &BTreeMap<NaiveDate, Vec<T>>,
+    symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
+    symbol_count: usize,
+) -> Option<DenseRowPositionIndex> {
+    let entries = groups.len().checked_mul(symbol_count)?;
+    let bytes = entries.checked_mul(std::mem::size_of::<u32>())?;
+    if bytes > MAX_DENSE_ROW_INDEX_BYTES {
+        return None;
+    }
+
+    let mut positions_by_date = BTreeMap::new();
+    for (date, rows) in groups {
+        let symbol_ids = symbol_ids_by_date.get(date)?;
+        if rows.len() != symbol_ids.len() {
+            return None;
+        }
+        let mut positions = vec![MISSING_ROW_POSITION; symbol_count];
+        for (row_index, symbol_id) in symbol_ids.iter().copied().enumerate() {
+            let position = positions.get_mut(usize::try_from(symbol_id).ok()?)?;
+            if *position != MISSING_ROW_POSITION {
+                return None;
+            }
+            *position = u32::try_from(row_index).ok()?;
+        }
+        positions_by_date.insert(*date, positions);
+    }
+    Some(positions_by_date)
+}
+
+fn dense_row_position(
+    positions_by_date: &Option<DenseRowPositionIndex>,
+    date: NaiveDate,
+    symbol_id: u32,
+) -> Option<usize> {
+    let position = positions_by_date
+        .as_ref()?
+        .get(&date)?
+        .get(usize::try_from(symbol_id).ok()?)
+        .copied()?;
+    (position != MISSING_ROW_POSITION).then_some(position as usize)
+}
+
 fn find_by_symbol_id<'a, T>(rows: &'a [T], symbol_ids: &[u32], symbol_id: u32) -> Option<&'a T> {
     find_by_symbol_id_with_preferred_index(rows, symbol_ids, symbol_id, None)
 }
@@ -4191,11 +4281,23 @@ mod tests {
 
         assert!(Arc::ptr_eq(&data.instruments, &run_data.instruments));
         assert!(Arc::ptr_eq(&data.market_by_date, &run_data.market_by_date));
+        assert!(Arc::ptr_eq(
+            &data.market_row_positions_by_date,
+            &run_data.market_row_positions_by_date
+        ));
         assert!(Arc::ptr_eq(&data.factor_by_date, &run_data.factor_by_date));
+        assert!(Arc::ptr_eq(
+            &data.factor_row_positions_by_date,
+            &run_data.factor_row_positions_by_date
+        ));
         assert!(Arc::ptr_eq(
             &data.candidate_by_date,
             &run_data.candidate_by_date
         ));
+        assert!(Arc::ptr_eq(
+            &data.candidate_row_positions_by_date,
+            &run_data.candidate_row_positions_by_date
+        ));
         assert!(Arc::ptr_eq(
             &data.benchmark_by_date,
             &run_data.benchmark_by_date