fix(data): data_platform硬化(增量merge/verify+raw_redownload/run_daily_update)+测试
merge_increment/verify_increment 增量staging→验证→合并工具; raw_redownload/run_daily_update/import_vnpy_daily 强化; 补 data_platform 与 index_downloader 测试.
This commit is contained in:
@@ -0,0 +1,191 @@
|
||||
#!/usr/bin/env python3
|
||||
"""合并 staging 增量到主库（task: 截断 bug 修复）。
|
||||
|
||||
**核心不变量（截断 bug 回归测试核心）**：
|
||||
合并后 main 的行数 >= 合并前 main 的行数（绝不截断）。
|
||||
|
||||
raw_redownload.py 的 save_one() 用"本次拉的几天"覆盖写整年文件 → 整年数据被截
|
||||
（NAS 2026 从 121 行截成 4 行就是这 bug）。本脚本负责按 symbol-year 安全合并：
|
||||
pd.concat([main, staging]).drop_duplicates(subset=['date'], keep='last')
|
||||
staging 的新数据/修订值优先（keep='last'），main 已有的历史保留。
|
||||
|
||||
用法：
|
||||
python3 merge_increment.py --staging data_cache/daily_update/raw --main data_cache/raw
|
||||
python3 merge_increment.py --staging ... --main ... --dry-run
|
||||
|
||||
约定：
|
||||
- staging 与 main 同构：`{root}/{year}/{sh|sz}{code}_daily.parquet`
|
||||
- staging 文件不修改/不删除（保留供排查），只写 main
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pandas as pd
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger("merge_inc")
|
||||
|
||||
# drop_duplicates 的去重键——日线按 date 唯一
|
||||
DEDUP_KEY = ["date"]
|
||||
|
||||
|
||||
def symbol_from_filename(fname: str) -> str:
|
||||
"""`sh600000_daily.parquet` → `600000`（剥 exchange 前缀和 _daily.parquet 后缀）。"""
|
||||
base = fname
|
||||
if base.endswith("_daily.parquet"):
|
||||
base = base[: -len("_daily.parquet")]
|
||||
if base[:2] in ("sh", "sz", "bj"):
|
||||
base = base[2:]
|
||||
return base
|
||||
|
||||
|
||||
def merge_one(staging_file: str, main_file: str, dry_run: bool = False) -> dict:
|
||||
"""合并一个 staging parquet 到 main。
|
||||
|
||||
返回 stats: {symbol, before, after, added, action}。
|
||||
若 main 不存在 → action="created"（直接搬 staging 过去）。
|
||||
若 main 存在 → action="merged"，按 DEDUP_KEY 去重，staging 值优先（keep='last'）。
|
||||
"""
|
||||
sym = symbol_from_filename(os.path.basename(staging_file))
|
||||
staging_df = pd.read_parquet(staging_file)
|
||||
|
||||
if not os.path.exists(main_file):
|
||||
if not dry_run:
|
||||
os.makedirs(os.path.dirname(main_file), exist_ok=True)
|
||||
staging_df.to_parquet(main_file, index=False)
|
||||
return {"symbol": sym, "before": 0, "after": len(staging_df),
|
||||
"added": len(staging_df), "action": "created"}
|
||||
|
||||
main_df = pd.read_parquet(main_file)
|
||||
before = len(main_df)
|
||||
|
||||
# concat → drop_duplicates(keep='last' 让 staging 的新/修订值覆盖 main) → sort
|
||||
merged = pd.concat([main_df, staging_df], ignore_index=True)
|
||||
merged = merged.drop_duplicates(subset=DEDUP_KEY, keep="last")
|
||||
merged = merged.sort_values(by="date").reset_index(drop=True)
|
||||
after = len(merged)
|
||||
|
||||
# 核心不变量：绝不能让 main 变少
|
||||
if after < before:
|
||||
raise RuntimeError(
|
||||
f"INVARIANT VIOLATED: {main_file} {before}→{after} (staging={staging_file})"
|
||||
)
|
||||
|
||||
if not dry_run:
|
||||
merged.to_parquet(main_file, index=False)
|
||||
|
||||
return {"symbol": sym, "before": before, "after": after,
|
||||
"added": after - before, "action": "merged"}
|
||||
|
||||
|
||||
def walk_staging(staging_root: str) -> list[tuple[str, str]]:
|
||||
"""收集 staging 下所有 `{year}/{sym}_daily.parquet`，返回 [(staging_file, main_file_relpath)]。
|
||||
|
||||
main_file_relpath 是相对 staging_root 的路径（如 `2026/sh600000_daily.parquet`），
|
||||
拼到 main_root 即得 main 文件全路径，保持两边同构。
|
||||
"""
|
||||
pairs: list[tuple[str, str]] = []
|
||||
for year in sorted(os.listdir(staging_root)):
|
||||
ydir = os.path.join(staging_root, year)
|
||||
if not os.path.isdir(ydir):
|
||||
continue
|
||||
for fname in sorted(os.listdir(ydir)):
|
||||
if not fname.endswith("_daily.parquet"):
|
||||
continue
|
||||
rel = os.path.join(year, fname)
|
||||
pairs.append((os.path.join(ydir, fname), rel))
|
||||
return pairs
|
||||
|
||||
|
||||
def run_merge(staging_root: str, main_root: str, dry_run: bool = False) -> dict:
|
||||
"""合并 staging → main，返回汇总统计。"""
|
||||
if not os.path.isdir(staging_root):
|
||||
raise FileNotFoundError(f"staging dir 不存在: {staging_root}")
|
||||
|
||||
pairs = walk_staging(staging_root)
|
||||
if not pairs:
|
||||
log.warning("staging 无 parquet: %s", staging_root)
|
||||
return {"merged": 0, "created": 0, "skipped": 0, "total_new_rows": 0,
|
||||
"details": [], "dry_run": dry_run, "ok": True,
|
||||
"invariant_violations": []}
|
||||
|
||||
os.makedirs(main_root, exist_ok=True)
|
||||
details: list[dict] = []
|
||||
merged_n = created_n = skipped_n = total_new = 0
|
||||
invariant_violations: list[str] = []
|
||||
|
||||
for i, (staging_file, rel) in enumerate(pairs, 1):
|
||||
main_file = os.path.join(main_root, rel)
|
||||
try:
|
||||
stat = merge_one(staging_file, main_file, dry_run=dry_run)
|
||||
except RuntimeError as e:
|
||||
# 不变式违反：立刻停（绝不能继续写入更小的 main）
|
||||
invariant_violations.append(str(e))
|
||||
log.error("[%d/%d] INVARIANT %s", i, len(pairs), e)
|
||||
continue
|
||||
except Exception as e: # noqa: BLE001
|
||||
log.error("[%d/%d] %s ERR %s", i, len(pairs), rel, e)
|
||||
skipped_n += 1
|
||||
continue
|
||||
|
||||
details.append(stat)
|
||||
if stat["action"] == "created":
|
||||
created_n += 1
|
||||
else:
|
||||
merged_n += 1
|
||||
total_new += stat["added"]
|
||||
|
||||
if i % 1000 == 0 or i == len(pairs):
|
||||
log.info("[%d/%d] %s %s +%d (main %d→%d)",
|
||||
i, len(pairs), stat["symbol"], stat["action"],
|
||||
stat["added"], stat["before"], stat["after"])
|
||||
|
||||
summary = {
|
||||
"merged": merged_n,
|
||||
"created": created_n,
|
||||
"skipped": skipped_n,
|
||||
"total_new_rows": total_new,
|
||||
"details": details,
|
||||
"dry_run": dry_run,
|
||||
"invariant_violations": invariant_violations,
|
||||
}
|
||||
if invariant_violations:
|
||||
# 不变式违反致命——即使个别合并成功也判失败
|
||||
summary["ok"] = False
|
||||
else:
|
||||
summary["ok"] = True
|
||||
return summary
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="合并 staging 增量到主库（绝不截断）")
|
||||
ap.add_argument("--staging", required=True, help="staging 根（如 data_cache/daily_update/raw）")
|
||||
ap.add_argument("--main", required=True, help="主库根（如 data_cache/raw）")
|
||||
ap.add_argument("--dry-run", action="store_true", help="只报不写")
|
||||
ap.add_argument("--summary-json", default=None, help="把汇总写到该 JSON 文件")
|
||||
args = ap.parse_args()
|
||||
|
||||
summary = run_merge(args.staging, args.main, dry_run=args.dry_run)
|
||||
|
||||
mode = "[DRY-RUN] " if args.dry_run else ""
|
||||
log.info("=== %s合并完成: merged=%d created=%d skipped=%d 新增行=%d ok=%s ===",
|
||||
mode, summary["merged"], summary["created"], summary["skipped"],
|
||||
summary["total_new_rows"], summary["ok"])
|
||||
if summary["invariant_violations"]:
|
||||
log.error("不变式违反 %d 条（main 被截），详情见上", len(summary["invariant_violations"]))
|
||||
|
||||
if args.summary_json:
|
||||
with open(args.summary_json, "w") as f:
|
||||
json.dump(summary, f, ensure_ascii=False, indent=2, default=str)
|
||||
|
||||
# 不变式违反 → 退 1（即便部分成功，也提示 main 可能已损坏需排查）
|
||||
sys.exit(0 if summary["ok"] else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user