diff --git a/crates/fidc-core/src/data.rs b/crates/fidc-core/src/data.rs index 0623626..a65002c 100644 --- a/crates/fidc-core/src/data.rs +++ b/crates/fidc-core/src/data.rs @@ -984,13 +984,10 @@ pub struct DataSet { instruments: HashMap, calendar: TradingCalendar, market_by_date: BTreeMap>>, - market_index: HashMap>>, factor_by_date: BTreeMap>>, - factor_index: HashMap>>, factor_text_by_date: BTreeMap>, factor_text_index: HashMap<(NaiveDate, String, String), FactorTextValue>, candidate_by_date: BTreeMap>>, - candidate_index: HashMap>>, corporate_actions_by_date: BTreeMap>, execution_quotes_by_date: HashMap>>, order_book_depth_index: HashMap<(NaiveDate, String), Vec>, @@ -1247,25 +1244,11 @@ impl DataSet { .collect::>(); let market = market.into_iter().map(Arc::new).collect::>(); - let market_by_date = group_arc_by_date(&market, |item| item.date); - let mut market_index = - HashMap::>>::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::>>::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::>(); - let candidate_by_date = group_arc_by_date(&candidates, |item| item.date); - let mut candidate_index = - HashMap::>>::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(groups: &mut BTreeMap>>, 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], 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 { let mut codes = benchmarks .iter() @@ -3672,8 +3646,8 @@ fn build_order_book_depth_index( fn build_eligible_universe( factor_by_date: &BTreeMap>>, - candidate_index: &HashMap>>, - market_index: &HashMap>>, + candidate_by_date: &BTreeMap>>, + market_by_date: &BTreeMap>>, instruments: &HashMap, ) -> BTreeMap> { let mut per_date = BTreeMap::>::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,