zsb/fidc-backtest-engine
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

加速按日股票快照查找

Browse Source
This commit is contained in:
boris
2026-08-28 14:39:12 +08:00
parent 2de84e88eb
commit 8691076cef
1 changed files with 142 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/fidc-core/src/data.rs
+142 -40
View File
@@ -566,6 +566,11 @@ struct SymbolPriceSeries {
valid_volume_start_by_count: Vec<usize>,
}
type DenseRowPositionIndex = BTreeMap<NaiveDate, Vec<u32>>;
const MISSING_ROW_POSITION: u32 = u32::MAX;
const MAX_DENSE_ROW_INDEX_BYTES: usize = 256 * 1024 * 1024;
#[derive(Debug, Clone)]
struct AdjustedCloseSeries {
dates: Vec<NaiveDate>,
@@ -1243,12 +1248,15 @@ pub struct DataSet {
calendar: Arc<TradingCalendar>,
market_by_date: Arc<BTreeMap<NaiveDate, Vec<DailyMarketSnapshot>>>,
market_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
market_row_positions_by_date: Arc<Option<DenseRowPositionIndex>>,
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_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>>>,
candidate_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
candidate_row_positions_by_date: Arc<Option<DenseRowPositionIndex>>,
corporate_actions_by_date: Arc<BTreeMap<NaiveDate, Vec<CorporateAction>>>,
execution_quotes_by_date: Arc<HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>>,
execution_quote_dates: Arc<Vec<NaiveDate>>,
@@ -1660,6 +1668,21 @@ impl DataSet {
build_group_symbol_ids(&candidate_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
});
let market_row_positions_by_date = build_dense_row_positions(
&market_by_date,
&market_symbol_ids_by_date,
symbol_id_by_code.len(),
);
let factor_row_positions_by_date = build_dense_row_positions(
&factor_by_date,
&factor_symbol_ids_by_date,
symbol_id_by_code.len(),
);
let candidate_row_positions_by_date = build_dense_row_positions(
&candidate_by_date,
&candidate_symbol_ids_by_date,
symbol_id_by_code.len(),
);
let mut market_series_by_symbol_id = vec![None; symbol_id_by_code.len()];
for (symbol, series) in &market_series_by_symbol {
if let Some(symbol_id) = symbol_id_by_code.get(symbol).copied() {
@@ -1685,12 +1708,15 @@ impl DataSet {
calendar: Arc::new(calendar),
market_by_date: Arc::new(market_by_date),
market_symbol_ids_by_date: Arc::new(market_symbol_ids_by_date),
market_row_positions_by_date: Arc::new(market_row_positions_by_date),
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_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),
candidate_symbol_ids_by_date: Arc::new(candidate_symbol_ids_by_date),
candidate_row_positions_by_date: Arc::new(candidate_row_positions_by_date),
corporate_actions_by_date: Arc::new(corporate_actions_by_date),
execution_quotes_by_date: Arc::new(execution_quotes_by_date),
execution_quote_dates: Arc::new(execution_quote_dates),
@@ -1759,8 +1785,13 @@ impl DataSet {
date: NaiveDate,
symbol_id: u32,
) -> Option<&DailyMarketSnapshot> {
let rows = self.market_by_date.get(&date)?;
if let Some(index) = dense_row_position(&self.market_row_positions_by_date, date, symbol_id)
{
return rows.get(index);
}
find_by_symbol_id(
self.market_by_date.get(&date)?,
rows,
self.market_symbol_ids_by_date.get(&date)?,
symbol_id,
)
@@ -1798,8 +1829,13 @@ impl DataSet {
date: NaiveDate,
symbol_id: u32,
) -> Option<&DailyFactorSnapshot> {
let rows = self.factor_by_date.get(&date)?;
if let Some(index) = dense_row_position(&self.factor_row_positions_by_date, date, symbol_id)
{
return rows.get(index);
}
find_by_symbol_id(
self.factor_by_date.get(&date)?,
rows,
self.factor_symbol_ids_by_date.get(&date)?,
symbol_id,
)
@@ -1815,8 +1851,14 @@ impl DataSet {
date: NaiveDate,
symbol_id: u32,
) -> Option<&CandidateEligibility> {
let rows = self.candidate_by_date.get(&date)?;
if let Some(index) =
dense_row_position(&self.candidate_row_positions_by_date, date, symbol_id)
{
return rows.get(index);
}
find_by_symbol_id(
self.candidate_by_date.get(&date)?,
rows,
self.candidate_symbol_ids_by_date.get(&date)?,
symbol_id,
)
@@ -1832,45 +1874,50 @@ impl DataSet {
.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_index = dense_row_position(
&self.market_row_positions_by_date,
date,
symbol_id,
)
.or_else(|| {
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 factor = self
.factor_by_date
.get(&date)
.map(Vec::as_slice)
.zip(
self.factor_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,
)
});
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,
)
});
let factor = self.factor_by_date.get(&date).and_then(|rows| {
dense_row_position(&self.factor_row_positions_by_date, date, symbol_id)
.and_then(|index| rows.get(index))
.or_else(|| {
self.factor_symbol_ids_by_date
.get(&date)
.and_then(|symbol_ids| {
find_by_symbol_id_with_preferred_index(
rows,
symbol_ids,
symbol_id,
market_index,
)
})
})
});
let candidate = self.candidate_by_date.get(&date).and_then(|rows| {
dense_row_position(&self.candidate_row_positions_by_date, date, symbol_id)
.and_then(|index| rows.get(index))
.or_else(|| {
self.candidate_symbol_ids_by_date
.get(&date)
.and_then(|symbol_ids| {
find_by_symbol_id_with_preferred_index(
rows,
symbol_ids,
symbol_id,
market_index,
)
})
})
});
SymbolSnapshotRefs {
market,
@@ -3807,6 +3854,49 @@ where
.collect()
}
fn build_dense_row_positions<T>(
groups: &BTreeMap<NaiveDate, Vec<T>>,
symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
symbol_count: usize,
) -> Option<DenseRowPositionIndex> {
let entries = groups.len().checked_mul(symbol_count)?;
let bytes = entries.checked_mul(std::mem::size_of::<u32>())?;
if bytes > MAX_DENSE_ROW_INDEX_BYTES {
return None;
}
let mut positions_by_date = BTreeMap::new();
for (date, rows) in groups {
let symbol_ids = symbol_ids_by_date.get(date)?;
if rows.len() != symbol_ids.len() {
return None;
}
let mut positions = vec![MISSING_ROW_POSITION; symbol_count];
for (row_index, symbol_id) in symbol_ids.iter().copied().enumerate() {
let position = positions.get_mut(usize::try_from(symbol_id).ok()?)?;
if *position != MISSING_ROW_POSITION {
return None;
}
*position = u32::try_from(row_index).ok()?;
}
positions_by_date.insert(*date, positions);
}
Some(positions_by_date)
}
fn dense_row_position(
positions_by_date: &Option<DenseRowPositionIndex>,
date: NaiveDate,
symbol_id: u32,
) -> Option<usize> {
let position = positions_by_date
.as_ref()?
.get(&date)?
.get(usize::try_from(symbol_id).ok()?)
.copied()?;
(position != MISSING_ROW_POSITION).then_some(position as usize)
}
fn find_by_symbol_id<'a, T>(rows: &'a [T], symbol_ids: &[u32], symbol_id: u32) -> Option<&'a T> {
find_by_symbol_id_with_preferred_index(rows, symbol_ids, symbol_id, None)
}
@@ -4191,11 +4281,23 @@ mod tests {
assert!(Arc::ptr_eq(&data.instruments, &run_data.instruments));
assert!(Arc::ptr_eq(&data.market_by_date, &run_data.market_by_date));
assert!(Arc::ptr_eq(
&data.market_row_positions_by_date,
&run_data.market_row_positions_by_date
));
assert!(Arc::ptr_eq(&data.factor_by_date, &run_data.factor_by_date));
assert!(Arc::ptr_eq(
&data.factor_row_positions_by_date,
&run_data.factor_row_positions_by_date
));
assert!(Arc::ptr_eq(
&data.candidate_by_date,
&run_data.candidate_by_date
));
assert!(Arc::ptr_eq(
&data.candidate_row_positions_by_date,
&run_data.candidate_row_positions_by_date
));
assert!(Arc::ptr_eq(
&data.benchmark_by_date,
&run_data.benchmark_by_date