akg-factor-bridge/common.py

108 lines
4.5 KiB
Python
Raw Normal View History

2026-07-24 14:20:54 +08:00
"""共用:股票码规范化、覆盖池、幂等写因子表、注册 factor_metadata自适应列"""
import json
import pandas as pd
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。"""
df = db.read_pg("SELECT ts_code FROM v_factor_universe")
return set(df["ts_code"].astype(str).str.strip())
def trading_days(start: str, end: str) -> list:
"""目标区间交易日历——用热度表(153日频、覆盖广)的 distinct trade_date 近似。"""
df = db.read_mysql("heat",
"SELECT DISTINCT trade_date FROM stock_fund_heat_scores "
"WHERE trade_date BETWEEN %s AND %s ORDER BY trade_date",
(start, end))
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) -> None:
"""df 列 [trade_date, stock_code, factor_value] → 幂等写因子表。
daily/history 都是删涉及日期区间 批插stock_code 统一转前缀式"""
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))
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))
cur.executemany(
f"INSERT INTO {table} (trade_date, stock_code, factor_value) VALUES (%s,%s,%s)",
rows)
conn.commit()
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:
cur.execute("SELECT column_name FROM information_schema.columns "
"WHERE table_schema=DATABASE() AND table_name='factor_metadata'")
cols = {r[0] for r in cur.fetchall()}
use = [k for k in values if k in cols]
if "factor_code" not in use:
raise RuntimeError("factor_metadata 无 factor_code 列?请核实线上 schema")
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') 过滤,缺列可能查不到本因子,请核实。")