diff --git a/scripts/data_platform/corpus_download.py b/scripts/data_platform/corpus_download.py index 8908f9c..6d7bb6c 100644 --- a/scripts/data_platform/corpus_download.py +++ b/scripts/data_platform/corpus_download.py @@ -349,9 +349,9 @@ class RecentIndex: self.keys = set() today = dt.date.today() for k in range(days + 1): - part = (CORPUS_ROOT / domain / f"dt={(today - dt.timedelta(days=k)).isoformat()}" - / "part-0.parquet") - if part.exists(): + day_dir = (CORPUS_ROOT / domain + / f"dt={(today - dt.timedelta(days=k)).isoformat()}") + for part in sorted(day_dir.glob("part-*.parquet")) if day_dir.exists() else []: df = pd.read_parquet(part, columns=ID_KEYS[domain]) for row in df.to_dict("records"): self.keys.add("|".join(str(row[c]) for c in ID_KEYS[domain])) @@ -363,37 +363,69 @@ class RecentIndex: self.keys.update(ids) -# ---------- 分区追加(append-only + 原子) ---------- +# ---------- 分区追加(多 part 文件追加写, O(1)/unit) ---------- -def append_parquet(domain, rows, id_col): - """追进今天分区 part-0.parquet; 同日同 id 只进一次; tmp+rename 原子。""" +_DAY_BUFFERS = {} + + +def _day_dir(domain): + return CORPUS_ROOT / domain / f"dt={dt.date.today().isoformat()}" + + +def append_parquet(domain, rows, id_col=None): + """进当日内存缓冲; flush_domains() 时写新 part-N.parquet(只写新行,tmp+rename)。 + + 历史: 初版整读整写当日唯一 part-0 → day-1 大流量下 O(n²)(09-07 首跑实锤 + news 段 5.3s/股劣化)。改多 part 追加: 旧数据永不重写;id 去重职责全在 + 调用方(账本+近窗);崩溃一致性=数据先落、marker 后落(Ctx.commit 次序)。 + id_col 参数保留仅为调用方兼容, 内部不再使用。 + """ if not rows: return 0 - cols = ID_KEYS[domain] if isinstance(id_col, list) else [id_col] - day = dt.date.today().isoformat() - d = CORPUS_ROOT / domain / f"dt={day}" - d.mkdir(parents=True, exist_ok=True) - part = d / "part-0.parquet" - df_new = pd.DataFrame(rows) - if part.exists(): - existing = pd.read_parquet(part) - keys = {"|".join(str(r[c]) for c in cols) - for r in existing[cols].to_dict("records")} - keep = ["|".join(str(r[c]) for c in cols) not in keys - for r in df_new[cols].to_dict("records")] - df_new = df_new[keep] - if df_new.empty: + buf = _DAY_BUFFERS.get(domain) + if buf is None: + buf = _DAY_BUFFERS[domain] = DayBuffer(domain) + buf.rows.extend(rows) + return len(rows) + + +class DayBuffer: + """当日新行内存缓冲 → flush 落一个新 part 文件(批量, 原子)。""" + + def __init__(self, domain): + self.domain = domain + self.rows = [] + d = _day_dir(domain) + idx = -1 + if d.exists(): + for p in d.glob("part-*.parquet"): + try: + idx = max(idx, int(p.stem.split("-")[1])) + except (ValueError, IndexError): + continue + self.next_idx = idx + 1 + + def flush(self): + if not self.rows: return 0 - df_new.index = range(len(existing), len(existing) + len(df_new)) - df_new = pd.concat([existing, df_new]) - tmp = d / f".part-0.{os.getpid()}.tmp" - try: - df_new.to_parquet(tmp, index=False) - os.replace(tmp, part) - except Exception: - tmp.unlink(missing_ok=True) - raise - return len(df_new) + d = _day_dir(self.domain) + d.mkdir(parents=True, exist_ok=True) + n = len(self.rows) + tmp = d / f".part-{self.next_idx}.{os.getpid()}.tmp" + try: + pd.DataFrame(self.rows).to_parquet(tmp, index=False) + os.replace(tmp, d / f"part-{self.next_idx}.parquet") + except Exception: + tmp.unlink(missing_ok=True) + raise + self.next_idx += 1 + self.rows = [] + return n + + +def flush_domains(): + for buf in _DAY_BUFFERS.values(): + buf.flush() # ---------- unit marker ---------- @@ -511,7 +543,7 @@ def _iter_partitions(domain): base = CORPUS_ROOT / domain if not base.exists(): return [] - return sorted(p for p in base.glob("dt=*/part-0.parquet") if p.is_file()) + return sorted(p for p in base.glob("dt=*/part-*.parquet") if p.is_file()) def _append_pdf_index(rows): @@ -570,6 +602,11 @@ class Ctx: self.depth = {} self.clients = {} self.t0 = time.monotonic() + self.ledgers = None + self.pending_marks = [] + + def set_stores(self, ledgers): + self.ledgers = ledgers def reset_stage(self): """--limit 预算按段独立(冒烟要覆盖每段;曾因预算共享+pdf 段无门控 @@ -579,6 +616,30 @@ class Ctx: def budget_exhausted(self): return self.limit is not None and self.stage_units >= self.limit + def unit_done(self, stage, unit, mark=True): + """unit 完成: 计数 + marker 延后入批; 每 FLUSH_EVERY unit 批量提交 + (数据→账本→marker 次序=崩溃时宁可重拉不产生洞)。""" + self.units += 1 + self.stage_units += 1 + if mark and self.limit is None: + self.pending_marks.append((self.lane, stage, unit)) + if self.units % FLUSH_EVERY == 0: + self.commit() + + def defer_mark(self, stage, unit): + if self.limit is None: + self.pending_marks.append((self.lane, stage, unit)) + + def commit(self): + """数据分区 → id 账本 → unit marker 依次落盘(原子粒度=各自 tmp+rename)。""" + flush_domains() + if self.ledgers: + for led in self.ledgers.values(): + led.flush() + for lane, stage, unit in self.pending_marks: + mark_done(lane, stage, unit) + self.pending_marks.clear() + def stop_now(self): """完成当前 unit 后调用: 过墙钟 → 抛 WallClockStop(checkpoint 语义)。""" if self.deadline and dt.datetime.now() >= self.deadline: @@ -640,14 +701,6 @@ def _save_depth(ctx): os.replace(tmp, p) -def _maybe_flush(ledgers, since_flush): - if since_flush >= FLUSH_EVERY: - for led in ledgers.values(): - led.flush() - return 0 - return since_flush - - # ---------- daily lane ---------- def _run_ann_increment(ctx, pool, cl_cn, ledgers, recents): @@ -656,7 +709,6 @@ def _run_ann_increment(ctx, pool, cl_cn, ledgers, recents): yday = today - dt.timedelta(days=1) first_seen = today.isoformat() candidates = [] - since_flush = 0 for code, org in pool: if ctx.limit is None: if is_done(ctx.lane, "ann", code): @@ -683,11 +735,7 @@ def _run_ann_increment(ctx, pool, cl_cn, ledgers, recents): _record_depth(ctx, "ann", r["sec_code"] or code, (r["ann_time"] or "")[:10]) log.info("ann %s: +%d/%d", code, len(new), len(rows)) - if ctx.limit is None: - mark_done(ctx.lane, "ann", code) - ctx.units += 1 - ctx.stage_units += 1 - since_flush = _maybe_flush(ledgers, since_flush + 1) + ctx.unit_done("ann", code, mark=ctx.limit is None) ctx.stop_now() return candidates @@ -718,10 +766,7 @@ def _run_news_increment(ctx, pool, cl_em, ledgers, recents, new_articles): _record_depth(ctx, "news", code, (r["show_time"] or "")[:10]) new_articles.append((r["art_code"], r["url"], r["show_time"])) log.info("news %s: +%d/%d", code, len(new), len(rows)) - if ctx.limit is None: - mark_done(ctx.lane, "news", code) - ctx.units += 1 - ctx.stage_units += 1 + ctx.unit_done("news", code, mark=ctx.limit is None) ctx.stop_now() @@ -742,8 +787,7 @@ def _run_fulltext(ctx, cl_em, ledgers, recents, new_articles): return except HttpDeterministicError as e: log.warning("fulltext %s 确定性失败(标done跳过): %s", art_code, e) - if ctx.limit is None: - mark_done(ctx.lane, "fulltext", art_code) + ctx.defer_mark("fulltext", art_code) seen.add(art_code) continue except TransportError as e: @@ -758,11 +802,8 @@ def _run_fulltext(ctx, cl_em, ledgers, recents, new_articles): append_parquet("news_fulltext", [row], "art_code") ledgers["news_fulltext"].add([art_code]) recents["news_fulltext"].add([art_code]) - if ctx.limit is None: - mark_done(ctx.lane, "fulltext", art_code) + ctx.unit_done("fulltext", art_code, mark=ctx.limit is None) seen.add(art_code) - ctx.units += 1 - ctx.stage_units += 1 ctx.stop_now() @@ -787,8 +828,7 @@ def _run_pdf_daily(ctx, pool, cl_cn, ledgers, candidates): log.warning("pdf fivecat %s 失败: %s", code, e) continue candidates.extend(rows) - ctx.units += 1 - ctx.stage_units += 1 + ctx.unit_done("pdf_scan", code, mark=False) ctx.stop_now() seen = set() index_rows = [] @@ -810,8 +850,7 @@ def _run_pdf_daily(ctx, pool, cl_cn, ledgers, candidates): except HttpDeterministicError as e: if e.status == 404: log.warning("pdf %s 404 missing(标done不重试): %s", aid, url) - if ctx.limit is None: - mark_done(ctx.lane, "pdf", aid) + ctx.defer_mark("pdf", aid) continue ctx.failed += 1 log.warning("pdf %s 失败: %s", aid, e) @@ -828,10 +867,7 @@ def _run_pdf_daily(ctx, pool, cl_cn, ledgers, candidates): ann_time.strftime("%Y-%m-%d") if ann_time else None), "local_path": str(dest), "bytes": size, "fetch_date": today.isoformat()}) - if ctx.limit is None: - mark_done(ctx.lane, "pdf", aid) - ctx.units += 1 - ctx.stage_units += 1 + ctx.unit_done("pdf", aid, mark=ctx.limit is None) ctx.stop_now() _append_pdf_index(index_rows) @@ -843,12 +879,11 @@ def run_daily(ctx, pool): cl_em = DomainClient("eastmoney") cl_art = DomainClient("article") ctx.clients = {"cninfo": cl_cn, "eastmoney": cl_em, "article": cl_art} + ctx.set_stores(ledgers) candidates = _run_ann_increment(ctx, pool, cl_cn, ledgers, recents) _run_news_increment(ctx, pool, cl_em, ledgers, recents, new_articles := []) _run_fulltext(ctx, cl_art, ledgers, recents, new_articles) _run_pdf_daily(ctx, pool, cl_cn, ledgers, candidates) - for led in ledgers.values(): - led.flush() # ---------- backfill lane(①ann_meta ②news_meta ③fulltext ④PDF;T5) ---------- @@ -859,7 +894,6 @@ def _run_ann_backfill(ctx, pool, cl_cn, ledgers, recents): ctx.reset_stage() today = dt.date.today() first_seen = today.isoformat() - since_flush = 0 for year in range(today.year, BACKFILL_FROM_YEAR - 1, -1): year_end = today.isoformat() if year == today.year else f"{year}-12-31" for code, org in pool: @@ -887,11 +921,7 @@ def _run_ann_backfill(ctx, pool, cl_cn, ledgers, recents): (r["ann_time"] or "")[:10]) if new: log.info("backfill ann %s: +%d/%d", unit, len(new), len(rows)) - if ctx.limit is None: - mark_done(ctx.lane, "ann", unit) - ctx.units += 1 - ctx.stage_units += 1 - since_flush = _maybe_flush(ledgers, since_flush + 1) + ctx.unit_done("ann", unit, mark=ctx.limit is None) ctx.stop_now() @@ -921,24 +951,20 @@ def _run_news_backfill(ctx, pool, cl_em, ledgers, recents): log.warning("backfill news %s 失败(断点留此页): %s", unit, e) break if not rows: - if ctx.limit is None: - mark_done(ctx.lane, "news", unit) - mark_done(ctx.lane, "news_end", code) + ctx.defer_mark("news", unit) + ctx.defer_mark("news_end", code) log.info("backfill news %s: 到头(空页@%d)", code, page) break norm = [norm_news_np_row(r, code) for r in rows] new = _absorb_new("news_meta", norm, ledgers, recents) for r in new: _record_depth(ctx, "news", code, (r["show_time"] or "")[:10]) - if ctx.limit is None: - mark_done(ctx.lane, "news", unit) - ctx.units += 1 - ctx.stage_units += 1 + ctx.unit_done("news", unit, mark=ctx.limit is None) if page == NEWS_BACKFILL_MAX_PAGE: capped = True ctx.stop_now() - if capped and ctx.limit is None: - mark_done(ctx.lane, "news_end", code) + if capped: + ctx.defer_mark("news_end", code) log.info("backfill news %s: page=%d 硬顶止", code, NEWS_BACKFILL_MAX_PAGE) @@ -952,8 +978,9 @@ def _run_fulltext_backfill(ctx, cl_art, ledgers, recents): rows_all.extend(df.to_dict("records")) if not rows_all: return - known = ledgers["news_fulltext"].has_any([str(r["art_code"]) - for r in rows_all]) + _keys = [str(r["art_code"]) for r in rows_all] + known = [l or r_ for l, r_ in zip(ledgers["news_fulltext"].has_any(_keys), + recents["news_fulltext"].has_any(_keys))] todo, seen = [], set() for r, k in zip(rows_all, known): art = str(r["art_code"]) @@ -962,7 +989,6 @@ def _run_fulltext_backfill(ctx, cl_art, ledgers, recents): seen.add(art) todo.append(r) log.info("fulltext backfill: %d 待补 / %d 库存行", len(todo), len(rows_all)) - since_flush = 0 today = dt.date.today().isoformat() for r in todo: if ctx.budget_exhausted(): @@ -988,9 +1014,7 @@ def _run_fulltext_backfill(ctx, cl_art, ledgers, recents): "art_code") ledgers["news_fulltext"].add([str(r["art_code"])]) recents["news_fulltext"].add([str(r["art_code"])]) - ctx.units += 1 - ctx.stage_units += 1 - since_flush = _maybe_flush(ledgers, since_flush + 1) + ctx.unit_done("fulltext", r["art_code"], mark=False) # 账本即进度 ctx.stop_now() @@ -1067,6 +1091,7 @@ def _run_pdf_backfill(ctx, cl_cn): log.warning("pdf %s 失败(下次重试): %s", aid, e) continue bytes_since_check += size + ctx.unit_done("pdf_bf", aid, mark=False) # 文件即进度 index_rows.append({"announcement_id": aid, "sec_code": r.get("sec_code"), "ann_type_name": _pdf_label(r["title"]), @@ -1074,8 +1099,6 @@ def _run_pdf_backfill(ctx, cl_cn): str(r["ann_time"] or "")), "local_path": str(dest), "bytes": size, "fetch_date": today}) - ctx.units += 1 - ctx.stage_units += 1 ctx.stop_now() _append_pdf_index(index_rows) @@ -1087,12 +1110,11 @@ def run_backfill(ctx, pool): cl_em = DomainClient("eastmoney") cl_art = DomainClient("article") ctx.clients = {"cninfo": cl_cn, "eastmoney": cl_em, "article": cl_art} + ctx.set_stores(ledgers) _run_ann_backfill(ctx, pool, cl_cn, ledgers, recents) _run_news_backfill(ctx, pool, cl_em, ledgers, recents) _run_fulltext_backfill(ctx, cl_art, ledgers, recents) _run_pdf_backfill(ctx, cl_cn) - for led in ledgers.values(): - led.flush() # ---------- lane 入口 ---------- @@ -1111,6 +1133,7 @@ def _ensure_log(): def run_lane(lane, until=None, limit=None): + _DAY_BUFFERS.clear() _ensure_log() ctx = Ctx(lane=lane, until=until, limit=limit) pool = load_stock_pool() @@ -1128,6 +1151,7 @@ def run_lane(lane, until=None, limit=None): except WallClockStop: log.info("墙钟到(--until %s): 完成当前 unit 后 checkpoint 退出", until) finally: + ctx.commit() _save_depth(ctx) for name, cl in ctx.clients.items(): if cl.n_requests: diff --git a/scripts/data_platform/verify_corpus.py b/scripts/data_platform/verify_corpus.py index af6a985..06f5fac 100644 --- a/scripts/data_platform/verify_corpus.py +++ b/scripts/data_platform/verify_corpus.py @@ -23,7 +23,7 @@ ID_KEYS = {"ann_meta": ["announcement_id"], def load_domain(root, domain): - parts = sorted(glob.glob(str(root / domain / "dt=*" / "part-0.parquet"))) + parts = sorted(glob.glob(str(root / domain / "dt=*" / "part-*.parquet"))) if not parts: return None, [] return pd.concat([pd.read_parquet(p, columns=ID_KEYS[domain]) for p in parts], diff --git a/tests/data_platform/test_corpus_download.py b/tests/data_platform/test_corpus_download.py index b79801c..3ffc71a 100644 --- a/tests/data_platform/test_corpus_download.py +++ b/tests/data_platform/test_corpus_download.py @@ -165,33 +165,59 @@ def test_recent_index_reads_partitions(root): # ---------- parquet 追加(append-only + 原子) ---------- -def test_append_parquet_dedup_and_atomic(root): +def test_append_parquet_buffers_then_flush_batch(root): + """多 part 语义(09-07 首跑 O(n²) 修复): append 进内存缓冲, flush 落新 part + 文件只写新行; id 去重职责在调用方(账本+近窗)。""" import pandas as pd today = dt.date.today().isoformat() - rows1 = [{"announcement_id": "a1", "title": "x"}, {"announcement_id": "a2", "title": "y"}] - cd.append_parquet("ann_meta", rows1, id_col="announcement_id") - part = root / "ann_meta" / f"dt={today}" / "part-0.parquet" - df = pd.read_parquet(part) - assert len(df) == 2 - # 同日追加含重复 id → 只进新 id - rows2 = [{"announcement_id": "a2", "title": "y2"}, {"announcement_id": "a3", "title": "z"}] - cd.append_parquet("ann_meta", rows2, id_col="announcement_id") - df = pd.read_parquet(part) - assert sorted(df["announcement_id"]) == ["a1", "a2", "a3"] - assert not list((root / "ann_meta" / f"dt={today}").glob("*.tmp")) # 无残留 tmp + cd._DAY_BUFFERS.clear() + cd.append_parquet("ann_meta", [{"announcement_id": "a1"}, {"announcement_id": "a2"}]) + d = root / "ann_meta" / f"dt={today}" + assert not list(d.glob("part-*.parquet")) # 未 flush 无文件 + cd.flush_domains() + assert sorted(p.name for p in d.glob("part-*.parquet")) == ["part-0.parquet"] + # 第二批 → 新 part-1(旧 part-0 永不重写) + cd.append_parquet("ann_meta", [{"announcement_id": "a3"}]) + cd.flush_domains() + df0 = pd.read_parquet(d / "part-0.parquet") + df1 = pd.read_parquet(d / "part-1.parquet") + assert sorted(df0["announcement_id"]) == ["a1", "a2"] and len(df1) == 1 + assert not list(d.glob("*.tmp")) def test_append_parquet_crash_leaves_no_partial(root): - """写中途抛错 → 不留 tmp/半截文件(下次重拉幂等前提)。""" + """flush 写中途抛错 → 不留 tmp/半截文件(下次重拉幂等前提)。""" import pandas as pd + cd._DAY_BUFFERS.clear() + cd.append_parquet("ann_meta", [{"announcement_id": "a1"}]) with patch.object(pd.DataFrame, "to_parquet", side_effect=RuntimeError("disk")): with pytest.raises(RuntimeError): - cd.append_parquet("ann_meta", [{"announcement_id": "a1"}], - id_col="announcement_id") + cd.flush_domains() today = dt.date.today().isoformat() assert not list((root / "ann_meta" / f"dt={today}").glob("*")) +def test_commit_batches_data_before_markers(root, monkeypatch): + """批量提交次序=数据→账本→marker: unit_done 只入批,commit 才双双落盘 + (崩溃时宁可整批重拉,不产生 marker-done-但数据缺的洞)。""" + import pandas as pd + ctx = cd.Ctx(lane="daily") + ledgers = {"ann_meta": cd.IdLedger("ann_meta")} + ctx.set_stores(ledgers) + cd._DAY_BUFFERS.clear() + cd.append_parquet("ann_meta", [{"announcement_id": "m1", "title": "x"}]) + ledgers["ann_meta"].add(["m1"]) # 真实流程 _absorb_new 负责 + ctx.unit_done("ann", "000001") + today = dt.date.today().isoformat() + d = root / "ann_meta" / f"dt={today}" + assert not d.exists() or not list(d.glob("part-*.parquet")) + assert not cd.is_done("daily", "ann", "000001") + ctx.commit() + assert (d / "part-0.parquet").exists() + assert cd.is_done("daily", "ann", "000001") + assert cd.IdLedger("ann_meta").has_any(["m1"]) == [True] + + # ---------- unit marker ---------- def test_marker_done_and_skip(root):