重构分钟线事件流与订阅加载

2026-08-25 01:41:50 +08:00
parent 4cf0224d2d
commit a147c495af
6 changed files with 316 additions and 47 deletions
@@ -1,5 +1,6 @@
 use std::borrow::Cow;
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::cmp::Reverse;
+use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet};
 use std::sync::{Arc, OnceLock};

 use ahash::AHashMap;
@@ -1618,19 +1619,38 @@ impl DataSet {
    }

    pub fn execution_quotes_on_date(&self, date: NaiveDate) -> Vec<IntradayExecutionQuote> {
-        let mut quotes = self
-            .execution_quotes_by_date
-            .get(&date)
-            .into_iter()
-            .flat_map(|rows_by_symbol| rows_by_symbol.values())
-            .flat_map(|rows| rows.iter().cloned())
+        let Some(rows_by_symbol) = self.execution_quotes_by_date.get(&date) else {
+            return Vec::new();
+        };
+        let mut streams = rows_by_symbol
+            .iter()
+            .map(|(symbol, rows)| (symbol.as_str(), rows.as_slice()))
            .collect::<Vec<_>>();
-        quotes.sort_by(|left, right| {
-            left.timestamp
-                .cmp(&right.timestamp)
-                .then_with(|| left.symbol.cmp(&right.symbol))
-        });
-        quotes
+        streams.sort_by_key(|(symbol, _)| *symbol);
+        let total_rows = streams.iter().map(|(_, rows)| rows.len()).sum();
+        let mut heap = BinaryHeap::<Reverse<(NaiveDateTime, usize, usize)>>::new();
+        for (stream_index, (_, rows)) in streams.iter().enumerate() {
+            if let Some(first) = rows.first() {
+                heap.push(Reverse((first.timestamp, stream_index, 0)));
+            }
+        }
+        let mut merged = Vec::with_capacity(total_rows);
+        while let Some(Reverse((_timestamp, stream_index, row_index))) = heap.pop() {
+            let (_, rows) = streams[stream_index];
+            merged.push(rows[row_index].clone());
+            let next_index = row_index + 1;
+            if let Some(next) = rows.get(next_index) {
+                heap.push(Reverse((next.timestamp, stream_index, next_index)));
+            }
+        }
+        merged
+    }
+
+    pub fn remove_execution_quotes_on_date(&mut self, date: NaiveDate) -> usize {
+        self.execution_quotes_by_date
+            .remove(&date)
+            .map(|rows_by_symbol| rows_by_symbol.into_values().map(|rows| rows.len()).sum())
+            .unwrap_or_default()
    }

    pub fn snapshot_components(&self) -> DataSetSnapshotComponents {
@@ -3728,6 +3748,68 @@ mod tests {
        assert_eq!(run_data.execution_quote_count(), 1);
    }

+    #[test]
+    fn execution_quotes_use_stable_k_way_merge_and_release_by_date() {
+        let date = NaiveDate::parse_from_str("2025-01-02", "%Y-%m-%d").unwrap();
+        let data = DataSet::from_components(
+            vec![Instrument {
+                symbol: "000001.SZ".to_string(),
+                name: "平安银行".to_string(),
+                board: "SZ".to_string(),
+                round_lot: 100,
+                listed_at: None,
+                delisted_at: None,
+                status: "active".to_string(),
+            }],
+            vec![market_row("2025-01-02", 10.0, 1_000_000)],
+            Vec::new(),
+            Vec::new(),
+            vec![benchmark_row("2025-01-02", 12.0)],
+        )
+        .unwrap();
+        let quote = |symbol: &str, time: &str| IntradayExecutionQuote {
+            date,
+            timestamp: NaiveDateTime::parse_from_str(
+                &format!("2025-01-02 {time}"),
+                "%Y-%m-%d %H:%M:%S",
+            )
+            .unwrap(),
+            symbol: symbol.to_string(),
+            last_price: 10.0,
+            bid1: 0.0,
+            ask1: 0.0,
+            bid1_volume: 0,
+            ask1_volume: 0,
+            volume_delta: 100,
+            amount_delta: 1_000.0,
+            trading_phase: Some("continuous".to_string()),
+        };
+        let mut run_data = data.clone();
+        run_data.add_execution_quotes(vec![
+            quote("000002.SZ", "09:31:00"),
+            quote("000001.SZ", "09:31:00"),
+            quote("000002.SZ", "09:30:00"),
+            quote("000001.SZ", "09:30:00"),
+        ]);
+
+        let merged = run_data.execution_quotes_on_date(date);
+        let keys = merged
+            .iter()
+            .map(|row| (row.timestamp.time().to_string(), row.symbol.clone()))
+            .collect::<Vec<_>>();
+        assert_eq!(
+            keys,
+            vec![
+                ("09:30:00".to_string(), "000001.SZ".to_string()),
+                ("09:30:00".to_string(), "000002.SZ".to_string()),
+                ("09:31:00".to_string(), "000001.SZ".to_string()),
+                ("09:31:00".to_string(), "000002.SZ".to_string()),
+            ]
+        );
+        assert_eq!(run_data.remove_execution_quotes_on_date(date), 4);
+        assert_eq!(run_data.execution_quote_count(), 0);
+    }
+
    #[test]
    fn baseline_selection_uses_structured_instrument_dates_and_status_only() {
        let date = NaiveDate::parse_from_str("2025-01-02", "%Y-%m-%d").unwrap();