sanguo/sanguo_vnpy_v2
3
0
Fork 0
Code Issues 4 Pull Requests Actions Packages Projects Releases Wiki Activity

fix(data): 分区改多 part 追加写+批量提交,根治 day-1 O(n²) 劣化 [nas]
CI/CD / test (push) Successful in 25s
Details
CI/CD / nas-deploy (push) Successful in 10s
Details
CI/CD / nas-verify (push) Successful in 15s
Details

Browse Source 
09-07 首跑实锤: append_parquet 每 unit 整读整写当日唯一 part-0, news 段
13万+行后 5.3s/股(应1.15s)且随行数恶化,10:30 墙钟必跑不完且逐日更糟。
修(两层):
- DayBuffer: 当日新行内存缓冲→flush 落新 part-N.parquet(只写新行,
  旧数据永不重写);id 去重职责收敛到调用方(账本+近窗双层,原就存在)
- Ctx 批量提交: unit_done 延后 marker 入批,每 FLUSH_EVERY unit 一次
  commit=数据→账本→marker 依次落盘(崩溃=整批重拉,绝不产生
  marker-done-但数据缺的洞);run_lane 段首清缓冲/finally 兜底 commit
- ③全文回补 known 检查补近窗兜底(账本 flush 崩溃窗防重)
- readers 全改 part-* glob(RecentIndex/_iter_partitions/verify_corpus)
- 测试 54:新增 多part语义/flush崩溃一致性/批量提交次序 三钉
This commit is contained in:
Claude Dev
2026-09-07 07:12:48 +08:00
parent c5db293b46
commit efcdc8d8e0
3 changed files with 156 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/data_platform/corpus_download.py
+114 -90
View File
@@ -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:
scripts/data_platform/verify_corpus.py
+1 -1
View File
@@ -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],
tests/data_platform/test_corpus_download.py
+41 -15
View File
@@ -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):