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::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
|
||||||
use std::sync::{Arc, OnceLock};
|
use std::sync::{Arc, OnceLock};
|
||||||
|
|
||||||
use ahash::AHashMap;
|
use ahash::{AHashMap, AHashSet};
|
||||||
use chrono::{NaiveDate, NaiveDateTime};
|
use chrono::{NaiveDate, NaiveDateTime};
|
||||||
use compact_str::CompactString;
|
use compact_str::CompactString;
|
||||||
use rayon::prelude::*;
|
use rayon::prelude::*;
|
||||||
@@ -4531,7 +4531,7 @@ fn build_symbol_id_index(
|
|||||||
factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
|
factor_by_date: &BTreeMap<NaiveDate, Vec<DailyFactorSnapshot>>,
|
||||||
candidate_by_date: &BTreeMap<NaiveDate, Vec<CandidateEligibility>>,
|
candidate_by_date: &BTreeMap<NaiveDate, Vec<CandidateEligibility>>,
|
||||||
) -> AHashMap<String, u32> {
|
) -> 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 rows in market_by_date.values() {
|
||||||
for row in rows {
|
for row in rows {
|
||||||
if !symbols.contains(row.symbol.as_str()) {
|
if !symbols.contains(row.symbol.as_str()) {
|
||||||
@@ -4573,10 +4573,11 @@ fn build_group_symbol_ids<T, F>(
|
|||||||
symbol_of: F,
|
symbol_of: F,
|
||||||
) -> BTreeMap<NaiveDate, Vec<u32>>
|
) -> BTreeMap<NaiveDate, Vec<u32>>
|
||||||
where
|
where
|
||||||
F: Fn(&T) -> &str + Copy,
|
T: Sync,
|
||||||
|
F: Fn(&T) -> &str + Copy + Send + Sync,
|
||||||
{
|
{
|
||||||
groups
|
groups
|
||||||
.iter()
|
.par_iter()
|
||||||
.map(|(date, rows)| {
|
.map(|(date, rows)| {
|
||||||
let symbol_ids = rows
|
let symbol_ids = rows
|
||||||
.iter()
|
.iter()
|
||||||
@@ -4589,6 +4590,8 @@ where
|
|||||||
debug_assert!(symbol_ids.windows(2).all(|window| window[0] < window[1]));
|
debug_assert!(symbol_ids.windows(2).all(|window| window[0] < window[1]));
|
||||||
(*date, symbol_ids)
|
(*date, symbol_ids)
|
||||||
})
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.into_iter()
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -4668,7 +4671,7 @@ fn build_factor_market_cap_order(
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_dense_row_positions<T>(
|
fn build_dense_row_positions<T: Sync>(
|
||||||
groups: &BTreeMap<NaiveDate, Vec<T>>,
|
groups: &BTreeMap<NaiveDate, Vec<T>>,
|
||||||
symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
|
symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
|
||||||
symbol_count: usize,
|
symbol_count: usize,
|
||||||
@@ -4679,8 +4682,11 @@ fn build_dense_row_positions<T>(
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut positions_by_date = BTreeMap::new();
|
// Each task owns one bounded day index. No partial index is published if
|
||||||
for (date, rows) in groups {
|
// 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)?;
|
let symbol_ids = symbol_ids_by_date.get(date)?;
|
||||||
if rows.len() != symbol_ids.len() {
|
if rows.len() != symbol_ids.len() {
|
||||||
return None;
|
return None;
|
||||||
@@ -4693,9 +4699,10 @@ fn build_dense_row_positions<T>(
|
|||||||
}
|
}
|
||||||
*position = u32::try_from(row_index).ok()?;
|
*position = u32::try_from(row_index).ok()?;
|
||||||
}
|
}
|
||||||
positions_by_date.insert(*date, positions);
|
Some((*date, positions))
|
||||||
}
|
})
|
||||||
Some(positions_by_date)
|
.collect::<Option<Vec<_>>>()
|
||||||
|
.map(|days| days.into_iter().collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_calendar_series_end_positions(
|
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]
|
#[test]
|
||||||
fn direct_symbol_id_snapshot_lookups_preserve_alignment_for_sparse_rows() {
|
fn direct_symbol_id_snapshot_lookups_preserve_alignment_for_sparse_rows() {
|
||||||
let date = NaiveDate::parse_from_str("2025-01-02", "%Y-%m-%d").unwrap();
|
let date = NaiveDate::parse_from_str("2025-01-02", "%Y-%m-%d").unwrap();
|
||||||
|
|||||||
Reference in New Issue
Block a user