160 lines
7.3 KiB
Python
160 lines
7.3 KiB
Python
"""共用：股票码规范化、覆盖池、交易日历、幂等写因子表、注册 factor_metadata（自适应列）。"""
|
||
import json
|
||
|
||
import pandas as pd
|
||
|
||
import config
|
||
import db
|
||
|
||
|
||
def to_prefix(ts_code: str) -> str:
|
||
"""600000.SH -> SH600000；已是前缀式(SZ002625)或纯代码则原样返回。"""
|
||
s = str(ts_code).strip()
|
||
if "." in s:
|
||
num, exch = s.split(".", 1)
|
||
return f"{exch.upper()}{num}"
|
||
return s
|
||
|
||
|
||
def load_universe() -> set[str]:
|
||
"""KG 覆盖池 ts_code 集合（600000.SH 形态），来自基座视图 v_factor_universe。
|
||
|
||
⚠️ 这是**当前态**：industry_pools 只存最新快照、每周一 refresh_pools 自动长大。
|
||
用它过滤历史数据会引入成员性前视——见 universe_filter 与 config.SUBFACTOR_UNIVERSE。
|
||
"""
|
||
df = db.read_pg("SELECT ts_code FROM v_factor_universe")
|
||
return set(df["ts_code"].astype(str).str.strip())
|
||
|
||
|
||
def universe_filter(df: pd.DataFrame, col: str = "stock_code",
|
||
as_prefix: bool = True) -> pd.DataFrame:
|
||
"""按 config.SUBFACTOR_UNIVERSE 决定子因子是否受覆盖池限制。
|
||
|
||
'pool'（默认，保持原行为）：只留池内成员。
|
||
'market'（评审 §4 建议）：全市场出行，池只在 akg_score 侧当过滤器。
|
||
|
||
为什么建议 market：industry_pools 只存最新态且每周一长大，用它过滤历史子因子，
|
||
值是点时正确的、**成员性却是前视的**（用今天的池去筛历史，而今天的池里有大量
|
||
主题是后来的研报才长出来的——正是 §2.2 批评传导时用的同一条论证）；
|
||
同时截面统计量会在每周一发生结构性跳变，IC 序列带周期性伪影。
|
||
子因子表是「可独立观察的仪表」，过滤该发生在消费端。
|
||
"""
|
||
if config.SUBFACTOR_UNIVERSE == "pool":
|
||
uni = load_universe()
|
||
if as_prefix:
|
||
uni = {to_prefix(x) for x in uni}
|
||
return df[df[col].astype(str).str.strip().isin(uni)]
|
||
return df
|
||
|
||
|
||
def trading_days(start: str, end: str) -> list:
|
||
"""目标区间交易日历。
|
||
|
||
⚠️ 曾用热度表 stock_fund_heat_scores —— 但它最早只有 2025-03-26，
|
||
导致 `build akg_event --mode history --start 2024-01-01`（README 里的示例命令）
|
||
拿到空日历 → 静默返回空 → 只打印「无数据（跳过）」，不报错。
|
||
改用行情表 gp_day_data（实测 5584 天，覆盖全历史）。
|
||
"""
|
||
df = db.read_mysql(
|
||
"price",
|
||
"SELECT DISTINCT `timestamp` AS trade_date FROM gp_day_data "
|
||
"WHERE `timestamp` BETWEEN %s AND %s ORDER BY `timestamp`",
|
||
(start, end))
|
||
if df.empty:
|
||
print(f" ⚠️ 交易日历为空（gp_day_data 在 {start}~{end} 无数据）"
|
||
f"——上层会跳过该区间，请核对区间与行情表覆盖")
|
||
return list(pd.to_datetime(df["trade_date"]))
|
||
|
||
|
||
_CREATE_FACTOR_TABLE = """
|
||
CREATE TABLE IF NOT EXISTS {t} (
|
||
trade_date DATE NOT NULL,
|
||
stock_code VARCHAR(15) NOT NULL,
|
||
factor_value DOUBLE,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
PRIMARY KEY (trade_date, stock_code),
|
||
KEY idx_stock (stock_code)
|
||
) ENGINE=InnoDB
|
||
"""
|
||
|
||
|
||
def ensure_table(table: str) -> None:
|
||
with db.factor_conn() as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute(_CREATE_FACTOR_TABLE.format(t=table))
|
||
conn.commit()
|
||
|
||
|
||
def write_factor(table: str, df: pd.DataFrame, mode: str = "daily") -> None:
|
||
"""df 列 [trade_date, stock_code, factor_value] → 幂等写因子表。
|
||
daily/history 都是「删涉及日期区间 → 批插」。stock_code 统一转前缀式。
|
||
|
||
分块提交（评审）：热度全史约 54 万行，原来单事务 executemany 走 ShardingSphere
|
||
代理有风险。改成 DELETE 一个事务 + INSERT 按 config.WRITE_CHUNK_ROWS 分块提交。
|
||
代价是中途失败会留下部分区间——但整段写入按日期区间幂等，重跑即修复。
|
||
"""
|
||
if df is None or df.empty:
|
||
print(f" {table}: 无数据（跳过）")
|
||
return
|
||
df = df.dropna(subset=["trade_date", "stock_code", "factor_value"]).copy()
|
||
df["stock_code"] = df["stock_code"].map(to_prefix)
|
||
df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date
|
||
df = df.drop_duplicates(["trade_date", "stock_code"], keep="last")
|
||
if df.empty:
|
||
print(f" {table}: 清洗后无数据")
|
||
return
|
||
dmin, dmax = df["trade_date"].min(), df["trade_date"].max()
|
||
ensure_table(table)
|
||
rows = list(df[["trade_date", "stock_code", "factor_value"]]
|
||
.itertuples(index=False, name=None))
|
||
chunk = max(1, config.WRITE_CHUNK_ROWS)
|
||
with db.factor_conn() as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute(f"DELETE FROM {table} WHERE trade_date BETWEEN %s AND %s",
|
||
(dmin, dmax))
|
||
conn.commit()
|
||
for i in range(0, len(rows), chunk):
|
||
with conn.cursor() as cur:
|
||
cur.executemany(
|
||
f"INSERT INTO {table} (trade_date, stock_code, factor_value) "
|
||
f"VALUES (%s,%s,%s)", rows[i:i + chunk])
|
||
conn.commit()
|
||
if len(rows) > chunk:
|
||
print(f" ...{min(i + chunk, len(rows))}/{len(rows)}")
|
||
print(f" {table}: 写入 {len(rows)} 行, 日期 {dmin}~{dmax}")
|
||
|
||
|
||
def register(factor_code, display_name, table, category, desc,
|
||
author="akg-factor-bridge") -> None:
|
||
"""注册 factor_metadata。自适应实际存在的列（避免猜死 schema）——
|
||
列名以线上 factor_metadata 为准，本函数只写它有的列。"""
|
||
values = {
|
||
"factor_code": factor_code, "display_name": display_name,
|
||
"target_ds_name": "ds_a", "target_table_name": table,
|
||
"author": author, "description": desc,
|
||
"category": json.dumps(category, ensure_ascii=False),
|
||
"frequency": "daily", "status": "active", "factor_type": "single",
|
||
}
|
||
with db.factor_conn() as conn:
|
||
with conn.cursor() as cur:
|
||
# ShardingSphere 代理不支持 information_schema 内省 → 用真实查询取列名
|
||
cur.execute("SELECT * FROM factor_metadata LIMIT 0")
|
||
cols = [d[0] for d in (cur.description or [])]
|
||
if not cols: # 极端兜底：代理连列元数据都不给时，用平台已知 schema
|
||
cols = ["factor_code", "display_name", "target_ds_name", "target_table_name",
|
||
"start_date", "end_date", "author", "description", "category",
|
||
"frequency", "status", "factor_type"]
|
||
use = [k for k in values if k in cols]
|
||
if "factor_code" not in use:
|
||
raise RuntimeError(f"factor_metadata 取不到 factor_code 列；实际列: {cols}")
|
||
ph = ",".join(["%s"] * len(use))
|
||
upd = ",".join(f"{k}=VALUES({k})" for k in use if k != "factor_code")
|
||
cur.execute(
|
||
f"INSERT INTO factor_metadata ({','.join(use)}) VALUES ({ph}) "
|
||
f"ON DUPLICATE KEY UPDATE {upd}", [values[k] for k in use])
|
||
conn.commit()
|
||
print(f" 注册 {factor_code} -> {table}（写入列: {sorted(use)}）")
|
||
if "factor_type" not in cols:
|
||
print(" ⚠️ factor_metadata 无 factor_type 列——平台 /mining/factors/all 按 "
|
||
"factor_type IN('single','multiple') 过滤，缺列可能查不到本因子，请核实。")
|