diff --git a/crates/fidc-core/src/data.rs b/crates/fidc-core/src/data.rs index 47ff342..ffa36ef 100644 --- a/crates/fidc-core/src/data.rs +++ b/crates/fidc-core/src/data.rs @@ -566,6 +566,11 @@ struct SymbolPriceSeries { valid_volume_start_by_count: Vec, } +type DenseRowPositionIndex = BTreeMap>; + +const MISSING_ROW_POSITION: u32 = u32::MAX; +const MAX_DENSE_ROW_INDEX_BYTES: usize = 256 * 1024 * 1024; + #[derive(Debug, Clone)] struct AdjustedCloseSeries { dates: Vec, @@ -1243,12 +1248,15 @@ pub struct DataSet { calendar: Arc, market_by_date: Arc>>, market_symbol_ids_by_date: Arc>>, + market_row_positions_by_date: Arc>, factor_by_date: Arc>>, factor_symbol_ids_by_date: Arc>>, + factor_row_positions_by_date: Arc>, factor_text_by_date: Arc>>, factor_text_index: Arc>, candidate_by_date: Arc>>, candidate_symbol_ids_by_date: Arc>>, + candidate_row_positions_by_date: Arc>, corporate_actions_by_date: Arc>>, execution_quotes_by_date: Arc>>>, execution_quote_dates: Arc>, @@ -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 - .zip(market_symbol_ids) - .and_then(|(rows, symbol_ids)| symbol_id_index(rows.len(), symbol_ids, symbol_id)); + let market_index = dense_row_position( + &self.market_row_positions_by_date, + 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 - .factor_by_date - .get(&date) - .map(Vec::as_slice) - .zip( - self.factor_symbol_ids_by_date - .get(&date) - .map(Vec::as_slice), - ) - .and_then(|(rows, symbol_ids)| { - find_by_symbol_id_with_preferred_index( - rows, - symbol_ids, - symbol_id, - market_index, - ) - }); - let candidate = self - .candidate_by_date - .get(&date) - .map(Vec::as_slice) - .zip( - self.candidate_symbol_ids_by_date - .get(&date) - .map(Vec::as_slice), - ) - .and_then(|(rows, symbol_ids)| { - find_by_symbol_id_with_preferred_index( - rows, - symbol_ids, - symbol_id, - market_index, - ) - }); + let factor = self.factor_by_date.get(&date).and_then(|rows| { + dense_row_position(&self.factor_row_positions_by_date, date, symbol_id) + .and_then(|index| rows.get(index)) + .or_else(|| { + self.factor_symbol_ids_by_date + .get(&date) + .and_then(|symbol_ids| { + find_by_symbol_id_with_preferred_index( + rows, + symbol_ids, + symbol_id, + market_index, + ) + }) + }) + }); + let candidate = self.candidate_by_date.get(&date).and_then(|rows| { + dense_row_position(&self.candidate_row_positions_by_date, date, symbol_id) + .and_then(|index| rows.get(index)) + .or_else(|| { + self.candidate_symbol_ids_by_date + .get(&date) + .and_then(|symbol_ids| { + find_by_symbol_id_with_preferred_index( + rows, + symbol_ids, + symbol_id, + market_index, + ) + }) + }) + }); SymbolSnapshotRefs { market, @@ -3807,6 +3854,49 @@ where .collect() } +fn build_dense_row_positions( + groups: &BTreeMap>, + symbol_ids_by_date: &BTreeMap>, + symbol_count: usize, +) -> Option { + let entries = groups.len().checked_mul(symbol_count)?; + let bytes = entries.checked_mul(std::mem::size_of::())?; + 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, + date: NaiveDate, + symbol_id: u32, +) -> Option { + 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