优化回测数据集索引查找

This commit is contained in:
boris
2026-06-18 20:09:25 +08:00
parent 3633905459
commit 02d4ea9ca7
+49 -26
View File
@@ -966,13 +966,13 @@ pub struct DataSet {
instruments: HashMap<String, Instrument>, instruments: HashMap<String, Instrument>,
calendar: TradingCalendar, calendar: TradingCalendar,
market_by_date: BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>, 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_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_by_date: BTreeMap<NaiveDate, Vec<FactorTextValue>>,
factor_text_index: HashMap<(NaiveDate, String, String), FactorTextValue>, factor_text_index: HashMap<(NaiveDate, String, String), FactorTextValue>,
candidate_by_date: BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>, 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>>, corporate_actions_by_date: BTreeMap<NaiveDate, Vec<CorporateAction>>,
execution_quotes_by_date: HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>, execution_quotes_by_date: HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>,
order_book_depth_index: HashMap<(NaiveDate, String), Vec<IntradayOrderBookDepthLevel>>, order_book_depth_index: HashMap<(NaiveDate, String), Vec<IntradayOrderBookDepthLevel>>,
@@ -1229,16 +1229,23 @@ impl DataSet {
.collect::<HashMap<_, _>>(); .collect::<HashMap<_, _>>();
let market_by_date = group_by_date(market.clone(), |item| item.date); 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_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 let factor_texts = factor_texts
.into_iter() .into_iter()
.filter_map(|mut item| { .filter_map(|mut item| {
@@ -1257,10 +1264,14 @@ impl DataSet {
.collect::<HashMap<_, _>>(); .collect::<HashMap<_, _>>();
let candidate_by_date = group_arc_by_date(&candidates, |item| item.date); 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 corporate_actions_by_date = group_by_date(corporate_actions, |item| item.date);
let execution_quotes_by_date = build_execution_quote_index(execution_quotes); let execution_quotes_by_date = build_execution_quote_index(execution_quotes);
let order_book_depth_index = build_order_book_depth_index(order_book_depth); 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> { 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> { pub fn factor(&self, date: NaiveDate, symbol: &str) -> Option<&DailyFactorSnapshot> {
self.factor_index self.factor_index
.get(&(date, symbol.to_string())) .get(&date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref) .map(Arc::as_ref)
} }
pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> { pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> {
self.candidate_index self.candidate_index
.get(&(date, symbol.to_string())) .get(&date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref) .map(Arc::as_ref)
} }
@@ -2154,7 +2169,9 @@ impl DataSet {
return None; return None;
} }
let previous_date = *series.dates.get(end - 1)?; 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> { pub fn factor_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyFactorSnapshot> {
@@ -2276,9 +2293,10 @@ impl DataSet {
.map(|day| { .map(|day| {
evaluator( evaluator(
self.candidate_index self.candidate_index
.get(&(*day, symbol.to_string())) .get(day)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref), .map(Arc::as_ref),
self.market_index.get(&(*day, symbol.to_string())), self.market_index.get(day).and_then(|rows| rows.get(symbol)),
) )
}) })
.collect() .collect()
@@ -3615,8 +3633,8 @@ fn build_order_book_depth_index(
fn build_eligible_universe( fn build_eligible_universe(
factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>, 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>, instruments: &HashMap<String, Instrument>,
) -> BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>> { ) -> BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>> {
let mut per_date = BTreeMap::<NaiveDate, Vec<EligibleUniverseSnapshot>>::new(); 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() { if factor.market_cap_bn <= 0.0 || !factor.market_cap_bn.is_finite() {
continue; 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; 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; continue;
}; };
if ChinaAShareRiskControl::selection_rejection_reason( if ChinaAShareRiskControl::selection_rejection_reason(