131 lines
5.5 KiB
Python
131 lines
5.5 KiB
Python
# sanguo_factor/fundamental_statements.py
|
||
"""三表读取层: NAS static 域 income/balance/cashflow parquet → 报告期宽表
|
||
(自 fundamental_adapter.py 拆出,纯结构重构).
|
||
|
||
按股文件路由 + 容错跳过(零行/缺文件/坏文件)+ (vt_symbol, REPORT_DATE)
|
||
去重取重述终值,有效披露日 notice_eff = 三表 NOTICE_DATE 最大值(保守).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
import polars as pl
|
||
|
||
from .fundamental_schema import _BALANCE_COLS, _CUM_MAP, _DATE_COLS, _VT_TO_FILE_SUFFIX
|
||
|
||
# ==================== 读取层 ====================
|
||
|
||
def _vt_to_file_code(vt_symbol: str) -> str | None:
|
||
"""``600000.SSE`` → ``600000.SH``;无法映射的(ETF/北交)返回 None."""
|
||
code, _, suffix = vt_symbol.partition(".")
|
||
file_suffix = _VT_TO_FILE_SUFFIX.get(suffix.upper())
|
||
return f"{code}.{file_suffix}" if file_suffix else None
|
||
|
||
|
||
def _read_static(path: str, want_cols: list[str]) -> pl.DataFrame | None:
|
||
"""读单股单表 parquet,缺列补 null;零行/缺文件/坏文件 → None(容错跳过)."""
|
||
if not os.path.exists(path):
|
||
return None
|
||
try:
|
||
schema = pl.read_parquet_schema(path)
|
||
if "REPORT_DATE" not in schema:
|
||
return None
|
||
cols = [c for c in _DATE_COLS + want_cols if c in schema]
|
||
df = pl.read_parquet(path, columns=cols)
|
||
except Exception:
|
||
return None
|
||
if df.height == 0:
|
||
return None
|
||
missing = [c for c in _DATE_COLS + want_cols if c not in df.columns]
|
||
if missing:
|
||
df = df.with_columns([pl.lit(None, dtype=pl.Utf8).alias(c) for c in missing])
|
||
for c in want_cols:
|
||
df = df.with_columns(pl.col(c).cast(pl.Float64, strict=False))
|
||
return df
|
||
|
||
|
||
def _norm_dates(df: pl.DataFrame, cols: list[str]) -> pl.DataFrame:
|
||
"""日期列归一为 pl.Date:字符串截前 10 位解析,Date/Datetime 直接 cast."""
|
||
exprs = []
|
||
for c in cols:
|
||
dtype = df.schema[c]
|
||
if dtype == pl.Utf8:
|
||
exprs.append(pl.col(c).str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False).alias(c))
|
||
elif dtype != pl.Date:
|
||
exprs.append(pl.col(c).cast(pl.Date, strict=False).alias(c))
|
||
return df.with_columns(exprs) if exprs else df
|
||
|
||
|
||
def _dedupe_reports(df: pl.DataFrame, value_cols: list[str]) -> pl.DataFrame:
|
||
"""按 (vt_symbol, REPORT_DATE) 去重:值取 (UPDATE_DATE, NOTICE_DATE) 排序末行
|
||
(重述取终值),有效披露日取组内 NOTICE_DATE 最大值(保守,不提前)."""
|
||
df = df.sort(["vt_symbol", "REPORT_DATE", "UPDATE_DATE", "NOTICE_DATE"],
|
||
nulls_last=False)
|
||
aggs = [pl.col(c).last().alias(c) for c in value_cols]
|
||
aggs.append(pl.col("NOTICE_DATE").max().alias("notice_eff"))
|
||
return df.group_by(["vt_symbol", "REPORT_DATE"]).agg(aggs)
|
||
|
||
|
||
# 各表显式读取清单(income/cashflow 列集不相交,join 不加后缀)
|
||
_TABLE_RAW = {
|
||
"income": ["TOTAL_OPERATE_INCOME", "OPERATE_COST", "PARENT_NETPROFIT",
|
||
"DEDUCT_PARENT_NETPROFIT", "TOTAL_PROFIT", "INVEST_INCOME",
|
||
"FAIRVALUE_CHANGE_INCOME", "ASSET_IMPAIRMENT_LOSS",
|
||
"CREDIT_IMPAIRMENT_LOSS",
|
||
"RESEARCH_EXPENSE", "SALE_EXPENSE", "FE_INTEREST_EXPENSE",
|
||
"INCOME_TAX", "BASIC_EPS"],
|
||
"balance": _BALANCE_COLS,
|
||
"cashflow": ["NETCASH_OPERATE", "SALES_SERVICES", "ACCEPT_INVEST_CASH",
|
||
"FA_IR_DEPR", "IA_AMORTIZE", "LPE_AMORTIZE",
|
||
"USERIGHT_ASSET_AMORTIZE", "CONSTRUCT_LONG_ASSET",
|
||
"RECEIVE_LOAN_CASH", "ISSUE_BOND", "PAY_DEBT_CASH"],
|
||
}
|
||
|
||
|
||
def _load_statements(codes: list[str], data_dir: str) -> pl.DataFrame:
|
||
"""全 codes 三表 → 报告期宽表(每 vt_symbol × REPORT_DATE 一行).
|
||
|
||
income/cashflow 值列重命名为短名(_CUM_MAP),balance 保持源列名;
|
||
anchor = 三表报告期并集,left join 保证单表缺期不拖垮其它表特征。
|
||
"""
|
||
table_cols = _TABLE_RAW
|
||
per_table: dict[str, pl.DataFrame] = {}
|
||
for table, raw_cols in table_cols.items():
|
||
frames = []
|
||
for vt in codes:
|
||
file_code = _vt_to_file_code(vt)
|
||
if file_code is None:
|
||
continue
|
||
df = _read_static(
|
||
os.path.join(data_dir, table, f"{file_code}_{table}.parquet"), raw_cols)
|
||
if df is None:
|
||
continue
|
||
df = _norm_dates(df, _DATE_COLS)
|
||
# 先统一列序/列集再入列(各股文件 schema 子集不同,concat 前必须对齐)
|
||
frames.append(
|
||
df.with_columns(pl.lit(vt).alias("vt_symbol"))
|
||
.select(["vt_symbol", *_DATE_COLS, *raw_cols]))
|
||
if not frames:
|
||
continue
|
||
merged = pl.concat(frames)
|
||
if table == "balance":
|
||
short = {c: c for c in raw_cols}
|
||
else:
|
||
short = {k: v for k, v in _CUM_MAP.items() if k in raw_cols}
|
||
dedup = _dedupe_reports(merged, list(raw_cols)).rename(short)
|
||
per_table[table] = dedup
|
||
|
||
empty = pl.DataFrame(schema={"vt_symbol": pl.Utf8, "REPORT_DATE": pl.Date})
|
||
if not per_table:
|
||
return empty
|
||
|
||
out = pl.concat([t.select(["vt_symbol", "REPORT_DATE"])
|
||
for t in per_table.values()]).unique()
|
||
notice_cols = []
|
||
for table, dedup in per_table.items():
|
||
renamed = dedup.rename({"notice_eff": f"_notice_{table}"})
|
||
notice_cols.append(f"_notice_{table}")
|
||
out = out.join(renamed, on=["vt_symbol", "REPORT_DATE"], how="left")
|
||
# 有效披露日 = 三表 NOTICE_DATE 行最大(最晚可见,保守不提前)
|
||
return out.with_columns(pl.max_horizontal(notice_cols).alias("notice_eff"))
|