diff --git a/sanguo_factor/universe.py b/sanguo_factor/universe.py index ce7fd13..aa7714e 100644 --- a/sanguo_factor/universe.py +++ b/sanguo_factor/universe.py @@ -10,9 +10,10 @@ from datetime import datetime, timedelta import polars as pl +from sanguo_data.datareader import guess_exchange + STOCK_PREFIXES = ("60", "00", "30") WARMUP_BARS = 120 -_CHUNK = 300 # 分块 IN 查询每块 symbol 数(控瞬时内存) _LOOKBACK_DAYS = 300 # start 前缓冲(覆盖最大 60 日窗口 + 节假日) _FORWARD_DAYS = 45 # end 后缓冲(覆盖 10 日前瞻收益) @@ -37,11 +38,16 @@ def load_universe_bars( if symbols is None: conn = sqlite3.connect(vnpy_db, timeout=60) try: - # 无过滤 DISTINCT symbol 走 (symbol,...) 前导索引顺序流式扫—— - # 带 WHERE(interval/datetime/LIKE)的版本会退化为 26G 全表扫(NAS 实测>5min)。 - # 前缀在 Python 侧滤;interval='d'/窗口过滤由下方分块数据查询天然承担 - # (无日线数据的 symbol 返回 0 行,不进最终 df,语义不变)。 - cur = conn.execute("SELECT DISTINCT symbol FROM dbbardata") + # 松散索引扫(递归 CTE):每符号一次 seek(NAS 实测 70s vs 朴素 DISTINCT 573s—— + # 34M 索引条目 2核NAS 走不完;seek 次数=符号数≈7700) + cur = conn.execute( + "WITH RECURSIVE s(sym) AS (" + " SELECT min(symbol) FROM dbbardata" + " UNION ALL" + " SELECT (SELECT min(symbol) FROM dbbardata WHERE symbol > s.sym)" + " FROM s WHERE s.sym IS NOT NULL" + ") SELECT sym FROM s WHERE sym IS NOT NULL" + ) symbols = [r[0] for r in cur if str(r[0]).startswith(STOCK_PREFIXES)] finally: conn.close() @@ -52,16 +58,16 @@ def load_universe_bars( symbols = sorted(random.Random(42).sample(symbols, limit)) chunks: list[pl.DataFrame] = [] - for i in range(0, len(symbols), _CHUNK): - part = symbols[i : i + _CHUNK] - ph = ",".join("?" * len(part)) - conn = sqlite3.connect(vnpy_db, timeout=30) + for sym in symbols: + ex = guess_exchange(sym).value # "SSE"/"SZSE";60→SSE,00/30→SZSE 与入库一致 + conn = sqlite3.connect(vnpy_db, timeout=60) conn.execute("PRAGMA busy_timeout=30000") try: cur = conn.execute( f"SELECT {_COLS} FROM dbbardata " - f"WHERE interval='d' AND datetime>=? AND datetime<=? AND symbol IN ({ph})", - (lookback_start, forward_end, *part), + "WHERE symbol=? AND exchange=? AND interval='d' AND datetime>=? AND datetime<=? " + "ORDER BY datetime", + (sym, ex, lookback_start, forward_end), ) rows = cur.fetchall() finally: