claude_dev
e3b688354f
merge_increment/verify_increment 增量staging→验证→合并工具; raw_redownload/run_daily_update/import_vnpy_daily 强化; 补 data_platform 与 index_downloader 测试.
192 lines
7.3 KiB
Python
192 lines
7.3 KiB
Python
#!/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()
|