claude_dev
d5582fd6f2
deep union 段 glob '*_union.parquet' 误匹配 000852_announce_union.parquet/ 932000_announce_union.parquet, 它们无 in_current/was_removed 列(用 adjust_type) -> KeyError 崩溃。显式 skip _announce_union.parquet 文件。 VPS 验证通过: - 000852: 1672 rows (in_current=1000, was_removed=672) - 932000: 2684 rows (in_current=2000, was_removed=684) - 000300/000905/000016/399001 行数不变(回归 OK)
245 lines
11 KiB
Python
245 lines
11 KiB
Python
#!/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 _read_announce_union_aggregated(hist_dir: str,
|
||
index_codes: list[str]) -> pd.DataFrame:
|
||
"""读 *_announce_union.parquet, 聚合成全集型(index_code/code/code_name/
|
||
in_current/was_removed/source)。
|
||
|
||
announce_union schema:
|
||
updateDate / index_code / code / code_name / adjust_type / notice_id / source
|
||
adjust_type: add | remove | current | initial
|
||
|
||
全集聚合(spec §14):
|
||
ever_codes = announce_union 所有 distinct code(任意 adjust_type)
|
||
current_codes = announce_union 中 adjust_type='current' 的 code
|
||
(∪ <idx>_snapshot.parquet 兜底, 兼容旧 announce 无 current 行)
|
||
in_current = code in current_codes
|
||
was_removed = not in_current(曾经入选已踢)
|
||
code_name = 优先 snapshot 当前名(announce 的历史名可能过时)
|
||
|
||
缺 announce_union 文件 -> 返回空 DataFrame(上层走 snapshot 兜底)。
|
||
"""
|
||
cols = ["index_code", "code", "code_name", "in_current", "was_removed", "source"]
|
||
rows = []
|
||
for idx in index_codes:
|
||
ann_path = os.path.join(hist_dir, f"{idx}_announce_union.parquet")
|
||
if not os.path.exists(ann_path):
|
||
continue
|
||
ann = pd.read_parquet(ann_path)
|
||
ann["code"] = ann["code"].apply(norm_code)
|
||
# current 兜底: announce 内 current 行 + snapshot 文件
|
||
current_codes = set(ann.loc[ann["adjust_type"] == "current", "code"])
|
||
snap_path = os.path.join(hist_dir, f"{idx}_snapshot.parquet")
|
||
snap_name_map: dict[str, str] = {}
|
||
if os.path.exists(snap_path):
|
||
snap = pd.read_parquet(snap_path)
|
||
snap["code"] = snap["code"].apply(norm_code)
|
||
current_codes |= set(snap["code"])
|
||
snap_name_map = dict(zip(snap["code"], snap["code_name"]))
|
||
# ever = distinct(code), name 取 announce 或 snapshot(优先 snapshot)
|
||
ever = (
|
||
ann[["code", "code_name"]]
|
||
.drop_duplicates(subset=["code"])
|
||
)
|
||
for _, r in ever.iterrows():
|
||
code = r["code"]
|
||
in_cur = code in current_codes
|
||
name = snap_name_map.get(code) or r["code_name"]
|
||
rows.append({
|
||
"index_code": idx, "code": code, "code_name": name,
|
||
"in_current": int(in_cur), "was_removed": int(not in_cur),
|
||
"source": "csindex_announce",
|
||
})
|
||
if not rows:
|
||
return pd.DataFrame(columns=cols)
|
||
return pd.DataFrame(rows, columns=cols)
|
||
|
||
|
||
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 (排除 *_announce_union.parquet, 不同 schema)
|
||
deep = []
|
||
for f in sorted(glob.glob(os.path.join(HIST, "*_union.parquet"))):
|
||
if f.endswith("_announce_union.parquet"):
|
||
continue
|
||
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. 中证 1000/2000 announce_union 全集(优先); 缺则回退 snapshot 当前
|
||
ANNOUNCE_IDX = ["000852", "932000"]
|
||
df_csi = _read_announce_union_aggregated(HIST, ANNOUNCE_IDX)
|
||
if df_csi.empty:
|
||
# 回退: 旧 snapshot-only 路径(announce_union 未生成时)
|
||
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_csi = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame(
|
||
columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"])
|
||
df_csi["code"] = df_csi["code"].apply(norm_code)
|
||
print(f"[中证 snapshot 回退] rows={len(df_csi)}")
|
||
else:
|
||
print(f"[中证 announce_union] rows={len(df_csi)}, "
|
||
f"indices={df_csi['index_code'].nunique()}")
|
||
|
||
# 合并 + 去重 (同 index+code+source)
|
||
all_df = pd.concat([pool, df_deep, df_csi], 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(*), SUM(in_current), SUM(was_removed) FROM "
|
||
f"{STAGING} WHERE index_code='000852'").fetchone())
|
||
print("sample 932000:", c.execute(
|
||
"SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM "
|
||
f"{STAGING} WHERE index_code='932000'").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)
|