移除DataSet行级Arc分配

This commit is contained in:
boris
2026-08-24 21:53:00 +08:00
parent 7503dc8517
commit 4cf0224d2d
+50 -83
View File
@@ -503,7 +503,7 @@ struct AdjustedCloseSeries {
impl AdjustedCloseSeries { impl AdjustedCloseSeries {
fn new( fn new(
market: &SymbolPriceSeries, market: &SymbolPriceSeries,
factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>, factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
) -> Option<Self> { ) -> Option<Self> {
let mut backward_factors = Vec::with_capacity(market.dates.len()); let mut backward_factors = Vec::with_capacity(market.dates.len());
let mut back_adjusted_closes = Vec::with_capacity(market.dates.len()); let mut back_adjusted_closes = Vec::with_capacity(market.dates.len());
@@ -514,9 +514,7 @@ impl AdjustedCloseSeries {
for (date, close) in market.dates.iter().zip(&market.closes) { for (date, close) in market.dates.iter().zip(&market.closes) {
let factor = factor_by_date let factor = factor_by_date
.get(date) .get(date)
.and_then(|rows| { .and_then(|rows| find_by_symbol(rows, &market.symbol, |row| row.symbol.as_str()))
find_arc_by_symbol(rows, &market.symbol, |row| row.symbol.as_str())
})
.and_then(|snapshot| factor_numeric_value(snapshot, "adjustment_factor_backward1")) .and_then(|snapshot| factor_numeric_value(snapshot, "adjustment_factor_backward1"))
.filter(|factor| factor.is_finite() && *factor > 0.0); .filter(|factor| factor.is_finite() && *factor > 0.0);
let back_adjusted_close = factor let back_adjusted_close = factor
@@ -1125,13 +1123,13 @@ impl BenchmarkPriceSeries {
pub struct DataSet { pub struct DataSet {
instruments: Arc<HashMap<String, Instrument>>, instruments: Arc<HashMap<String, Instrument>>,
calendar: Arc<TradingCalendar>, calendar: Arc<TradingCalendar>,
market_by_date: Arc<BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>>, market_by_date: Arc<BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>>,
market_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>, market_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
factor_by_date: Arc<BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>>, factor_by_date: Arc<BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>>,
factor_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>, factor_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
factor_text_by_date: Arc<BTreeMap<NaiveDate, Vec<FactorTextValue>>>, factor_text_by_date: Arc<BTreeMap<NaiveDate, Vec<FactorTextValue>>>,
factor_text_index: Arc<HashMap<(NaiveDate, String, String), FactorTextValue>>, factor_text_index: Arc<HashMap<(NaiveDate, String, String), FactorTextValue>>,
candidate_by_date: Arc<BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>>, candidate_by_date: Arc<BTreeMap<NaiveDate, Vec<CandidateEligibility>>>,
candidate_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>, candidate_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
corporate_actions_by_date: Arc<BTreeMap<NaiveDate, Vec<CorporateAction>>>, corporate_actions_by_date: Arc<BTreeMap<NaiveDate, Vec<CorporateAction>>>,
execution_quotes_by_date: HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>, execution_quotes_by_date: HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>,
@@ -1292,26 +1290,23 @@ impl DataSet {
let benchmark_code = collect_benchmark_code(&benchmarks)?; let benchmark_code = collect_benchmark_code(&benchmarks)?;
let calendar = TradingCalendar::new(benchmarks.iter().map(|item| item.date).collect()); let calendar = TradingCalendar::new(benchmarks.iter().map(|item| item.date).collect());
let factors = normalize_factor_snapshots(factors); let factors = normalize_factor_snapshots(factors);
let factors = factors.into_iter().map(Arc::new).collect::<Vec<_>>();
let candidates = candidates.into_iter().map(Arc::new).collect::<Vec<_>>();
let instruments = instruments let instruments = instruments
.into_iter() .into_iter()
.map(|instrument| (instrument.symbol.clone(), instrument)) .map(|instrument| (instrument.symbol.clone(), instrument))
.collect::<HashMap<_, _>>(); .collect::<HashMap<_, _>>();
let market = market.into_iter().map(Arc::new).collect::<Vec<_>>(); let mut market_by_date = group_by_date(market, |item| item.date);
let mut market_by_date = group_arc_by_date(&market, |item| item.date); sort_groups_by_symbol(&mut market_by_date, |item| item.symbol.as_str());
sort_arc_groups_by_symbol(&mut market_by_date, |item| item.symbol.as_str());
let mut factor_by_date = group_arc_by_date(&factors, |item| item.date); let mut factor_by_date = group_by_date(factors, |item| item.date);
sort_arc_groups_by_symbol(&mut factor_by_date, |item| item.symbol.as_str()); sort_groups_by_symbol(&mut factor_by_date, |item| item.symbol.as_str());
let mut market_rows_by_symbol = AHashMap::<String, Vec<&DailyMarketSnapshot>>::new(); let mut market_rows_by_symbol = AHashMap::<String, Vec<&DailyMarketSnapshot>>::new();
for row in &market { for row in market_by_date.values().flatten() {
market_rows_by_symbol market_rows_by_symbol
.entry(row.symbol.clone()) .entry(row.symbol.clone())
.or_default() .or_default()
.push(row.as_ref()); .push(row);
} }
let market_rows_by_symbol = market_rows_by_symbol.into_iter().collect::<Vec<_>>(); let market_rows_by_symbol = market_rows_by_symbol.into_iter().collect::<Vec<_>>();
let market_series_by_symbol = market_rows_by_symbol let market_series_by_symbol = market_rows_by_symbol
@@ -1349,8 +1344,8 @@ impl DataSet {
.map(|item| ((item.date, item.symbol.clone(), item.field.clone()), item)) .map(|item| ((item.date, item.symbol.clone(), item.field.clone()), item))
.collect::<HashMap<_, _>>(); .collect::<HashMap<_, _>>();
let mut candidate_by_date = group_arc_by_date(&candidates, |item| item.date); let mut candidate_by_date = group_by_date(candidates, |item| item.date);
sort_arc_groups_by_symbol(&mut candidate_by_date, |item| item.symbol.as_str()); sort_groups_by_symbol(&mut candidate_by_date, |item| item.symbol.as_str());
let symbol_id_by_code = build_symbol_id_index( let symbol_id_by_code = build_symbol_id_index(
&instruments, &instruments,
&market_by_date, &market_by_date,
@@ -1471,7 +1466,7 @@ impl DataSet {
date: NaiveDate, date: NaiveDate,
symbol_id: u32, symbol_id: u32,
) -> Option<&DailyMarketSnapshot> { ) -> Option<&DailyMarketSnapshot> {
find_arc_by_symbol_id( find_by_symbol_id(
self.market_by_date.get(&date)?, self.market_by_date.get(&date)?,
self.market_symbol_ids_by_date.get(&date)?, self.market_symbol_ids_by_date.get(&date)?,
symbol_id, symbol_id,
@@ -1510,7 +1505,7 @@ impl DataSet {
date: NaiveDate, date: NaiveDate,
symbol_id: u32, symbol_id: u32,
) -> Option<&DailyFactorSnapshot> { ) -> Option<&DailyFactorSnapshot> {
find_arc_by_symbol_id( find_by_symbol_id(
self.factor_by_date.get(&date)?, self.factor_by_date.get(&date)?,
self.factor_symbol_ids_by_date.get(&date)?, self.factor_symbol_ids_by_date.get(&date)?,
symbol_id, symbol_id,
@@ -1527,7 +1522,7 @@ impl DataSet {
date: NaiveDate, date: NaiveDate,
symbol_id: u32, symbol_id: u32,
) -> Option<&CandidateEligibility> { ) -> Option<&CandidateEligibility> {
find_arc_by_symbol_id( find_by_symbol_id(
self.candidate_by_date.get(&date)?, self.candidate_by_date.get(&date)?,
self.candidate_symbol_ids_by_date.get(&date)?, self.candidate_symbol_ids_by_date.get(&date)?,
symbol_id, symbol_id,
@@ -1645,17 +1640,17 @@ impl DataSet {
let market = self let market = self
.market_by_date .market_by_date
.values() .values()
.flat_map(|rows| rows.iter().map(|row| row.as_ref().clone())) .flat_map(|rows| rows.iter().cloned())
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let factors = self let factors = self
.factor_by_date .factor_by_date
.values() .values()
.flat_map(|rows| rows.iter().map(|row| row.as_ref().clone())) .flat_map(|rows| rows.iter().cloned())
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let candidates = self let candidates = self
.candidate_by_date .candidate_by_date
.values() .values()
.flat_map(|rows| rows.iter().map(|row| row.as_ref().clone())) .flat_map(|rows| rows.iter().cloned())
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let benchmarks = self.benchmark_by_date.values().cloned().collect::<Vec<_>>(); let benchmarks = self.benchmark_by_date.values().cloned().collect::<Vec<_>>();
let corporate_actions = self let corporate_actions = self
@@ -2286,7 +2281,6 @@ impl DataSet {
.range(start..=end) .range(start..=end)
.flat_map(|(_, rows)| rows.iter()) .flat_map(|(_, rows)| rows.iter())
.filter(|row| row.symbol == symbol) .filter(|row| row.symbol == symbol)
.map(Arc::as_ref)
.map(daily_market_price_bar) .map(daily_market_price_bar)
.collect(), .collect(),
Some("1m") => { Some("1m") => {
@@ -2337,11 +2331,11 @@ impl DataSet {
pub fn factor_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyFactorSnapshot> { pub fn factor_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyFactorSnapshot> {
self.factor_by_date self.factor_by_date
.get(&date) .get(&date)
.map(|rows| rows.iter().map(Arc::as_ref).collect()) .map(|rows| rows.iter().collect())
.unwrap_or_default() .unwrap_or_default()
} }
pub fn factor_snapshot_rows_on(&self, date: NaiveDate) -> &[Arc<DailyFactorSnapshot>] { pub fn factor_snapshot_rows_on(&self, date: NaiveDate) -> &[DailyFactorSnapshot] {
self.factor_by_date self.factor_by_date
.get(&date) .get(&date)
.map(Vec::as_slice) .map(Vec::as_slice)
@@ -2365,14 +2359,14 @@ impl DataSet {
pub fn market_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyMarketSnapshot> { pub fn market_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyMarketSnapshot> {
self.market_by_date self.market_by_date
.get(&date) .get(&date)
.map(|rows| rows.iter().map(Arc::as_ref).collect()) .map(|rows| rows.iter().collect())
.unwrap_or_default() .unwrap_or_default()
} }
pub fn candidate_snapshots_on(&self, date: NaiveDate) -> Vec<&CandidateEligibility> { pub fn candidate_snapshots_on(&self, date: NaiveDate) -> Vec<&CandidateEligibility> {
self.candidate_by_date self.candidate_by_date
.get(&date) .get(&date)
.map(|rows| rows.iter().map(Arc::as_ref).collect()) .map(|rows| rows.iter().collect())
.unwrap_or_default() .unwrap_or_default()
} }
@@ -2384,20 +2378,12 @@ impl DataSet {
Ok(DailySnapshotBundle { Ok(DailySnapshotBundle {
date, date,
benchmark, benchmark,
market: self market: self.market_by_date.get(&date).cloned().unwrap_or_default(),
.market_by_date factors: self.factor_by_date.get(&date).cloned().unwrap_or_default(),
.get(&date)
.map(|rows| rows.iter().map(|row| row.as_ref().clone()).collect())
.unwrap_or_default(),
factors: self
.factor_by_date
.get(&date)
.map(|rows| rows.iter().map(|row| row.as_ref().clone()).collect())
.unwrap_or_default(),
candidates: self candidates: self
.candidate_by_date .candidate_by_date
.get(&date) .get(&date)
.map(|rows| rows.iter().map(|row| row.as_ref().clone()).collect()) .cloned()
.unwrap_or_default(), .unwrap_or_default(),
corporate_actions: self corporate_actions: self
.corporate_actions_by_date .corporate_actions_by_date
@@ -3276,34 +3262,20 @@ where
grouped grouped
} }
fn group_arc_by_date<T, F>(rows: &[Arc<T>], mut date_of: F) -> BTreeMap<NaiveDate, Vec<Arc<T>>> fn sort_groups_by_symbol<T, F>(groups: &mut BTreeMap<NaiveDate, Vec<T>>, symbol_of: F)
where
F: FnMut(&T) -> NaiveDate,
{
let mut grouped = BTreeMap::<NaiveDate, Vec<Arc<T>>>::new();
for row in rows {
grouped
.entry(date_of(row.as_ref()))
.or_default()
.push(Arc::clone(row));
}
grouped
}
fn sort_arc_groups_by_symbol<T, F>(groups: &mut BTreeMap<NaiveDate, Vec<Arc<T>>>, symbol_of: F)
where where
F: Fn(&T) -> &str + Copy, F: Fn(&T) -> &str + Copy,
{ {
for rows in groups.values_mut() { for rows in groups.values_mut() {
rows.sort_by(|left, right| symbol_of(left.as_ref()).cmp(symbol_of(right.as_ref()))); rows.sort_by(|left, right| symbol_of(left).cmp(symbol_of(right)));
} }
} }
fn build_symbol_id_index( fn build_symbol_id_index(
instruments: &HashMap<String, Instrument>, instruments: &HashMap<String, Instrument>,
market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>, market_by_date: &BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>,
factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>, factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
candidate_by_date: &BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>, candidate_by_date: &BTreeMap<NaiveDate, Vec<CandidateEligibility>>,
) -> AHashMap<String, u32> { ) -> AHashMap<String, u32> {
let mut symbols = instruments.keys().cloned().collect::<HashSet<_>>(); let mut symbols = instruments.keys().cloned().collect::<HashSet<_>>();
for rows in market_by_date.values() { for rows in market_by_date.values() {
@@ -3342,7 +3314,7 @@ fn build_symbol_id_index(
} }
fn build_group_symbol_ids<T, F>( fn build_group_symbol_ids<T, F>(
groups: &BTreeMap<NaiveDate, Vec<Arc<T>>>, groups: &BTreeMap<NaiveDate, Vec<T>>,
symbol_id_by_code: &AHashMap<String, u32>, symbol_id_by_code: &AHashMap<String, u32>,
symbol_of: F, symbol_of: F,
) -> BTreeMap<NaiveDate, Vec<u32>> ) -> BTreeMap<NaiveDate, Vec<u32>>
@@ -3356,7 +3328,7 @@ where
.iter() .iter()
.map(|row| { .map(|row| {
*symbol_id_by_code *symbol_id_by_code
.get(symbol_of(row.as_ref())) .get(symbol_of(row))
.expect("snapshot symbol missing from FIDC symbol index") .expect("snapshot symbol missing from FIDC symbol index")
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
@@ -3366,11 +3338,7 @@ where
.collect() .collect()
} }
fn find_arc_by_symbol_id<'a, T>( fn find_by_symbol_id<'a, T>(rows: &'a [T], symbol_ids: &[u32], symbol_id: u32) -> Option<&'a T> {
rows: &'a [Arc<T>],
symbol_ids: &[u32],
symbol_id: u32,
) -> Option<&'a T> {
if rows.len() != symbol_ids.len() { if rows.len() != symbol_ids.len() {
return None; return None;
} }
@@ -3378,16 +3346,15 @@ fn find_arc_by_symbol_id<'a, T>(
.binary_search(&symbol_id) .binary_search(&symbol_id)
.ok() .ok()
.and_then(|index| rows.get(index)) .and_then(|index| rows.get(index))
.map(Arc::as_ref)
} }
fn find_arc_by_symbol<'a, T, F>(rows: &'a [Arc<T>], symbol: &str, symbol_of: F) -> Option<&'a T> fn find_by_symbol<'a, T, F>(rows: &'a [T], symbol: &str, symbol_of: F) -> Option<&'a T>
where where
F: Fn(&T) -> &str, F: Fn(&T) -> &str,
{ {
rows.binary_search_by(|row| symbol_of(row.as_ref()).cmp(symbol)) rows.binary_search_by(|row| symbol_of(row).cmp(symbol))
.ok() .ok()
.map(|index| rows[index].as_ref()) .map(|index| &rows[index])
} }
fn collect_benchmark_code(benchmarks: &[BenchmarkSnapshot]) -> Result<String, DataSetError> { fn collect_benchmark_code(benchmarks: &[BenchmarkSnapshot]) -> Result<String, DataSetError> {
@@ -3512,8 +3479,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<DailyFactorSnapshot>>,
market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>, market_by_date: &BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>,
) -> 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();
@@ -3527,8 +3494,8 @@ fn build_eligible_universe(
fn build_fundamental_universe_for_date( fn build_fundamental_universe_for_date(
date: NaiveDate, date: NaiveDate,
factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>, factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>, market_by_date: &BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>,
) -> Vec<EligibleUniverseSnapshot> { ) -> Vec<EligibleUniverseSnapshot> {
let mut rows = Vec::new(); let mut rows = Vec::new();
let Some(factors) = factor_by_date.get(&date) else { let Some(factors) = factor_by_date.get(&date) else {
@@ -3537,7 +3504,7 @@ fn build_fundamental_universe_for_date(
for factor in factors { for factor in factors {
if market_by_date if market_by_date
.get(&date) .get(&date)
.and_then(|rows| find_arc_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str())) .and_then(|rows| find_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str()))
.is_none() .is_none()
{ {
continue; continue;
@@ -3563,9 +3530,9 @@ fn build_fundamental_universe_for_date(
fn build_eligible_universe_for_date( fn build_eligible_universe_for_date(
date: NaiveDate, date: NaiveDate,
factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>, factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
candidate_by_date: &BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>, candidate_by_date: &BTreeMap<NaiveDate, Vec<CandidateEligibility>>,
market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>, market_by_date: &BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>,
instruments: &HashMap<String, Instrument>, instruments: &HashMap<String, Instrument>,
risk_config: &FidcRiskControlConfig, risk_config: &FidcRiskControlConfig,
) -> Vec<EligibleUniverseSnapshot> { ) -> Vec<EligibleUniverseSnapshot> {
@@ -3586,9 +3553,9 @@ fn build_eligible_universe_for_date(
fn build_eligible_universe_for_date_from_factors( fn build_eligible_universe_for_date_from_factors(
date: NaiveDate, date: NaiveDate,
factors: &[Arc<DailyFactorSnapshot>], factors: &[DailyFactorSnapshot],
candidate_by_date: &BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>, candidate_by_date: &BTreeMap<NaiveDate, Vec<CandidateEligibility>>,
market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>, market_by_date: &BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>,
instruments: &HashMap<String, Instrument>, instruments: &HashMap<String, Instrument>,
risk_config: &FidcRiskControlConfig, risk_config: &FidcRiskControlConfig,
) -> Vec<EligibleUniverseSnapshot> { ) -> Vec<EligibleUniverseSnapshot> {
@@ -3600,7 +3567,7 @@ fn build_eligible_universe_for_date_from_factors(
let synthetic_candidate; let synthetic_candidate;
let candidate = if let Some(candidate) = candidate_by_date let candidate = if let Some(candidate) = candidate_by_date
.get(&date) .get(&date)
.and_then(|rows| find_arc_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str())) .and_then(|rows| find_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str()))
{ {
candidate candidate
} else { } else {
@@ -3609,7 +3576,7 @@ fn build_eligible_universe_for_date_from_factors(
}; };
let Some(market) = market_by_date let Some(market) = market_by_date
.get(&date) .get(&date)
.and_then(|rows| find_arc_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str())) .and_then(|rows| find_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str()))
else { else {
continue; continue;
}; };