Compare commits

...

3 Commits

Author SHA1 Message Date
boris c83526a6a4 懒加载日线序列缓存降低回测内存 2026-06-21 03:57:07 +08:00
boris 9bd19aa042 瘦身回测数据集按日索引内存 2026-06-21 03:48:22 +08:00
boris 2f62d82420 优化回测数据集内存并修复rolling依赖识别 2026-06-21 03:31:45 +08:00
2 changed files with 194 additions and 133 deletions
+103 -133
View File
@@ -1,7 +1,7 @@
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs;
use std::path::Path;
use std::sync::Arc;
use std::sync::{Arc, OnceLock, RwLock};
use chrono::{NaiveDate, NaiveDateTime};
use serde::{Deserialize, Serialize};
@@ -984,20 +984,17 @@ pub struct DataSet {
instruments: HashMap<String, Instrument>,
calendar: TradingCalendar,
market_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
market_index: HashMap<NaiveDate, HashMap<String, Arc<DailyMarketSnapshot>>>,
factor_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
factor_index: HashMap<NaiveDate, HashMap<String, Arc<DailyFactorSnapshot>>>,
factor_text_by_date: BTreeMap<NaiveDate, Vec<FactorTextValue>>,
factor_text_index: HashMap<(NaiveDate, String, String), FactorTextValue>,
candidate_by_date: BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>,
candidate_index: HashMap<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>,
corporate_actions_by_date: BTreeMap<NaiveDate, Vec<CorporateAction>>,
execution_quotes_by_date: HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>,
order_book_depth_index: HashMap<(NaiveDate, String), Vec<IntradayOrderBookDepthLevel>>,
benchmark_by_date: BTreeMap<NaiveDate, BenchmarkSnapshot>,
market_series_by_symbol: HashMap<String, SymbolPriceSeries>,
market_series_by_symbol: Arc<RwLock<HashMap<String, Arc<SymbolPriceSeries>>>>,
benchmark_series_cache: BenchmarkPriceSeries,
eligible_universe_by_date: BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>>,
eligible_universe_by_date: Arc<OnceLock<BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>>>>,
benchmark_code: String,
futures_params_by_symbol: HashMap<String, Vec<FuturesTradingParameter>>,
}
@@ -1247,25 +1244,11 @@ impl DataSet {
.collect::<HashMap<_, _>>();
let market = market.into_iter().map(Arc::new).collect::<Vec<_>>();
let market_by_date = group_arc_by_date(&market, |item| item.date);
let mut market_index =
HashMap::<NaiveDate, HashMap<String, Arc<DailyMarketSnapshot>>>::new();
for item in market {
market_index
.entry(item.date)
.or_default()
.insert(item.symbol.clone(), item);
}
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 factor_by_date = group_arc_by_date(&factors, |item| item.date);
let mut factor_index =
HashMap::<NaiveDate, HashMap<String, Arc<DailyFactorSnapshot>>>::new();
for item in factors {
factor_index
.entry(item.date)
.or_default()
.insert(item.symbol.clone(), item);
}
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 factor_texts = factor_texts
.into_iter()
.filter_map(|mut item| {
@@ -1283,15 +1266,8 @@ impl DataSet {
.map(|item| ((item.date, item.symbol.clone(), item.field.clone()), item))
.collect::<HashMap<_, _>>();
let candidate_by_date = group_arc_by_date(&candidates, |item| item.date);
let mut candidate_index =
HashMap::<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>::new();
for item in candidates {
candidate_index
.entry(item.date)
.or_default()
.insert(item.symbol.clone(), item);
}
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 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);
@@ -1300,35 +1276,25 @@ impl DataSet {
.into_iter()
.map(|item| (item.date, item))
.collect::<BTreeMap<_, _>>();
let market_series_by_symbol = build_market_series(&market_by_date);
let benchmark_series_cache =
BenchmarkPriceSeries::new(&benchmark_by_date.values().cloned().collect::<Vec<_>>());
let eligible_universe_by_date = build_eligible_universe(
&factor_by_date,
&candidate_index,
&market_index,
&instruments,
);
let futures_params_by_symbol = build_futures_params_index(futures_params);
Ok(Self {
instruments,
calendar,
market_by_date,
market_index,
factor_by_date,
factor_index,
factor_text_by_date,
factor_text_index,
candidate_by_date,
candidate_index,
corporate_actions_by_date,
execution_quotes_by_date,
order_book_depth_index,
benchmark_by_date,
market_series_by_symbol,
market_series_by_symbol: Arc::new(RwLock::new(HashMap::new())),
benchmark_series_cache,
eligible_universe_by_date,
eligible_universe_by_date: Arc::new(OnceLock::new()),
benchmark_code,
futures_params_by_symbol,
})
@@ -1372,24 +1338,54 @@ impl DataSet {
}
pub fn market(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
self.market_index
self.market_by_date
.get(&date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref)
.and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
}
fn market_series(&self, symbol: &str) -> Option<Arc<SymbolPriceSeries>> {
if let Some(series) = self
.market_series_by_symbol
.read()
.expect("market series cache lock poisoned")
.get(symbol)
.cloned()
{
return Some(series);
}
let rows = self
.market_by_date
.values()
.filter_map(|day_rows| find_arc_by_symbol(day_rows, symbol, |row| row.symbol.as_str()))
.collect::<Vec<_>>();
if rows.is_empty() {
return None;
}
let series = Arc::new(SymbolPriceSeries::new(symbol.to_string(), rows));
let mut cache = self
.market_series_by_symbol
.write()
.expect("market series cache lock poisoned");
Some(
cache
.entry(symbol.to_string())
.or_insert_with(|| Arc::clone(&series))
.clone(),
)
}
pub fn factor(&self, date: NaiveDate, symbol: &str) -> Option<&DailyFactorSnapshot> {
self.factor_index
self.factor_by_date
.get(&date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref)
.and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
}
pub fn candidate(&self, date: NaiveDate, symbol: &str) -> Option<&CandidateEligibility> {
self.candidate_index
self.candidate_by_date
.get(&date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref)
.and_then(|rows| find_arc_by_symbol(rows, symbol, |row| row.symbol.as_str()))
}
pub fn benchmark(&self, date: NaiveDate) -> Option<&BenchmarkSnapshot> {
@@ -1621,8 +1617,7 @@ impl DataSet {
bar_count: usize,
include_now: bool,
) -> Vec<DailyMarketSnapshot> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.map(|series| series.trailing_snapshots(date, bar_count, include_now))
.unwrap_or_default()
}
@@ -2179,22 +2174,18 @@ impl DataSet {
symbol: &str,
field: PriceField,
) -> Option<f64> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.and_then(|series| series.price_on_or_before(date, field))
}
pub fn market_before(&self, date: NaiveDate, symbol: &str) -> Option<&DailyMarketSnapshot> {
let series = self.market_series_by_symbol.get(symbol)?;
let series = self.market_series(symbol)?;
let end = series.previous_completed_end_index(date)?;
if end == 0 {
return None;
}
let previous_date = *series.dates.get(end - 1)?;
self.market_index
.get(&previous_date)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref)
self.market(previous_date, symbol)
}
pub fn factor_snapshots_on(&self, date: NaiveDate) -> Vec<&DailyFactorSnapshot> {
@@ -2261,8 +2252,7 @@ impl DataSet {
}
pub fn market_closes_up_to(&self, date: NaiveDate, symbol: &str, lookback: usize) -> Vec<f64> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.map(|series| series.trailing_values(date, lookback, PriceField::Close))
.unwrap_or_default()
}
@@ -2275,8 +2265,7 @@ impl DataSet {
field: &str,
include_now: bool,
) -> Vec<f64> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.map(|series| series.trailing_numeric_values(date, bar_count, field, include_now))
.unwrap_or_default()
}
@@ -2317,24 +2306,12 @@ impl DataSet {
let start = days.len().saturating_sub(count);
days[start..]
.iter()
.map(|day| {
evaluator(
self.candidate_index
.get(day)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref),
self.market_index
.get(day)
.and_then(|rows| rows.get(symbol))
.map(Arc::as_ref),
)
})
.map(|day| evaluator(self.candidate(*day, symbol), self.market(*day, symbol)))
.collect()
}
pub fn market_decision_close(&self, date: NaiveDate, symbol: &str) -> Option<f64> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.and_then(|series| series.decision_price_on_or_before(date))
}
@@ -2344,8 +2321,7 @@ impl DataSet {
symbol: &str,
lookback: usize,
) -> Option<f64> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.and_then(|series| series.decision_close_moving_average(date, lookback))
}
@@ -2355,8 +2331,7 @@ impl DataSet {
symbol: &str,
lookback: usize,
) -> Option<f64> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.and_then(|series| series.decision_volume_moving_average(date, lookback))
}
@@ -2443,12 +2418,10 @@ impl DataSet {
let field = normalize_field(field);
match field.as_str() {
"close" | "prev_close" | "stock_close" | "price" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.and_then(|series| series.decision_close_moving_average(date, lookback)),
"volume" | "stock_volume" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.and_then(|series| series.decision_volume_moving_average(date, lookback)),
"day_open" | "dayopen" => {
self.market_moving_average(date, symbol, lookback, PriceField::DayOpen)
@@ -2474,8 +2447,7 @@ impl DataSet {
self.market_moving_average(date, symbol, lookback, PriceField::Close)
}
"volume" | "stock_volume" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.and_then(|series| series.current_volume_moving_average(date, lookback))
.or_else(|| self.factor_moving_average(date, symbol, "daily_volume", lookback)),
"day_open" | "dayopen" => {
@@ -2502,28 +2474,23 @@ impl DataSet {
let field = normalize_field(field);
match field.as_str() {
"close" | "prev_close" | "stock_close" | "price" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.and_then(|series| series.decision_prev_close_values(date, lookback))
.unwrap_or_default(),
"volume" | "stock_volume" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.and_then(|series| series.decision_volume_values(date, lookback))
.unwrap_or_default(),
"day_open" | "dayopen" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.map(|series| series.trailing_values(date, lookback, PriceField::DayOpen))
.unwrap_or_default(),
"open" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.map(|series| series.trailing_values(date, lookback, PriceField::Open))
.unwrap_or_default(),
"last" | "last_price" => self
.market_series_by_symbol
.get(symbol)
.market_series(symbol)
.map(|series| series.trailing_values(date, lookback, PriceField::Last))
.unwrap_or_default(),
other => self.factor_numeric_values(date, symbol, other, lookback),
@@ -2555,8 +2522,7 @@ impl DataSet {
lookback: usize,
field: PriceField,
) -> Option<f64> {
self.market_series_by_symbol
.get(symbol)
self.market_series(symbol)
.and_then(|series| series.moving_average(date, lookback, field))
}
@@ -2625,6 +2591,14 @@ impl DataSet {
pub fn eligible_universe_on(&self, date: NaiveDate) -> &[EligibleUniverseSnapshot] {
self.eligible_universe_by_date
.get_or_init(|| {
build_eligible_universe(
&self.factor_by_date,
&self.candidate_by_date,
&self.market_by_date,
&self.instruments,
)
})
.get(&date)
.map(Vec::as_slice)
.unwrap_or(&[])
@@ -3521,6 +3495,24 @@ where
grouped
}
fn sort_arc_groups_by_symbol<T, F>(groups: &mut BTreeMap<NaiveDate, Vec<Arc<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())));
}
}
fn find_arc_by_symbol<'a, T, F>(rows: &'a [Arc<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))
.ok()
.map(|index| rows[index].as_ref())
}
fn collect_benchmark_code(benchmarks: &[BenchmarkSnapshot]) -> Result<String, DataSetError> {
let mut codes = benchmarks
.iter()
@@ -3589,25 +3581,6 @@ mod optional_date_format {
}
}
fn build_market_series(
market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
) -> HashMap<String, SymbolPriceSeries> {
let mut grouped = HashMap::<String, Vec<&DailyMarketSnapshot>>::new();
for rows in market_by_date.values() {
for row in rows {
grouped.entry(row.symbol.clone()).or_default().push(row.as_ref());
}
}
grouped
.into_iter()
.map(|(symbol, rows)| {
let series = SymbolPriceSeries::new(symbol.clone(), rows);
(symbol, series)
})
.collect()
}
fn build_futures_params_index(
rows: Vec<FuturesTradingParameter>,
) -> HashMap<String, Vec<FuturesTradingParameter>> {
@@ -3667,8 +3640,8 @@ fn build_order_book_depth_index(
fn build_eligible_universe(
factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
candidate_index: &HashMap<NaiveDate, HashMap<String, Arc<CandidateEligibility>>>,
market_index: &HashMap<NaiveDate, HashMap<String, Arc<DailyMarketSnapshot>>>,
candidate_by_date: &BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>,
market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
instruments: &HashMap<String, Instrument>,
) -> BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>> {
let mut per_date = BTreeMap::<NaiveDate, Vec<EligibleUniverseSnapshot>>::new();
@@ -3679,19 +3652,16 @@ fn build_eligible_universe(
if factor.market_cap_bn <= 0.0 || !factor.market_cap_bn.is_finite() {
continue;
}
let Some(candidate) = candidate_index
.get(date)
.and_then(|rows| rows.get(&factor.symbol))
else {
let Some(candidate) = candidate_by_date.get(date).and_then(|rows| {
find_arc_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str())
}) else {
continue;
};
let Some(market) = market_index
.get(date)
.and_then(|rows| rows.get(&factor.symbol))
else {
let Some(market) = market_by_date.get(date).and_then(|rows| {
find_arc_by_symbol(rows, &factor.symbol, |row| row.symbol.as_str())
}) else {
continue;
};
let market = market.as_ref();
if ChinaAShareRiskControl::selection_rejection_reason(
*date,
candidate,
@@ -6374,6 +6374,7 @@ impl PlatformExprStrategy {
return requirements;
}
Self::require_stock_rollings_for_identifiers(&mut requirements, config, &normalized);
Self::require_stock_rollings_for_helper_calls(&mut requirements, &normalized);
}
Self::require_stock_rollings_for_fast_filter(
&mut requirements,
@@ -6612,6 +6613,96 @@ impl PlatformExprStrategy {
}
}
fn require_stock_rollings_for_helper_calls(
requirements: &mut StockRollingRequirements,
expr: &str,
) {
let compact = Self::compact_expr(expr);
Self::require_stock_rollings_for_named_helper(requirements, &compact, "rolling_mean");
Self::require_stock_rollings_for_named_helper(requirements, &compact, "sma");
Self::require_stock_rollings_for_named_helper(requirements, &compact, "ma");
Self::require_stock_rollings_for_vma_helper(requirements, &compact);
}
fn require_stock_rollings_for_named_helper(
requirements: &mut StockRollingRequirements,
compact: &str,
helper: &str,
) {
let needle = format!("{helper}(");
let mut offset = 0usize;
while let Some(relative) = compact[offset..].find(&needle) {
let args_start = offset + relative + needle.len();
let args = &compact[args_start..];
let Some((field, rest)) = Self::parse_quoted_arg(args) else {
offset = args_start;
continue;
};
let Some(rest) = rest.strip_prefix(',') else {
offset = args_start;
continue;
};
let Some(lookback) = Self::parse_usize_prefix(rest) else {
offset = args_start;
continue;
};
let normalized_field = field.trim().to_ascii_lowercase();
match normalized_field.as_str() {
"close" | "prev_close" | "stock_close" | "price" => {
requirements.require(StockRollingField::Close, lookback)
}
"volume" | "stock_volume" => {
requirements.require(StockRollingField::Volume, lookback)
}
_ => {}
}
offset = args_start;
}
}
fn require_stock_rollings_for_vma_helper(
requirements: &mut StockRollingRequirements,
compact: &str,
) {
let needle = "vma(";
let mut offset = 0usize;
while let Some(relative) = compact[offset..].find(needle) {
let args_start = offset + relative + needle.len();
if let Some(lookback) = Self::parse_usize_prefix(&compact[args_start..]) {
requirements.require(StockRollingField::Volume, lookback);
}
offset = args_start;
}
}
fn parse_quoted_arg(input: &str) -> Option<(String, &str)> {
let mut chars = input.char_indices();
let (_, quote) = chars.next()?;
if quote != '"' && quote != '\'' {
return None;
}
for (idx, ch) in chars {
if ch == quote {
let value = input[1..idx].to_string();
let rest = &input[idx + ch.len_utf8()..];
return Some((value, rest));
}
}
None
}
fn parse_usize_prefix(input: &str) -> Option<usize> {
let end = input
.char_indices()
.take_while(|(_, ch)| ch.is_ascii_digit())
.map(|(idx, ch)| idx + ch.len_utf8())
.last()?;
input[..end]
.parse::<usize>()
.ok()
.filter(|value| *value > 0)
}
fn require_stock_rollings_for_fast_filter(
requirements: &mut StockRollingRequirements,
config: &PlatformExprStrategyConfig,