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