fix(data): avoid destructive enriched rebuild cleanup

This commit is contained in:
shy3130
2026-07-15 21:34:44 +08:00
parent 0b3b769ca8
commit 77876c7ee4
2 changed files with 92 additions and 5 deletions
+12 -5
View File
@@ -1097,8 +1097,7 @@ def run_pipeline(data_dir: Path | None = None,
SYM_BATCH = prefs_mod.get_enriched_batch_size() # 每批 N 只 × ~244 天, 可在设置中调整
total_batches = (total_syms + SYM_BATCH - 1) // SYM_BATCH
# 全量模式: 先清理旧 enriched 目录, 最后一次性按日期写入
# 收集所有批次结果, 按日期分区写入
# 全量模式: 收集所有批次结果, 最后按日期分区覆盖写入
from collections import defaultdict
date_buffers: dict[str, list[pl.DataFrame]] = defaultdict(list)
@@ -1165,9 +1164,17 @@ def run_pipeline(data_dir: Path | None = None,
# 全量模式: 按日期分区写入
if not symbols and date_buffers:
if base.exists():
import shutil
shutil.rmtree(base)
existing_dates = {
p.name.removeprefix("date=")
for p in base.glob("date=*")
if p.is_dir()
}
rebuilt_dates = set(date_buffers)
missing_dates = existing_dates - rebuilt_dates
if missing_dates:
sample = ", ".join(sorted(missing_dates)[:5])
raise RuntimeError(f"全量重建结果缺少已有日期分区,拒绝覆盖: {sample}")
base.mkdir(parents=True, exist_ok=True)
for ds, dfs in date_buffers.items():
@@ -0,0 +1,80 @@
from __future__ import annotations
from datetime import date
import polars as pl
import pytest
from app.indicators import pipeline
def _write_daily(data_dir, ds: str, close: float) -> None:
out = data_dir / "kline_daily" / f"date={ds}" / "part.parquet"
out.parent.mkdir(parents=True, exist_ok=True)
pl.DataFrame({
"symbol": ["600000.SH"],
"date": [date.fromisoformat(ds)],
"open": [close],
"high": [close],
"low": [close],
"close": [close],
"volume": [100.0],
"amount": [1000.0],
"quote_ts": [0],
}).write_parquet(out)
def _write_existing(data_dir, ds: str, close: float) -> None:
out = data_dir / "kline_daily_enriched" / f"date={ds}" / "part.parquet"
out.parent.mkdir(parents=True, exist_ok=True)
pl.DataFrame({
"symbol": ["600000.SH"],
"date": [date.fromisoformat(ds)],
"close": [close],
}).write_parquet(out)
def _fake_compute_enriched(raw: pl.DataFrame, **_kwargs) -> pl.DataFrame:
return raw.with_columns(
pl.col("close").alias("raw_close"),
pl.col("high").alias("raw_high"),
pl.col("low").alias("raw_low"),
pl.lit(None, dtype=pl.Float64).alias("turnover_rate"),
pl.lit(0, dtype=pl.UInt32).alias("consecutive_limit_ups"),
pl.lit(0, dtype=pl.UInt32).alias("consecutive_limit_downs"),
)
def test_full_rebuild_overwrites_existing_partitions_without_deleting_base(tmp_path, monkeypatch):
_write_daily(tmp_path, "2026-07-14", 14.0)
_write_daily(tmp_path, "2026-07-15", 15.0)
_write_existing(tmp_path, "2026-07-15", 1.0)
marker = tmp_path / "kline_daily_enriched" / "keep.txt"
marker.write_text("keep", encoding="utf-8")
monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
written = pipeline.run_pipeline(data_dir=tmp_path)
assert written == 2
assert marker.read_text(encoding="utf-8") == "keep"
assert pl.read_parquet(
tmp_path / "kline_daily_enriched" / "date=2026-07-14" / "part.parquet"
)["close"].to_list() == [14.0]
assert pl.read_parquet(
tmp_path / "kline_daily_enriched" / "date=2026-07-15" / "part.parquet"
)["close"].to_list() == [15.0]
def test_full_rebuild_rejects_missing_existing_dates_before_writing(tmp_path, monkeypatch):
_write_daily(tmp_path, "2026-07-15", 15.0)
_write_existing(tmp_path, "2026-07-14", 14.0)
_write_existing(tmp_path, "2026-07-15", 1.0)
monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
with pytest.raises(RuntimeError, match="缺少已有日期分区"):
pipeline.run_pipeline(data_dir=tmp_path)
existing = pl.read_parquet(
tmp_path / "kline_daily_enriched" / "date=2026-07-15" / "part.parquet"
)
assert existing["close"].to_list() == [1.0]