用快速哈希优化回测内部索引

2026-08-24 12:09:09 +08:00
parent 1d7ac19886
commit c52478708f
5 changed files with 40 additions and 26 deletions
@@ -2,6 +2,7 @@ use std::borrow::Cow;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::sync::{Arc, OnceLock};

+use ahash::AHashMap;
 use chrono::{NaiveDate, NaiveDateTime};
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
@@ -1136,12 +1137,12 @@ pub struct DataSet {
    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: Arc<HashMap<String, Arc<SymbolPriceSeries>>>,
-    adjusted_close_series_by_symbol: Arc<HashMap<String, Arc<AdjustedCloseSeries>>>,
+    market_series_by_symbol: Arc<AHashMap<String, Arc<SymbolPriceSeries>>>,
+    adjusted_close_series_by_symbol: Arc<AHashMap<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>>,
+    symbol_id_by_code: Arc<AHashMap<String, u32>>,
    eligible_universe_by_date: Arc<OnceLock<BTreeMap<NaiveDate, Vec<EligibleUniverseSnapshot>>>>,
    benchmark_code: String,
    futures_params_by_symbol: HashMap<String, Vec<FuturesTradingParameter>>,
@@ -1305,27 +1306,32 @@ impl DataSet {

        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 mut market_rows_by_symbol = HashMap::<String, Vec<&DailyMarketSnapshot>>::new();
+        let mut market_rows_by_symbol = AHashMap::<String, Vec<&DailyMarketSnapshot>>::new();
        for row in &market {
            market_rows_by_symbol
                .entry(row.symbol.clone())
                .or_default()
                .push(row.as_ref());
        }
+        let market_rows_by_symbol = market_rows_by_symbol.into_iter().collect::<Vec<_>>();
        let market_series_by_symbol = market_rows_by_symbol
            .into_par_iter()
            .map(|(symbol, rows)| {
                let series = Arc::new(SymbolPriceSeries::new(symbol.clone(), rows));
                (symbol, series)
            })
-            .collect::<HashMap<_, _>>();
+            .collect::<Vec<_>>()
+            .into_iter()
+            .collect::<AHashMap<_, _>>();
        let adjusted_close_series_by_symbol = market_series_by_symbol
            .par_iter()
            .filter_map(|(symbol, market)| {
                AdjustedCloseSeries::new(market, &factor_by_date)
                    .map(|series| (symbol.clone(), Arc::new(series)))
            })
-            .collect::<HashMap<_, _>>();
+            .collect::<Vec<_>>()
+            .into_iter()
+            .collect::<AHashMap<_, _>>();
        let factor_texts = factor_texts
            .into_iter()
            .filter_map(|mut item| {
@@ -3298,7 +3304,7 @@ fn build_symbol_id_index(
    market_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
    factor_by_date: &BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
    candidate_by_date: &BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>,
-) -> HashMap<String, u32> {
+) -> AHashMap<String, u32> {
    let mut symbols = instruments.keys().cloned().collect::<HashSet<_>>();
    for rows in market_by_date.values() {
        for row in rows {
@@ -3337,7 +3343,7 @@ fn build_symbol_id_index(

 fn build_group_symbol_ids<T, F>(
    groups: &BTreeMap<NaiveDate, Vec<Arc<T>>>,
-    symbol_id_by_code: &HashMap<String, u32>,
+    symbol_id_by_code: &AHashMap<String, u32>,
    symbol_of: F,
 ) -> BTreeMap<NaiveDate, Vec<u32>>
 where
@@ -1,7 +1,8 @@
 use std::cell::RefCell;
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeMap, BTreeSet};
 use std::sync::Arc;

+use ahash::{AHashMap, AHashSet};
 use chrono::{Datelike, Duration, NaiveDate, NaiveDateTime, NaiveTime};
 use rhai::{AST, Dynamic, Engine, Map, Scope};

@@ -768,24 +769,19 @@ fn framework_stock_rolling_factor_requirement(key: &str) -> Option<(&'static str

 struct SelectiveExpressionScope<'a> {
    inner: Scope<'static>,
-    expression_identifiers: &'a BTreeSet<String>,
-    prelude_identifiers: &'a BTreeSet<String>,
+    required_identifiers: &'a AHashSet<String>,
 }

 impl<'a> SelectiveExpressionScope<'a> {
-    fn new(
-        expression_identifiers: &'a BTreeSet<String>,
-        prelude_identifiers: &'a BTreeSet<String>,
-    ) -> Self {
+    fn new(required_identifiers: &'a AHashSet<String>) -> Self {
        Self {
            inner: Scope::new(),
-            expression_identifiers,
-            prelude_identifiers,
+            required_identifiers,
        }
    }

    fn requires(&self, name: &str) -> bool {
-        self.expression_identifiers.contains(name) || self.prelude_identifiers.contains(name)
+        self.required_identifiers.contains(name)
    }

    fn push<T: Into<Dynamic>>(&mut self, name: &str, value: T) -> &mut Self {
@@ -811,6 +807,7 @@ impl<'a> SelectiveExpressionScope<'a> {

 struct ExpressionEvalPlan {
    identifiers: BTreeSet<String>,
+    scope_identifiers: AHashSet<String>,
    runtime_template: Result<RuntimeExpressionTemplate, String>,
    prelude_source: String,
    prelude_identifiers: BTreeSet<String>,
@@ -910,11 +907,11 @@ pub struct PlatformExprStrategy {
    /// Value 是 Rhai 编译产物。命中后 eval 走 eval_ast_with_scope，避免重复
    /// parsing。一次回测里同一表达式（stock_filter / stop_loss / rank_expr 等）
    /// 会被反复执行，重复解析的常数级开销在大规模回测里不可忽略。
-    compiled_cache: RefCell<HashMap<String, AST>>,
+    compiled_cache: RefCell<AHashMap<String, AST>>,
    /// 命中计数与未命中计数，便于在 unit test 中验证缓存生效；非生产指标。
    cache_hits: RefCell<u64>,
    cache_misses: RefCell<u64>,
-    expression_plan_cache: RefCell<HashMap<String, Arc<ExpressionEvalPlan>>>,
+    expression_plan_cache: RefCell<AHashMap<String, Arc<ExpressionEvalPlan>>>,
    prelude_dependency_plan: PreludeDependencyPlan,
    prelude_identifier_candidates: BTreeSet<String>,
    prelude_declared_identifiers: BTreeSet<String>,
@@ -926,7 +923,7 @@ pub struct PlatformExprStrategy {
    stock_text_factors_required: bool,
    stock_state_cache_date: RefCell<Option<NaiveDate>>,
    stock_state_cache: RefCell<
-        HashMap<(NaiveDate, NaiveDate, u32, Option<NaiveTime>, bool), StockExpressionState>,
+        AHashMap<(NaiveDate, NaiveDate, u32, Option<NaiveTime>, bool), StockExpressionState>,
    >,
 }

@@ -1222,10 +1219,10 @@ impl PlatformExprStrategy {
            position_entry_dates: BTreeMap::new(),
            position_holding_days: BTreeMap::new(),
            position_holding_days_last_counted: BTreeMap::new(),
-            compiled_cache: RefCell::new(HashMap::new()),
+            compiled_cache: RefCell::new(AHashMap::new()),
            cache_hits: RefCell::new(0),
            cache_misses: RefCell::new(0),
-            expression_plan_cache: RefCell::new(HashMap::new()),
+            expression_plan_cache: RefCell::new(AHashMap::new()),
            prelude_dependency_plan,
            prelude_identifier_candidates,
            prelude_declared_identifiers,
@@ -1236,7 +1233,7 @@ impl PlatformExprStrategy {
            stock_extra_factor_identifiers,
            stock_text_factors_required,
            stock_state_cache_date: RefCell::new(None),
-            stock_state_cache: RefCell::new(HashMap::new()),
+            stock_state_cache: RefCell::new(AHashMap::new()),
        }
    }

@@ -3973,13 +3970,14 @@ impl PlatformExprStrategy {
        day: &DayExpressionState,
        stock: Option<&StockExpressionState>,
        position: Option<&PositionExpressionState>,
+        scope_identifiers: &AHashSet<String>,
        identifiers: &BTreeSet<String>,
        prelude_identifiers: &BTreeSet<String>,
        include_day_factors: bool,
        include_factors_map: bool,
        include_process_event_counts: bool,
    ) -> Scope<'static> {
-        let mut scope = SelectiveExpressionScope::new(identifiers, prelude_identifiers);
+        let mut scope = SelectiveExpressionScope::new(scope_identifiers);
        let trade_date = day.date.format("%Y-%m-%d").to_string();
        let decision_date = ctx.decision_date.format("%Y-%m-%d").to_string();
        let execution_date = ctx.execution_date.format("%Y-%m-%d").to_string();
@@ -4686,7 +4684,8 @@ impl PlatformExprStrategy {
            day,
            stock,
            position,
-            &normalized_identifiers,
+            &expression_plan.scope_identifiers,
+            normalized_identifiers,
            prelude_identifiers,
            include_day_factors,
            include_factors_map,
@@ -4756,10 +4755,16 @@ impl PlatformExprStrategy {
            .prelude_dependency_plan
            .source_for_expression(&identifiers);
        let prelude_identifiers = Self::extract_identifier_candidates(&prelude_source);
+        let scope_identifiers = identifiers
+            .iter()
+            .chain(&prelude_identifiers)
+            .cloned()
+            .collect::<AHashSet<_>>();
        let prelude_runtime_template = (!prelude_source.trim().is_empty())
            .then(|| Self::compile_runtime_helper_template(&prelude_source));
        let plan = Arc::new(ExpressionEvalPlan {
            identifiers,
+            scope_identifiers,
            runtime_template: Self::compile_runtime_helper_template(&normalized),
            prelude_source,
            prelude_identifiers,