perf: share immutable market-cap order index

This commit is contained in:
boris
2026-09-05 00:37:39 +08:00
parent 6538742dfa
commit 05953f857a
2 changed files with 103 additions and 23 deletions
+82 -14
View File
@@ -1309,6 +1309,7 @@ pub struct DataSet {
factor_by_date: Arc<BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>>,
factor_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
factor_row_positions_by_date: Arc<Option<DenseRowPositionIndex>>,
factor_market_cap_order_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
factor_text_by_date: Arc<BTreeMap<NaiveDate, Vec<FactorTextValue>>>,
factor_text_index: Arc<HashMap<(NaiveDate, String, String), FactorTextValue>>,
candidate_by_date: Arc<BTreeMap<NaiveDate, Vec<CandidateEligibility>>>,
@@ -1357,8 +1358,7 @@ impl<'a, T> DailySymbolRows<'a, T> {
/// the already indexed slices once and keeps all lookups read-only.
pub(crate) struct DailySnapshotView<'a> {
market: DailySymbolRows<'a, DailyMarketSnapshot>,
factor_rows: &'a [DailyFactorSnapshot],
factor_symbol_ids: &'a [u32],
factors: DailySymbolRows<'a, DailyFactorSnapshot>,
candidates: DailySymbolRows<'a, CandidateEligibility>,
}
@@ -1371,12 +1371,16 @@ impl<'a> DailySnapshotView<'a> {
self.candidates.get(symbol_id)
}
pub(crate) fn factor(&self, symbol_id: u32) -> Option<&'a DailyFactorSnapshot> {
self.factors.get(symbol_id)
}
pub(crate) fn factor_rows(&self) -> &'a [DailyFactorSnapshot] {
self.factor_rows
self.factors.rows
}
pub(crate) fn factor_symbol_ids(&self) -> &'a [u32] {
self.factor_symbol_ids
self.factors.symbol_ids
}
}
@@ -1781,6 +1785,8 @@ impl DataSet {
build_group_symbol_ids(&factor_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
});
let factor_market_cap_order_by_date =
build_factor_market_cap_order(&factor_by_date, &factor_symbol_ids_by_date);
let candidate_symbol_ids_by_date =
build_group_symbol_ids(&candidate_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
@@ -1833,6 +1839,7 @@ impl DataSet {
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_market_cap_order_by_date: Arc::new(factor_market_cap_order_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),
@@ -1905,6 +1912,17 @@ impl DataSet {
self.symbol_id_by_code.get(symbol).copied()
}
pub(crate) fn symbol_count(&self) -> usize {
self.symbol_id_by_code.len()
}
pub(crate) fn factor_symbol_ids_by_market_cap_on(&self, date: NaiveDate) -> &[u32] {
self.factor_market_cap_order_by_date
.get(&date)
.map(Vec::as_slice)
.unwrap_or(&[])
}
pub fn market(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
let symbol_id = self.symbol_id(symbol)?;
self.market_by_symbol_id(date, symbol_id)
@@ -1950,16 +1968,12 @@ impl DataSet {
&self.market_symbol_ids_by_date,
&self.market_row_positions_by_date,
),
factor_rows: self
.factor_by_date
.get(&date)
.map(Vec::as_slice)
.unwrap_or(&[]),
factor_symbol_ids: self
.factor_symbol_ids_by_date
.get(&date)
.map(Vec::as_slice)
.unwrap_or(&[]),
factors: rows_on(
date,
&self.factor_by_date,
&self.factor_symbol_ids_by_date,
&self.factor_row_positions_by_date,
),
candidates: rows_on(
date,
&self.candidate_by_date,
@@ -4202,6 +4216,48 @@ where
.collect()
}
fn build_factor_market_cap_order(
factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
factor_symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
) -> BTreeMap<NaiveDate, Vec<u32>> {
factor_by_date
.par_iter()
.map(|(date, rows)| {
let symbol_ids = factor_symbol_ids_by_date
.get(date)
.expect("factor symbol ids missing for market-cap order");
assert_eq!(
rows.len(),
symbol_ids.len(),
"factor rows and symbol ids diverged for {date}"
);
let mut row_indices = rows
.iter()
.enumerate()
.filter_map(|(index, row)| {
let market_cap_bn = decision_market_cap_bn(row);
(market_cap_bn.is_finite() && market_cap_bn > 0.0).then_some(index)
})
.collect::<Vec<_>>();
row_indices.sort_by(|left, right| {
let left = &rows[*left];
let right = &rows[*right];
decision_market_cap_bn(left)
.partial_cmp(&decision_market_cap_bn(right))
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| left.symbol.cmp(&right.symbol))
});
let ordered = row_indices
.into_iter()
.map(|index| symbol_ids[index])
.collect::<Vec<_>>();
(*date, ordered)
})
.collect::<Vec<_>>()
.into_iter()
.collect()
}
fn build_dense_row_positions<T>(
groups: &BTreeMap<NaiveDate, Vec<T>>,
symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
@@ -5005,6 +5061,10 @@ mod tests {
.map(|row| row.symbol.as_str()),
Some(symbol)
);
assert_eq!(
day.factor(symbol_id).map(|row| row.symbol.as_str()),
Some(symbol)
);
assert_eq!(
data.candidate_by_symbol_id(date, symbol_id)
.map(|row| row.symbol.as_str()),
@@ -5024,6 +5084,7 @@ mod tests {
Some("000300.SH")
);
assert!(data.factor_by_symbol_id(date, signal_id).is_none());
assert!(day.factor(signal_id).is_none());
assert!(data.candidate_by_symbol_id(date, signal_id).is_none());
assert!(day.candidate(signal_id).is_none());
assert_eq!(
@@ -6373,6 +6434,13 @@ mod tests {
assert_eq!(rows[1].symbol, "000001.SZ");
assert!((rows[1].market_cap_bn - 12.0).abs() < 1e-9);
assert!((rows[1].free_float_cap_bn - 4.0).abs() < 1e-9);
assert_eq!(
data.factor_symbol_ids_by_market_cap_on(date),
&[
data.symbol_id("000002.SZ").unwrap(),
data.symbol_id("000001.SZ").unwrap(),
]
);
}
#[test]
+21 -9
View File
@@ -9247,6 +9247,7 @@ impl PlatformExprStrategy {
) -> (Vec<EligibleUniverseSnapshot>, Vec<FidcRiskDecisionAudit>) {
let mut rows = Vec::new();
let mut decisions = Vec::new();
let mut eligible_symbols = vec![false; ctx.data.symbol_count()];
let execution_day = ctx.data.daily_snapshot_view(date);
let factor_day = ctx.data.daily_snapshot_view(factor_date);
let factor_rows = factor_day.factor_rows();
@@ -9312,19 +9313,30 @@ impl PlatformExprStrategy {
if market_cap_bn <= 0.0 || !market_cap_bn.is_finite() {
continue;
}
let free_float_cap_bn = decision_free_float_cap_bn(factor);
eligible_symbols[symbol_id as usize] = true;
}
for symbol_id in ctx
.data
.factor_symbol_ids_by_market_cap_on(factor_date)
.iter()
.copied()
{
if !eligible_symbols
.get(symbol_id as usize)
.copied()
.unwrap_or(false)
{
continue;
}
let factor = factor_day
.factor(symbol_id)
.expect("market-cap order references missing factor row");
rows.push(EligibleUniverseSnapshot {
symbol: factor.symbol.clone(),
market_cap_bn,
free_float_cap_bn,
market_cap_bn: decision_market_cap_bn(factor),
free_float_cap_bn: decision_free_float_cap_bn(factor),
});
}
rows.sort_by(|left, right| {
left.market_cap_bn
.partial_cmp(&right.market_cap_bn)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| left.symbol.cmp(&right.symbol))
});
(rows, decisions)
}