perf: build calendar series boundaries in one pass

2026-09-06 22:18:54 +08:00
parent 0d831c4ca6
commit f927ef8c0f
1 changed files with 93 additions and 44 deletions
@@ -1397,9 +1397,12 @@ impl DataSet {
        let mut calendar_dates = self.calendar.days().to_vec();
        calendar_dates.extend(dates);
        let calendar = Arc::new(TradingCalendar::new(calendar_dates));
-        self.market_series_end_positions_by_calendar_index = Arc::new(
-            build_calendar_series_end_positions(&self.market_series_by_symbol_id, &calendar),
-        );
+        self.market_series_end_positions_by_calendar_index =
+            Arc::new(build_calendar_series_end_positions(
+                &self.market_series_by_symbol_id,
+                &self.market_symbol_ids_by_date,
+                &calendar,
+            ));
        self.calendar = calendar;
        self
    }
@@ -1821,8 +1824,11 @@ impl DataSet {
                adjusted_close_series_by_symbol_id[symbol_id as usize] = Some(Arc::clone(series));
            }
        }
-        let market_series_end_positions_by_calendar_index =
-            build_calendar_series_end_positions(&market_series_by_symbol_id, &calendar);
+        let market_series_end_positions_by_calendar_index = build_calendar_series_end_positions(
+            &market_series_by_symbol_id,
+            &market_symbol_ids_by_date,
+            &calendar,
+        );
        let execution_quotes_by_date = build_execution_quote_index(execution_quotes);
        let mut execution_quote_dates =
            execution_quotes_by_date.keys().copied().collect::<Vec<_>>();
