按证券索引优化表达式数据访问
This commit is contained in:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user