perf(core): specialize market candidate snapshot lookup

This commit is contained in:
boris
2026-08-28 08:38:22 +08:00
parent 096115bd94
commit e677a73f95
2 changed files with 58 additions and 1 deletions
+55
View File
@@ -1246,6 +1246,12 @@ pub(crate) struct SymbolSnapshotRefs<'a> {
pub candidate: Option<&'a CandidateEligibility>, pub candidate: Option<&'a CandidateEligibility>,
} }
#[derive(Debug, Clone, Copy)]
pub(crate) struct MarketCandidateSnapshotRefs<'a> {
pub market: Option<&'a DailyMarketSnapshot>,
pub candidate: Option<&'a CandidateEligibility>,
}
impl DataSet { impl DataSet {
pub fn with_additional_trading_dates( pub fn with_additional_trading_dates(
mut self, mut self,
@@ -1711,6 +1717,40 @@ impl DataSet {
} }
} }
pub(crate) fn market_candidate_snapshots_by_id(
&self,
date: NaiveDate,
symbol_id: u32,
) -> MarketCandidateSnapshotRefs<'_> {
let market_rows = self.market_by_date.get(&date).map(Vec::as_slice);
let market_symbol_ids = self
.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 = market_index.and_then(|index| market_rows?.get(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,
)
});
MarketCandidateSnapshotRefs { market, candidate }
}
pub fn benchmark(&self, date: NaiveDate) -> Option<&BenchmarkSnapshot> { pub fn benchmark(&self, date: NaiveDate) -> Option<&BenchmarkSnapshot> {
self.benchmark_by_date.get(&date) self.benchmark_by_date.get(&date)
} }
@@ -4108,11 +4148,16 @@ mod tests {
for symbol in ["000001.SZ", "600000.SH"] { for symbol in ["000001.SZ", "600000.SH"] {
let symbol_id = data.symbol_id(symbol).unwrap(); let symbol_id = data.symbol_id(symbol).unwrap();
let combined = data.symbol_snapshots_by_id(date, symbol_id); let combined = data.symbol_snapshots_by_id(date, symbol_id);
let market_candidate = data.market_candidate_snapshots_by_id(date, symbol_id);
assert_eq!( assert_eq!(
combined.market.map(|row| row.symbol.as_str()), combined.market.map(|row| row.symbol.as_str()),
data.market_by_symbol_id(date, symbol_id) data.market_by_symbol_id(date, symbol_id)
.map(|row| row.symbol.as_str()) .map(|row| row.symbol.as_str())
); );
assert_eq!(
market_candidate.market.map(|row| row.symbol.as_str()),
combined.market.map(|row| row.symbol.as_str())
);
assert_eq!( assert_eq!(
combined.factor.map(|row| row.symbol.as_str()), combined.factor.map(|row| row.symbol.as_str()),
data.factor_by_symbol_id(date, symbol_id) data.factor_by_symbol_id(date, symbol_id)
@@ -4123,6 +4168,10 @@ mod tests {
data.candidate_by_symbol_id(date, symbol_id) data.candidate_by_symbol_id(date, symbol_id)
.map(|row| row.symbol.as_str()) .map(|row| row.symbol.as_str())
); );
assert_eq!(
market_candidate.candidate.map(|row| row.symbol.as_str()),
combined.candidate.map(|row| row.symbol.as_str())
);
} }
let signal_id = data.symbol_id("000300.SH").unwrap(); let signal_id = data.symbol_id("000300.SH").unwrap();
@@ -4130,6 +4179,12 @@ mod tests {
assert_eq!(signal.market.map(|row| row.symbol.as_str()), Some("000300.SH")); assert_eq!(signal.market.map(|row| row.symbol.as_str()), Some("000300.SH"));
assert!(signal.factor.is_none()); assert!(signal.factor.is_none());
assert!(signal.candidate.is_none()); assert!(signal.candidate.is_none());
let signal_market_candidate = data.market_candidate_snapshots_by_id(date, signal_id);
assert_eq!(
signal_market_candidate.market.map(|row| row.symbol.as_str()),
Some("000300.SH")
);
assert!(signal_market_candidate.candidate.is_none());
} }
#[test] #[test]
@@ -8889,7 +8889,9 @@ impl PlatformExprStrategy {
if ctx.has_dynamic_universe() && !ctx.dynamic_universe_contains(&factor.symbol) { if ctx.has_dynamic_universe() && !ctx.dynamic_universe_contains(&factor.symbol) {
continue; continue;
} }
let snapshots = ctx.data.symbol_snapshots_by_id(date, symbol_id); let snapshots = ctx
.data
.market_candidate_snapshots_by_id(date, symbol_id);
let synthetic_candidate; let synthetic_candidate;
let candidate = if let Some(candidate) = snapshots.candidate { let candidate = if let Some(candidate) = snapshots.candidate {
candidate candidate