Files
sanguo_vnpy_v2/scripts/factor_research/preaggregate_top_holders.py
T
claude_dev 6e719de0c6 feat(factor): top_holders预聚合一次性脚本——108,610小文件→factor_cache聚合产物 [nas]
E13 前置: adapter 逐文件读全市场不可行,先聚合成单 parquet
(file_code/period/hold_pct/n_holders;null占比按0,坏文件名跳过,幂等--force)。
NAS 真跑已验证: 449.7s/108,610文件 → 97,836 行(skipped=4370 empty=6404)。
2026-09-09 12:12:00 +08:00

112 lines
4.5 KiB
Python

#!/usr/bin/env python
# scripts/factor_research/preaggregate_top_holders.py
"""top_holders 十大流通股东 → (code, 期) 合计占比一次性预聚合.
背景(P1-B 批 E13): static/top_holders/ 按股×期 108,610 个小文件
(每文件 10 行),adapter 逐文件读全市场不可行——先聚合成单 parquet,
adapter 只读聚合产物(期→法定披露截止的 PIT 映射在 adapter 做)。
契约:
- 输入文件名 {code}.{SH|SZ}_{YYYYMMDD}_top_holders.parquet(北交无文件,
非契约文件名/坏文件跳过并计数)
- hold_pct = 单文件「占总流通股本持股比例」(%,54.07=54%)求和,null 按 0
- 输出列: file_code("600519.SH") / period(Date) / hold_pct(Float64) /
n_holders(原始行数,sanity 用)
- 幂等: dst 存在默认跳过返回既有内容(--force 重算)
- 绝不写 static 树,输出到独立 factor_cache 目录
用法(NAS 容器内真跑):
python scripts/factor_research/preaggregate_top_holders.py \
--src /volume1/stock/sanguo_vnpy_v2/data/static/top_holders \
--dst /volume1/stock/sanguo_vnpy_v2/data/factor_cache/top_holders_agg.parquet
预计耗时 ~5-15 分钟(108,610 个小文件顺序读)。
"""
from __future__ import annotations
import argparse
import os
import re
import sys
import time
from datetime import datetime
import polars as pl
DEFAULT_SRC = "/volume1/stock/sanguo_vnpy_v2/data/static/top_holders"
DEFAULT_DST = "/volume1/stock/sanguo_vnpy_v2/data/factor_cache/top_holders_agg.parquet"
_RATIO_COL = "占总流通股本持股比例"
# 契约文件名: 6 位代码 + .SH/.SZ + 8 位期 + _top_holders.parquet(北交排除)
_FNAME_RE = re.compile(r"^(\d{6}\.(?:SH|SZ))_(\d{8})_top_holders\.parquet$")
_LOG_EVERY = 10000
def _parse_filename(fname: str) -> tuple[str, object] | tuple[None, None]:
"""文件名 → (file_code, period);非契约名返回 (None, None)."""
m = _FNAME_RE.match(fname)
if not m:
return None, None
try:
period = datetime.strptime(m.group(2), "%Y%m%d").date()
except ValueError:
return None, None
return m.group(1), period
def aggregate(src: str, dst: str, force: bool = False,
log_every: int = _LOG_EVERY) -> pl.DataFrame:
"""扫描 src 全量文件按 (file_code, period) 求和占比,写出/返回聚合表."""
if os.path.exists(dst) and not force:
print(f"[preaggregate_top_holders] dst 已存在,跳过重算(--force 重算): {dst}")
return pl.read_parquet(dst)
if not os.path.isdir(src):
raise SystemExit(f"[preaggregate_top_holders] src 不存在: {src}")
files = sorted(f for f in os.listdir(src) if f.endswith(".parquet"))
rows: list[dict] = []
skipped = empty = 0
t0 = time.time()
for i, fname in enumerate(files, 1):
file_code, period = _parse_filename(fname)
if file_code is None:
skipped += 1
continue
try:
df = pl.read_parquet(os.path.join(src, fname), columns=[_RATIO_COL])
except Exception:
skipped += 1
continue
if df.height == 0:
empty += 1
continue
# null 占比按 0(缺披露的席位不贡献,也不清零整期)
hold_pct = df[_RATIO_COL].cast(pl.Float64, strict=False).fill_null(0.0).sum()
rows.append({"file_code": file_code, "period": period,
"hold_pct": float(hold_pct), "n_holders": df.height})
if i % log_every == 0:
el = time.time() - t0
print(f"[preaggregate_top_holders] {i}/{len(files)} "
f"({el:.0f}s, 均 {el / i * 1000:.1f}ms/文件)", file=sys.stderr)
if not rows:
raise SystemExit(f"[preaggregate_top_holders] src 无可用文件: {src}")
out = pl.DataFrame(rows).sort(["file_code", "period"])
os.makedirs(os.path.dirname(dst) or ".", exist_ok=True)
out.write_parquet(dst)
print(f"[preaggregate_top_holders] rows={out.height} files={len(files)} "
f"skipped={skipped} empty={empty}{dst} "
f"({time.time() - t0:.1f}s)")
return out
def main(argv: list[str] | None = None) -> None:
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
ap.add_argument("--src", default=DEFAULT_SRC, help="top_holders 原始目录")
ap.add_argument("--dst", default=DEFAULT_DST, help="聚合产物 parquet 路径")
ap.add_argument("--force", action="store_true", help="dst 存在也重算")
args = ap.parse_args(argv)
aggregate(args.src, args.dst, force=args.force)
if __name__ == "__main__":
main()