优化回测数据集索引查找

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>,
calendar: TradingCalendar,
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_index: HashMap<(NaiveDate, String), Arc<DailyFactorSnapshot>>,
factor_index: HashMap<NaiveDate, HashMap<String, Arc<DailyFactorSnapshot>>>,
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_index: HashMap<(NaiveDate, String), Arc<CandidateEligibility>>,
candidate_index: HashMap<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>,
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>>,
@@ -1229,16 +1229,23 @@ impl DataSet {
.collect::<HashMap<_, _>>();
let market_by_date = group_by_date(market.clone(), |item| item.date);
let market_index = market
.into_iter()
.map(|item| ((item.date, item.symbol.clone()), item))
.collect::<HashMap<_, _>>();
let mut market_index = HashMap::<NaiveDate, HashMap<String, DailyMarketSnapshot>>::new();
for item in market {
market_index
.entry(item.date)
.or_default()
.insert(item.symbol.clone(), item);
}
let factor_by_date = group_arc_by_date(&factors, |item| item.date);
let factor_index = factors
.into_iter()
.map(|item| ((item.date, item.symbol.clone()), item))
.collect::<HashMap<_, _>>();
let mut factor_index =
HashMap::<NaiveDate, HashMap<String, Arc<DailyFactorSnapshot>>>::new();
for item in factors {
factor_index
.entry(item.date)
.or_default()
.insert(item.symbol.clone(), item);
}
let factor_texts = factor_texts
.into_iter()
.filter_map(|mut item| {
@@ -1257,10 +1264,14 @@ impl DataSet {
.collect::<HashMap<_, _>>();
let candidate_by_date = group_arc_by_date(&candidates, |item| item.date);
let candidate_index = candidates
.into_iter()
.map(|item| ((item.date, item.symbol.clone()), item))
.collect::<HashMap<_, _>>();
let mut candidate_index =
HashMap::<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>::new();
for item in candidates {
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 execution_quotes_by_date = build_execution_quote_index(execution_quotes);
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> {
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> {
self.factor_index
.get(&(date, symbol.to_string()))
.get(&date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref)
}
pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> {
self.candidate_index
.get(&(date, symbol.to_string()))
.get(&date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref)
}
@@ -2154,7 +2169,9 @@ impl DataSet {
return None;
}
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> {
@@ -2276,9 +2293,10 @@ impl DataSet {
.map(|day| {
evaluator(
self.candidate_index
.get(&(*day, symbol.to_string()))
.get(day)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref),
self.market_index.get(&(*day, symbol.to_string())),
self.market_index.get(day).and_then(|rows| rows.get(symbol)),
)
})
.collect()
@@ -3615,8 +3633,8 @@ fn build_order_book_depth_index(
fn build_eligible_universe(
factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
candidate_index: &HashMap<(NaiveDate, String), Arc<CandidateEligibility>>,
market_index: &HashMap<(NaiveDate, String), DailyMarketSnapshot>,
candidate_index: &HashMap<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>,
market_index: &HashMap<NaiveDate, HashMap<String, DailyMarketSnapshot>>,
instruments: &HashMap<String, Instrument>,
) -> BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>> {
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() {
continue;
}
let key = (*date, factor.symbol.clone());
let Some(candidate) = candidate_index.get(&key) else {
let Some(candidate) = candidate_index
.get(date)
.and_then(|rows| rows.get(&factor.symbol))
else {
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;
};
if ChinaAShareRiskControl::selection_rejection_reason(