perf(data): parallelize bounded daily symbol indices
This commit is contained in:
+100
-10
@@ -3,7 +3,7 @@ use std::cmp::Reverse;
|
||||
use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use ahash::AHashMap;
|
||||
use ahash::{AHashMap, AHashSet};
|
||||
use chrono::{NaiveDate, NaiveDateTime};
|
||||
use compact_str::CompactString;
|
||||
use rayon::prelude::*;
|
||||
@@ -4531,7 +4531,7 @@ fn build_symbol_id_index(
|
||||
factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
|
||||
candidate_by_date: &BTreeMap<NaiveDate, Vec<CandidateEligibility>>,
|
||||
) -> AHashMap<String, u32> {
|
||||
let mut symbols = instruments.keys().cloned().collect::<HashSet<_>>();
|
||||
let mut symbols = instruments.keys().cloned().collect::<AHashSet<_>>();
|
||||
for rows in market_by_date.values() {
|
||||
for row in rows {
|
||||
if !symbols.contains(row.symbol.as_str()) {
|
||||
@@ -4573,10 +4573,11 @@ fn build_group_symbol_ids<T, F>(
|
||||
symbol_of: F,
|
||||
) -> BTreeMap<NaiveDate, Vec<u32>>
|
||||
where
|
||||
F: Fn(&T) -> &str + Copy,
|
||||
T: Sync,
|
||||
F: Fn(&T) -> &str + Copy + Send + Sync,
|
||||
{
|
||||
groups
|
||||
.iter()
|
||||
.par_iter()
|
||||
.map(|(date, rows)| {
|
||||
let symbol_ids = rows
|
||||
.iter()
|
||||
@@ -4589,6 +4590,8 @@ where
|
||||
debug_assert!(symbol_ids.windows(2).all(|window| window[0] < window[1]));
|
||||
(*date, symbol_ids)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
@@ -4668,7 +4671,7 @@ fn build_factor_market_cap_order(
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn build_dense_row_positions<T>(
|
||||
fn build_dense_row_positions<T: Sync>(
|
||||
groups: &BTreeMap<NaiveDate, Vec<T>>,
|
||||
symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
|
||||
symbol_count: usize,
|
||||
@@ -4679,8 +4682,11 @@ fn build_dense_row_positions<T>(
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut positions_by_date = BTreeMap::new();
|
||||
for (date, rows) in groups {
|
||||
// Each task owns one bounded day index. No partial index is published if
|
||||
// any day has a missing, duplicate, or misaligned symbol identifier.
|
||||
groups
|
||||
.par_iter()
|
||||
.map(|(date, rows)| {
|
||||
let symbol_ids = symbol_ids_by_date.get(date)?;
|
||||
if rows.len() != symbol_ids.len() {
|
||||
return None;
|
||||
@@ -4693,9 +4699,10 @@ fn build_dense_row_positions<T>(
|
||||
}
|
||||
*position = u32::try_from(row_index).ok()?;
|
||||
}
|
||||
positions_by_date.insert(*date, positions);
|
||||
}
|
||||
Some(positions_by_date)
|
||||
Some((*date, positions))
|
||||
})
|
||||
.collect::<Option<Vec<_>>>()
|
||||
.map(|days| days.into_iter().collect())
|
||||
}
|
||||
|
||||
fn build_calendar_series_end_positions(
|
||||
@@ -5574,6 +5581,89 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parallel_daily_symbol_indices_match_scalar_for_sparse_and_empty_days() {
|
||||
let symbols = ["000001.SZ", "159915.SZ", "600000.SH", "932000.CSI", "custom-long-instrument"];
|
||||
let index = symbols.iter().enumerate()
|
||||
.map(|(id, symbol)| (symbol.to_string(), id as u32))
|
||||
.collect::<AHashMap<_, _>>();
|
||||
let groups = (1..29).map(|day| {
|
||||
let date = NaiveDate::from_ymd_opt(2025, 1, day).unwrap();
|
||||
let rows = symbols.iter().enumerate()
|
||||
.filter(|(id, _)| day % 7 != 0 && (*id + day as usize) % 3 != 0)
|
||||
.map(|(_, symbol)| symbol.to_string()).collect::<Vec<_>>();
|
||||
(date, rows)
|
||||
}).collect::<BTreeMap<_, _>>();
|
||||
let expected_ids = groups.iter().map(|(date, rows)| {
|
||||
(*date, rows.iter().map(|symbol| index[symbol]).collect::<Vec<_>>())
|
||||
}).collect::<BTreeMap<_, _>>();
|
||||
let expected_positions = expected_ids.iter().map(|(date, ids)| {
|
||||
let mut positions = vec![super::MISSING_ROW_POSITION; symbols.len()];
|
||||
for (row, id) in ids.iter().enumerate() { positions[*id as usize] = row as u32; }
|
||||
(*date, positions)
|
||||
}).collect::<BTreeMap<_, _>>();
|
||||
for threads in [1, 2, 8] {
|
||||
rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap().install(|| {
|
||||
let ids = super::build_group_symbol_ids(&groups, &index, String::as_str);
|
||||
assert_eq!(ids, expected_ids);
|
||||
assert_eq!(super::build_dense_row_positions(&groups, &ids, symbols.len()), Some(expected_positions.clone()));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parallel_dense_index_rejects_invalid_days_without_publishing_partial_index() {
|
||||
let day1 = NaiveDate::from_ymd_opt(2025, 1, 2).unwrap();
|
||||
let day2 = NaiveDate::from_ymd_opt(2025, 1, 3).unwrap();
|
||||
let groups = BTreeMap::from([(day1, vec![0, 1]), (day2, vec![0, 1])]);
|
||||
let valid = BTreeMap::from([(day1, vec![0, 2]), (day2, vec![1, 2])]);
|
||||
for threads in [1, 2, 8] {
|
||||
rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap().install(|| {
|
||||
for invalid in [vec![], vec![1], vec![1, 1], vec![1, 3], vec![1, u32::MAX]] {
|
||||
let mut ids = valid.clone();
|
||||
ids.insert(day2, invalid);
|
||||
assert!(super::build_dense_row_positions(&groups, &ids, 3).is_none());
|
||||
}
|
||||
let mut missing = valid.clone();
|
||||
missing.remove(&day2);
|
||||
assert!(super::build_dense_row_positions(&groups, &missing, 3).is_none());
|
||||
assert!(super::build_dense_row_positions(&groups, &valid, usize::MAX).is_none());
|
||||
assert!(super::build_dense_row_positions(&groups, &valid, super::MAX_DENSE_ROW_INDEX_BYTES).is_none());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn symbol_id_union_preserves_lexical_order_and_all_component_sources() {
|
||||
let date = NaiveDate::from_ymd_opt(2025, 1, 2).unwrap();
|
||||
let instrument = Instrument {
|
||||
symbol: "932000.CSI".into(), name: "index".into(), board: "CSI".into(),
|
||||
round_lot: 100, listed_at: None, delisted_at: None, status: "active".into(),
|
||||
};
|
||||
let mut market = market_row("2025-01-02", -0.0, 0);
|
||||
market.symbol = "custom-long-instrument".into();
|
||||
let factor = DailyFactorSnapshot {
|
||||
date, symbol: "159915.SZ".into(), market_cap_bn: 0.0, free_float_cap_bn: 0.0,
|
||||
pe_ttm: 0.0, turnover_ratio: None, effective_turnover_ratio: None,
|
||||
adjustment_factor_backward1: None, extra_factors: NumericFactorMap::new(),
|
||||
};
|
||||
let candidate = CandidateEligibility {
|
||||
date, symbol: "000001.SZ".into(), is_st: true, is_star_st: true,
|
||||
is_new_listing: false, is_paused: true, allow_buy: false, allow_sell: false,
|
||||
is_kcb: false, is_one_yuan: false, risk_level_code: Some("test".into()),
|
||||
};
|
||||
let ids = super::build_symbol_id_index(
|
||||
&HashMap::from([(instrument.symbol.clone(), instrument)]),
|
||||
&BTreeMap::from([(date, vec![market.clone(), market])]),
|
||||
&BTreeMap::from([(date, vec![factor])]),
|
||||
&BTreeMap::from([(date, vec![candidate])]),
|
||||
);
|
||||
assert_eq!(ids, AHashMap::from_iter([
|
||||
("000001.SZ".to_string(), 0), ("159915.SZ".to_string(), 1),
|
||||
("932000.CSI".to_string(), 2), ("custom-long-instrument".to_string(), 3),
|
||||
]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_symbol_id_snapshot_lookups_preserve_alignment_for_sparse_rows() {
|
||||
let date = NaiveDate::parse_from_str("2025-01-02", "%Y-%m-%d").unwrap();
|
||||
|
||||
Reference in New Issue
Block a user