@@ -4326,6 +4332,7 @@ fn build_dense_row_positions<T>(

 fn build_calendar_series_end_positions(
    series_by_symbol_id: &[Option<Arc<SymbolPriceSeries>>],
+    market_symbol_ids_by_date: &BTreeMap<NaiveDate, Vec<u32>>,
    calendar: &TradingCalendar,
 ) -> Option<CalendarSeriesEndPositions> {
    let entries = series_by_symbol_id.len().checked_mul(calendar.len())?;
@@ -4340,54 +4347,65 @@ fn build_calendar_series_end_positions(
        return None;
    }

-    let calendar_days = calendar.days();
-    let positions_by_symbol = series_by_symbol_id
-        .par_iter()
+    let mut completed_rows_by_symbol = series_by_symbol_id
+        .iter()
        .map(|series| {
-            let series = series.as_deref()?;
-            let mut decision = Vec::with_capacity(calendar_days.len());
-            let mut current = Vec::with_capacity(calendar_days.len());
-            let mut series_index = 0usize;
-            for date in calendar_days {
-                while series
-                    .dates
-                    .get(series_index)
-                    .is_some_and(|series_date| *series_date < *date)
-                {
-                    series_index += 1;
-                }
-                decision.push(series_index as u32);
-                let current_index = if series.dates.get(series_index) == Some(date) {
-                    series_index + 1
-                } else {
-                    series_index
-                };
-                current.push(current_index as u32);
+            if series.is_some() {
+                0
+            } else {
+                MISSING_ROW_POSITION
            }
-            Some((decision, current))
        })
        .collect::<Vec<_>>();
-    let positions_by_calendar = (0..calendar_days.len())
-        .into_par_iter()
-        .map(|calendar_index| {
-            let mut decision = Vec::with_capacity(series_by_symbol_id.len());
-            let mut current = Vec::with_capacity(series_by_symbol_id.len());
-            for positions in &positions_by_symbol {
-                if let Some((symbol_decision, symbol_current)) = positions {
-                    decision.push(symbol_decision[calendar_index]);
-                    current.push(symbol_current[calendar_index]);
-                } else {
-                    decision.push(MISSING_ROW_POSITION);
-                    current.push(MISSING_ROW_POSITION);
+    let mut market_dates = market_symbol_ids_by_date.iter().peekable();
+    let mut decision = Vec::with_capacity(calendar.len());
+    let mut current = Vec::with_capacity(calendar.len());
+
+    for date in calendar.days() {
+        while market_dates
+            .peek()
+            .is_some_and(|(market_date, _)| *market_date < date)
+        {
+            let (_, symbol_ids) = market_dates.next()?;
+            advance_completed_series_rows(&mut completed_rows_by_symbol, symbol_ids)?;
+        }
+
+        decision.push(completed_rows_by_symbol.clone());
+        let mut current_on_date = completed_rows_by_symbol.clone();
+        if market_dates
+            .peek()
+            .is_some_and(|(market_date, _)| *market_date == date)
+        {
+            let (_, symbol_ids) = market_dates.next()?;
+            for symbol_id in symbol_ids {
+                let end = current_on_date.get_mut(*symbol_id as usize)?;
+                if *end == MISSING_ROW_POSITION {
+                    return None;
                }
+                *end = completed_rows_by_symbol[*symbol_id as usize].checked_add(1)?;
            }
-            (decision, current)
-        })
-        .collect::<Vec<_>>();
-    let (decision, current) = positions_by_calendar.into_iter().unzip();
+            advance_completed_series_rows(&mut completed_rows_by_symbol, symbol_ids)?;
+        }
+        current.push(current_on_date);
+    }
+
    Some(CalendarSeriesEndPositions { decision, current })
 }

+fn advance_completed_series_rows(
+    completed_rows_by_symbol: &mut [u32],
+    symbol_ids: &[u32],
+) -> Option<()> {
+    for symbol_id in symbol_ids {
+        let completed = completed_rows_by_symbol.get_mut(*symbol_id as usize)?;
+        if *completed == MISSING_ROW_POSITION {
+            return None;
+        }
+        *completed = completed.checked_add(1)?;
+    }
+    Some(())
+}
+
 fn dense_row_position(
    positions_by_date: &Option<DenseRowPositionIndex>,
    date: NaiveDate,
@@ -6110,6 +6128,37 @@ mod tests {
        );
    }

+    #[test]
+    fn series_end_position_index_preserves_sparse_and_duplicate_date_boundaries() {
+        let first = NaiveDate::from_ymd_opt(2025, 1, 2).unwrap();
+        let duplicate_date = NaiveDate::from_ymd_opt(2025, 1, 3).unwrap();
+        let last = NaiveDate::from_ymd_opt(2025, 1, 6).unwrap();
+        let rows = vec![
+            market_row("2025-01-03", 10.0, 100),
+            market_row("2025-01-03", 10.1, 110),
+            market_row("2025-01-06", 10.2, 120),
+        ];
+        let series = Arc::new(SymbolPriceSeries::new("000001.SZ".to_string(), rows.iter()));
+        let series_by_symbol_id = vec![Some(series), None];
+        let market_symbol_ids_by_date =
+            BTreeMap::from([(duplicate_date, vec![0, 0]), (last, vec![0])]);
+        let calendar = TradingCalendar::new(vec![first, duplicate_date, last]);
+
+        let positions = build_calendar_series_end_positions(
+            &series_by_symbol_id,
+            &market_symbol_ids_by_date,
+            &calendar,
+        )
+        .expect("series position index");
+
+        assert_eq!(positions.decision[0], vec![0, MISSING_ROW_POSITION]);
+        assert_eq!(positions.current[0], vec![0, MISSING_ROW_POSITION]);
+        assert_eq!(positions.decision[1], vec![0, MISSING_ROW_POSITION]);
+        assert_eq!(positions.current[1], vec![1, MISSING_ROW_POSITION]);
+        assert_eq!(positions.decision[2], vec![2, MISSING_ROW_POSITION]);
+        assert_eq!(positions.current[2], vec![3, MISSING_ROW_POSITION]);
+    }
+
    #[test]
    fn source_volume_contract_rejects_windows_containing_missing_values() {
        let data = volume_contract_data(Some([1.0, 0.0, 1.0]));