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

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]
+126 -19
View File
@@ -615,6 +615,7 @@ struct DayExpressionState {
#[derive(Debug, Clone)]
struct StockExpressionState {
symbol: String,
symbol_id: u32,
market_cap: f64,
market_cap_bn: f64,
free_float_cap: f64,
@@ -925,7 +926,7 @@ pub struct PlatformExprStrategy {
stock_text_factors_required: bool,
stock_state_cache_date: RefCell<Option<NaiveDate>>,
stock_state_cache: RefCell<
HashMap<(NaiveDate, NaiveDate, String, Option<NaiveTime>, bool), StockExpressionState>,
HashMap<(NaiveDate, NaiveDate, u32, Option<NaiveTime>, bool), StockExpressionState>,
>,
}
@@ -3612,24 +3613,29 @@ impl PlatformExprStrategy {
&self,
ctx: &StrategyContext<'_>,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
field: &str,
lookback: usize,
) -> Option<f64> {
ctx.data
.market_decision_numeric_moving_average(date, symbol, field, lookback)
.market_decision_numeric_moving_average_by_symbol_id(
date, symbol_id, symbol, field, lookback,
)
}
fn stock_current_rolling_mean(
&self,
ctx: &StrategyContext<'_>,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
field: &str,
lookback: usize,
) -> Option<f64> {
ctx.data
.market_current_numeric_moving_average(date, symbol, field, lookback)
ctx.data.market_current_numeric_moving_average_by_symbol_id(
date, symbol_id, symbol, field, lookback,
)
}
fn stock_state_at_time(
@@ -3703,6 +3709,13 @@ impl PlatformExprStrategy {
execution_time: Option<NaiveTime>,
use_intraday_quote: bool,
) -> Result<StockExpressionState, BacktestError> {
let symbol_id = ctx.data.symbol_id(symbol).ok_or_else(|| {
BacktestError::Data(crate::data::DataSetError::MissingSnapshot {
kind: "symbol_index",
date,
symbol: symbol.to_string(),
})
})?;
{
let mut cache_date = self.stock_state_cache_date.borrow_mut();
if *cache_date != Some(date) {
@@ -3713,7 +3726,7 @@ impl PlatformExprStrategy {
let cache_key = (
date,
factor_date,
symbol.to_string(),
symbol_id,
execution_time,
use_intraday_quote,
);
@@ -3721,8 +3734,13 @@ impl PlatformExprStrategy {
return Ok(state.clone());
}
let market = ctx.data.require_market(date, symbol)?;
let feature_market = ctx.data.market(factor_date, symbol).unwrap_or(market);
let market = ctx
.data
.require_market_by_symbol_id(date, symbol_id, symbol)?;
let feature_market = ctx
.data
.market_by_symbol_id(factor_date, symbol_id)
.unwrap_or(market);
let intraday_same_day_factor = self.uses_intraday_execution_quotes()
&& factor_date == date
&& !ctx.is_lagged_execution();
@@ -3731,14 +3749,18 @@ impl PlatformExprStrategy {
} else {
None
};
let factor = ctx.data.require_factor(factor_date, symbol)?;
let candidate = ctx.data.require_candidate(date, symbol)?;
let factor = ctx
.data
.require_factor_by_symbol_id(factor_date, symbol_id, symbol)?;
let candidate = ctx
.data
.require_candidate_by_symbol_id(date, symbol_id, symbol)?;
let instrument = ctx.data.instrument(symbol);
let rolling = |field: &'static str, lookback: usize| -> f64 {
if !self.stock_rolling_requirements.requires(field, lookback) {
return f64::NAN;
}
self.stock_decision_rolling_mean(ctx, date, symbol, field, lookback)
self.stock_decision_rolling_mean(ctx, date, symbol_id, symbol, field, lookback)
.unwrap_or(f64::NAN)
};
let stock_ma_short = rolling("close", self.config.stock_short_ma_days);
@@ -3864,6 +3886,7 @@ impl PlatformExprStrategy {
let state = StockExpressionState {
symbol: symbol.to_string(),
symbol_id,
market_cap,
market_cap_bn,
free_float_cap,
@@ -3970,6 +3993,7 @@ impl PlatformExprStrategy {
scope.push("benchmark_open", day.benchmark_open);
scope.push("benchmark_close", day.benchmark_close);
scope.push("benchmark_signal_close", day.benchmark_signal_close);
if self.market_scope_maps_requested(identifiers, prelude_identifiers) {
for field in [
"day_open",
"open",
@@ -4005,6 +4029,7 @@ impl PlatformExprStrategy {
identifiers,
);
}
}
scope.push("signal_ma5", day.signal_ma5);
scope.push("signal_ma10", day.signal_ma10);
scope.push("signal_ma20", day.signal_ma20);
@@ -5345,7 +5370,14 @@ impl PlatformExprStrategy {
))
})?;
let value = self
.stock_decision_rolling_mean(ctx, day.date, &stock.symbol, field, lookback)
.stock_decision_rolling_mean(
ctx,
day.date,
stock.symbol_id,
&stock.symbol,
field,
lookback,
)
.ok_or_else(|| {
BacktestError::Execution(format!(
"missing framework rolling factor {key} for {} on {}",
@@ -6002,7 +6034,14 @@ impl PlatformExprStrategy {
"rolling_mean(\"{other}\", {lookback}) requires stock context"
))
})?;
self.stock_decision_rolling_mean(ctx, day.date, &stock.symbol, other, lookback)
self.stock_decision_rolling_mean(
ctx,
day.date,
stock.symbol_id,
&stock.symbol,
other,
lookback,
)
}
};
value.ok_or_else(|| {
@@ -6044,7 +6083,14 @@ impl PlatformExprStrategy {
"rolling_mean_current(\"{other}\", {lookback}) requires stock context"
))
})?;
self.stock_current_rolling_mean(ctx, day.date, &stock.symbol, other, lookback)
self.stock_current_rolling_mean(
ctx,
day.date,
stock.symbol_id,
&stock.symbol,
other,
lookback,
)
}
};
value.ok_or_else(|| {
@@ -7060,8 +7106,44 @@ impl PlatformExprStrategy {
value.is_finite().then_some(value)
}
fn scope_identifier_requested(&self, identifiers: &BTreeSet<String>, name: &str) -> bool {
identifiers.contains(name) || self.prelude_identifier_candidates.contains(name)
fn is_market_scope_map_identifier(name: &str) -> bool {
let Some(field) = name
.strip_prefix("decision_")
.or_else(|| name.strip_prefix("execution_"))
else {
return false;
};
matches!(
field,
"day_open"
| "open"
| "high"
| "low"
| "close"
| "last_price"
| "prev_close"
| "upper_limit"
| "lower_limit"
| "volume"
| "minute_volume"
| "bid1"
| "ask1"
| "bid1_volume"
| "ask1_volume"
| "price_tick"
)
}
fn market_scope_maps_requested(
&self,
identifiers: &BTreeSet<String>,
prelude_identifiers: &BTreeSet<String>,
) -> bool {
identifiers
.iter()
.chain(prelude_identifiers)
.chain(&self.prelude_identifier_candidates)
.any(|name| Self::is_market_scope_map_identifier(name))
}
fn push_market_scope_map(
@@ -7073,17 +7155,24 @@ impl PlatformExprStrategy {
field: &str,
identifiers: &BTreeSet<String>,
) {
let name = format!("{prefix}_{field}");
if !self.scope_identifier_requested(identifiers, &name) {
let requested_name = identifiers
.iter()
.chain(&self.prelude_identifier_candidates)
.find(|name| {
name.strip_prefix(prefix)
.and_then(|suffix| suffix.strip_prefix('_'))
== Some(field)
});
let Some(name) = requested_name else {
return;
}
};
let mut map = Map::new();
for snapshot in ctx.data.market_snapshots_on(date) {
if let Some(value) = Self::market_scope_value(snapshot, field) {
map.insert(snapshot.symbol.clone().into(), Dynamic::from(value));
}
}
scope.push_dynamic(name, Dynamic::from(map));
scope.push_dynamic(name.clone(), Dynamic::from(map));
}
fn action_stock_state(
@@ -31453,6 +31542,24 @@ fn passes_threshold(value) { value > stock_threshold }
);
}
#[test]
fn expression_plan_skips_unreferenced_market_scope_maps() {
let strategy = PlatformExprStrategy::new(PlatformExprStrategyConfig::microcap_rotation());
let stock_plan = strategy.expression_eval_plan("close > 0 && !is_st");
assert!(
!strategy.market_scope_maps_requested(
&stock_plan.identifiers,
&stock_plan.prelude_identifiers,
)
);
let market_map_plan = strategy.expression_eval_plan("decision_close[stock] > 0");
assert!(strategy.market_scope_maps_requested(
&market_map_plan.identifiers,
&market_map_plan.prelude_identifiers,
));
}
#[test]
fn ast_cache_reuses_rolling_helper_scripts_across_dates() {
let dates = [d(2025, 2, 3), d(2025, 2, 4)];