按证券索引优化表达式数据访问

This commit is contained in:
boris
2026-08-24 11:21:39 +08:00
parent a79077af17
commit 555f2ab9bd
2 changed files with 350 additions and 56 deletions
+190 -3
View File
@@ -1129,6 +1129,8 @@ pub struct DataSet {
benchmark_by_date: BTreeMap<NaiveDate, BenchmarkSnapshot>,
market_series_by_symbol: Arc<HashMap<String, Arc<SymbolPriceSeries>>>,
adjusted_close_series_by_symbol: Arc<HashMap<String, Arc<AdjustedCloseSeries>>>,
market_series_by_symbol_id: Arc<Vec<Option<Arc<SymbolPriceSeries>>>>,
adjusted_close_series_by_symbol_id: Arc<Vec<Option<Arc<AdjustedCloseSeries>>>>,
benchmark_series_cache: BenchmarkPriceSeries,
symbol_id_by_code: Arc<HashMap<String, u32>>,
eligible_universe_by_date: Arc<OnceLock<BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>>>>,
@@ -1352,6 +1354,18 @@ impl DataSet {
build_group_symbol_ids(&candidate_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
});
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() {
market_series_by_symbol_id[symbol_id as usize] = Some(Arc::clone(series));
}
}
let mut adjusted_close_series_by_symbol_id = vec![None; symbol_id_by_code.len()];
for (symbol, series) in &adjusted_close_series_by_symbol {
if let Some(symbol_id) = symbol_id_by_code.get(symbol).copied() {
adjusted_close_series_by_symbol_id[symbol_id as usize] = Some(Arc::clone(series));
}
}
let corporate_actions_by_date = group_by_date(corporate_actions, |item| item.date);
let execution_quotes_by_date = build_execution_quote_index(execution_quotes);
let order_book_depth_index = build_order_book_depth_index(order_book_depth);
@@ -1381,6 +1395,8 @@ impl DataSet {
benchmark_by_date,
market_series_by_symbol: Arc::new(market_series_by_symbol),
adjusted_close_series_by_symbol: Arc::new(adjusted_close_series_by_symbol),
market_series_by_symbol_id: Arc::new(market_series_by_symbol_id),
adjusted_close_series_by_symbol_id: Arc::new(adjusted_close_series_by_symbol_id),
benchmark_series_cache,
symbol_id_by_code: Arc::new(symbol_id_by_code),
eligible_universe_by_date: Arc::new(OnceLock::new()),
@@ -1426,8 +1442,20 @@ impl DataSet {
self.instruments.get(symbol)
}
pub fn symbol_id(&self, symbol: &str) -> Option<u32> {
self.symbol_id_by_code.get(symbol).copied()
}
pub fn market(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
let symbol_id = *self.symbol_id_by_code.get(symbol)?;
let symbol_id = self.symbol_id(symbol)?;
self.market_by_symbol_id(date, symbol_id)
}
pub fn market_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
) -> Option<&DailyMarketSnapshot> {
find_arc_by_symbol_id(
self.market_by_date.get(&date)?,
self.market_symbol_ids_by_date.get(&date)?,
@@ -1439,14 +1467,34 @@ impl DataSet {
self.market_series_by_symbol.get(symbol).map(Arc::as_ref)
}
fn market_series_by_symbol_id(&self, symbol_id: u32) -> Option<&SymbolPriceSeries> {
self.market_series_by_symbol_id
.get(symbol_id as usize)?
.as_deref()
}
fn adjusted_close_series(&self, symbol: &str) -> Option<&AdjustedCloseSeries> {
self.adjusted_close_series_by_symbol
.get(symbol)
.map(Arc::as_ref)
}
fn adjusted_close_series_by_symbol_id(&self, symbol_id: u32) -> Option<&AdjustedCloseSeries> {
self.adjusted_close_series_by_symbol_id
.get(symbol_id as usize)?
.as_deref()
}
pub fn factor(&self, date: NaiveDate, symbol: &str) -> Option<&DailyFactorSnapshot> {
let symbol_id = *self.symbol_id_by_code.get(symbol)?;
let symbol_id = self.symbol_id(symbol)?;
self.factor_by_symbol_id(date, symbol_id)
}
pub fn factor_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
) -> Option<&DailyFactorSnapshot> {
find_arc_by_symbol_id(
self.factor_by_date.get(&date)?,
self.factor_symbol_ids_by_date.get(&date)?,
@@ -1455,7 +1503,15 @@ impl DataSet {
}
pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> {
let symbol_id = *self.symbol_id_by_code.get(symbol)?;
let symbol_id = self.symbol_id(symbol)?;
self.candidate_by_symbol_id(date, symbol_id)
}
pub fn candidate_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
) -> Option<&CandidateEligibility> {
find_arc_by_symbol_id(
self.candidate_by_date.get(&date)?,
self.candidate_symbol_ids_by_date.get(&date)?,
@@ -2509,6 +2565,35 @@ impl DataSet {
}
}
pub fn market_decision_numeric_moving_average_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
field: &str,
lookback: usize,
) -> Option<f64> {
let field = normalized_field(field);
match field.as_ref() {
"close" | "prev_close" | "stock_close" | "price" => self
.adjusted_close_series_by_symbol_id(symbol_id)
.and_then(|series| series.decision_moving_average(date, lookback)),
"volume" | "stock_volume" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.decision_volume_moving_average(date, lookback)),
"day_open" | "dayopen" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.moving_average(date, lookback, PriceField::DayOpen)),
"open" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.moving_average(date, lookback, PriceField::Open)),
"last" | "last_price" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.moving_average(date, lookback, PriceField::Last)),
other => self.factor_moving_average(date, symbol, other, lookback),
}
}
pub fn market_current_numeric_moving_average(
&self,
date: NaiveDate,
@@ -2535,6 +2620,35 @@ impl DataSet {
}
}
pub fn market_current_numeric_moving_average_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
field: &str,
lookback: usize,
) -> Option<f64> {
let field = normalized_field(field);
match field.as_ref() {
"close" | "prev_close" | "stock_close" | "price" => self
.adjusted_close_series_by_symbol_id(symbol_id)
.and_then(|series| series.current_moving_average(date, lookback)),
"volume" | "stock_volume" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.current_volume_moving_average(date, lookback)),
"day_open" | "dayopen" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.moving_average(date, lookback, PriceField::DayOpen)),
"open" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.moving_average(date, lookback, PriceField::Open)),
"last" | "last_price" => self
.market_series_by_symbol_id(symbol_id)
.and_then(|series| series.moving_average(date, lookback, PriceField::Last)),
other => self.factor_moving_average(date, symbol, other, lookback),
}
}
pub fn market_latest_back_adjusted_close(&self, date: NaiveDate, symbol: &str) -> Option<f64> {
self.adjusted_close_series(symbol)
.and_then(|series| series.latest_back_adjusted_close(date))
@@ -2736,6 +2850,20 @@ impl DataSet {
})
}
pub fn require_market_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
) -> Result<&DailyMarketSnapshot, DataSetError> {
self.market_by_symbol_id(date, symbol_id)
.ok_or_else(|| DataSetError::MissingSnapshot {
kind: "market",
date,
symbol: symbol.to_string(),
})
}
pub fn require_candidate(
&self,
date: NaiveDate,
@@ -2749,6 +2877,20 @@ impl DataSet {
})
}
pub fn require_candidate_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
) -> Result<&CandidateEligibility, DataSetError> {
self.candidate_by_symbol_id(date, symbol_id)
.ok_or_else(|| DataSetError::MissingSnapshot {
kind: "candidate",
date,
symbol: symbol.to_string(),
})
}
pub fn require_factor(
&self,
date: NaiveDate,
@@ -2761,6 +2903,20 @@ impl DataSet {
symbol: symbol.to_string(),
})
}
pub fn require_factor_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
) -> Result<&DailyFactorSnapshot, DataSetError> {
self.factor_by_symbol_id(date, symbol_id)
.ok_or_else(|| DataSetError::MissingSnapshot {
kind: "factor",
date,
symbol: symbol.to_string(),
})
}
}
fn normalized_aliases(values: &[String]) -> Vec<String> {
@@ -3703,10 +3859,31 @@ mod tests {
volume_contract_data(Some([1.0, 1.0, 1.0])),
volume_contract_data(None),
] {
let symbol_id = data.symbol_id("000001.SZ").expect("symbol id");
assert!(std::ptr::eq(
data.market_by_symbol_id(date, symbol_id)
.expect("market by id"),
data.market(date, "000001.SZ").expect("market by code"),
));
assert!(std::ptr::eq(
data.factor_by_symbol_id(date, symbol_id)
.expect("factor by id"),
data.factor(date, "000001.SZ").expect("factor by code"),
));
assert_eq!(
data.market_current_numeric_moving_average(date, "000001.SZ", "volume", 2),
Some(200.0)
);
assert_eq!(
data.market_current_numeric_moving_average_by_symbol_id(
date,
symbol_id,
"000001.SZ",
"volume",
2,
),
Some(200.0)
);
assert_eq!(
data.market_current_numeric_values(date, "000001.SZ", "volume", 2),
vec![100.0, 300.0]
@@ -3719,6 +3896,16 @@ mod tests {
data.market_decision_numeric_moving_average(date, "000001.SZ", "volume", 1),
Some(100.0)
);
assert_eq!(
data.market_decision_numeric_moving_average_by_symbol_id(
date,
symbol_id,
"000001.SZ",
"volume",
1,
),
Some(100.0)
);
assert_eq!(
data.market_decision_numeric_values(date, "000001.SZ", "volume", 1),
vec![100.0]