#!/usr/bin/env python3 # -*- coding: utf-8 -*- """migrate_constituent.py — 单元3: 合并成份股 -> constituent_unified_staging (全集型, 幂等)。 设计 (spec §14, 治幸存者偏差选股池): - baostock (``bs_index_constituent_old`` 988 时点全集, 方案A 后的权威历史源) -> 聚合成全集 (hs300/zz500/sz50 -> 000300/000905/000016), in_current=最后时点成份, was_removed=历史入选过但已踢出 - 深证/国证 _union.parquet (399001/399006/399005/399330) -> 直入 (akshare cni) - 中证1000/2000 _snapshot.parquet (000852/932000) -> 当前 (akshare csindex) - 新浪 _sina.parquet -> 丢弃 (baostock 300/500/50 已权威) - code 统一 6 位无前缀 (sh.600000 / 600519.SH -> 600519) 幂等(可重跑, 月度 schtask 安全): - ``DROP TABLE IF EXISTS constituent_unified_staging`` + ``CREATE TABLE ...`` 每次重建 - baostock 源表读取**自动适配**: * 方案A 后正常只有 ``bs_index_constituent_old``(权威全集) * 若将来重建了 live ``bs_index_constituent``(新时点增量), UNION ALL 两者去重, 兼容两种状态 环境变量: - ``SANGUO_DB``: quant_trading.db 路径, 默认 ``C:\\sanguo_vnpy_v2\\data\\quant_trading.db``(VPS) - ``HIST``(模块常量, 测试可 monkeypatch): 深证/中证 parquet 目录 """ import glob import os import re import sqlite3 import pandas as pd HIST = r"C:\sanguo_vnpy_v2\data\index_const_hist" STAGING = "constituent_unified_staging" BS_MAP = {"hs300": "000300", "zz500": "000905", "sz50": "000016"} def norm_code(code): s = re.sub(r"^(sh|sz|SH|SZ)\.?", "", str(code)) s = re.sub(r"\.(SH|SZ|sh|sz)$", "", s) return s.zfill(6) if s.isdigit() and len(s) <= 6 else s def _read_baostock_constituent(c: sqlite3.Connection) -> pd.DataFrame: """读 baostock 成份股历史, 自动适配方案A后(_old) / 未来(live)两种状态。 - 只有 ``bs_index_constituent_old``: 读它(方案A 后常态) - 只有 ``bs_index_constituent``: 读它(未来重建 live 表) - 两者都在: UNION ALL 后 drop_duplicates(兼容过渡期) - 都没有: 返回空 DataFrame(不崩, 由上层决定是否报错) 返回字段: ``updateDate, index_code, code, code_name``。 """ cur = c.execute( "SELECT name FROM sqlite_master WHERE type='table' AND " "name IN ('bs_index_constituent_old', 'bs_index_constituent')" ) tables = {row[0] for row in cur.fetchall()} parts = [] if "bs_index_constituent_old" in tables: parts.append(pd.read_sql( "SELECT updateDate, index_code, code, code_name " "FROM bs_index_constituent_old", c, )) print("[baostock] read from bs_index_constituent_old (方案A 权威历史全集)") if "bs_index_constituent" in tables: parts.append(pd.read_sql( "SELECT updateDate, index_code, code, code_name " "FROM bs_index_constituent", c, )) print("[baostock] read from bs_index_constituent (live 增量)") if not parts: print("[baostock] WARN: 既无 _old 也无 live 表, baostock 段产出 0 行") return pd.DataFrame(columns=["updateDate", "index_code", "code", "code_name"]) df = pd.concat(parts, ignore_index=True).drop_duplicates() return df def migrate(db_path: str) -> None: """(幂等)重建 constituent_unified_staging。 Args: db_path: quant_trading.db 路径(测试可传 tmp sqlite; 生产读 ``SANGUO_DB``)。 """ c = sqlite3.connect(db_path, timeout=60) try: c.execute("PRAGMA busy_timeout = 60000") # 1. baostock -> 全集(自动适配 _old / live / 两者皆在) df_bs = _read_baostock_constituent(c) df_bs["index_code"] = df_bs["index_code"].map(BS_MAP) df_bs["code"] = df_bs["code"].apply(norm_code) last_sets = {} for idx, grp in df_bs.groupby("index_code"): last_d = grp["updateDate"].max() last_sets[idx] = set(grp[grp["updateDate"] == last_d]["code"]) pool = (df_bs.groupby(["index_code", "code"])["code_name"] .first().reset_index()) pool["in_current"] = pool.apply( lambda r: r["code"] in last_sets.get(r["index_code"], set()), axis=1) pool["was_removed"] = ~pool["in_current"] pool["source"] = "baostock" print(f"[baostock] pool rows={len(pool)} (300/500/50 全集)") # 2. 深证 union deep = [] for f in sorted(glob.glob(os.path.join(HIST, "*_union.parquet"))): d = pd.read_parquet(f)[["index_code", "code", "code_name", "in_current", "was_removed"]] d["source"] = "akshare_cni" deep.append(d) df_deep = pd.concat(deep, ignore_index=True) if deep else pd.DataFrame( columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"]) df_deep["code"] = df_deep["code"].apply(norm_code) print(f"[深证 union] rows={len(df_deep)}") # 3. 中证 snapshot snap = [] for f in [os.path.join(HIST, "000852_snapshot.parquet"), os.path.join(HIST, "932000_snapshot.parquet")]: if os.path.exists(f): d = pd.read_parquet(f)[["index_code", "code", "code_name"]] d["in_current"] = True d["was_removed"] = False d["source"] = "akshare_csindex" snap.append(d) df_snap = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame( columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"]) df_snap["code"] = df_snap["code"].apply(norm_code) print(f"[中证 snapshot] rows={len(df_snap)}") # 合并 + 去重 (同 index+code+source) all_df = pd.concat([pool, df_deep, df_snap], ignore_index=True) all_df = all_df.drop_duplicates(["index_code", "code", "source"]) print(f"\n[TOTAL] constituent_unified: {len(all_df)} rows, " f"{all_df['index_code'].nunique()} indices") print("\n各指数分布:") print(all_df.groupby("index_code").agg( n=("code", "count"), src=("source", "first"), in_cur=("in_current", "sum"), removed=("was_removed", "sum"))) # 写 staging(幂等: DROP+CREATE) c.execute(f"DROP TABLE IF EXISTS {STAGING}") c.execute(f"""CREATE TABLE {STAGING} ( index_code TEXT, code TEXT, code_name TEXT, source TEXT, in_current INTEGER, was_removed INTEGER)""") work = all_df[["index_code", "code", "code_name", "source", "in_current", "was_removed"]].copy() work["in_current"] = work["in_current"].astype(int) work["was_removed"] = work["was_removed"].astype(int) c.executemany(f"INSERT INTO {STAGING} VALUES (?,?,?,?,?,?)", work.itertuples(index=False, name=None)) c.commit() n = c.execute(f"SELECT COUNT(*) FROM {STAGING}").fetchone()[0] # 抽样验证 print(f"\n[staging] {STAGING}: {n} rows") print("sample 300:", c.execute( "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM " f"{STAGING} WHERE index_code='000300'").fetchone()) print("sample 399001:", c.execute( "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM " f"{STAGING} WHERE index_code='399001'").fetchone()) print("sample 000852:", c.execute( "SELECT COUNT(*) FROM " f"{STAGING} WHERE index_code='000852'").fetchone()) finally: c.close() print("\nMIGRATE STAGING DONE (未 rename, 验证 OK 后单独合并)") if __name__ == "__main__": _db = os.environ.get( "SANGUO_DB", r"C:\sanguo_vnpy_v2\data\quant_trading.db" ) migrate(_db)