diff --git a/app/harness/store.py b/app/harness/store.py index 49509d9..87fba07 100644 --- a/app/harness/store.py +++ b/app/harness/store.py @@ -266,7 +266,9 @@ def read_seed(store_dir: Path, name: str) -> dict: return json.loads(seed_json.read_text()) -def extract_run_db(src_db: Path, dst_db: Path, run_id: str) -> None: +def extract_run_db( + src_db: Path, dst_db: Path, run_id: str, *, dedupe_per_question: bool = False +) -> None: """从 src_db 抽出某 run_id 的 _runs + predictions 行,写一个最小 db(种子 baseline.db)。 用源表的**原始 CREATE 语句**重建目标表,保留主键/列类型/约束—— @@ -277,6 +279,9 @@ def extract_run_db(src_db: Path, dst_db: Path, run_id: str) -> None: src_db: 源 harness.db。 dst_db: 目标 db(不得已存在)。 run_id: 要抽取的 run。 + dedupe_per_question: True 时 predictions 表每 question_id 仅保留 rowid 最小 + 的首行(对齐 canonical「每 question_id 取第一行 ORDER BY rowid」口径, + 902→900)。_runs 表不受影响。 异常: RuntimeError: 源中无该表或无该 run 的行。 @@ -294,9 +299,17 @@ def extract_run_db(src_db: Path, dst_db: Path, run_id: str) -> None: dst.execute(create_sql[0]) cols = [r[1] for r in src.execute(f"PRAGMA table_info({table})")] col_sql = ", ".join(cols) - rows = src.execute( - f"SELECT {col_sql} FROM {table} WHERE run_id=?", (run_id,) - ).fetchall() + if table == "predictions" and dedupe_per_question: + rows = src.execute( + f"SELECT {col_sql} FROM {table} WHERE run_id=? " + "AND rowid IN (SELECT MIN(rowid) FROM predictions " + "WHERE run_id=? GROUP BY question_id)", + (run_id, run_id), + ).fetchall() + else: + rows = src.execute( + f"SELECT {col_sql} FROM {table} WHERE run_id=?", (run_id,) + ).fetchall() if not rows: raise RuntimeError(f"{table} 中无 run_id={run_id} 的行") ph = ", ".join("?" * len(cols)) diff --git a/tests/unit/test_harness_store.py b/tests/unit/test_harness_store.py index 63bbaf0..721366b 100644 --- a/tests/unit/test_harness_store.py +++ b/tests/unit/test_harness_store.py @@ -338,6 +338,44 @@ class TestExtractRunDb: with pytest.raises(RuntimeError, match="无 run_id="): extract_run_db(src, dst, "nonexistent") + def test_dedupe_per_question_keeps_first_row(self, tmp_path): + """dedupe_per_question=True 时每 question_id 只保留 rowid 最小的首行。""" + import sqlite3 + + src = tmp_path / "src.db" + conn = sqlite3.connect(src) + conn.execute( + "CREATE TABLE _runs (run_id TEXT PRIMARY KEY, started_at TEXT)" + ) + conn.execute("INSERT INTO _runs VALUES ('r1', 't0')") + conn.execute( + "CREATE TABLE predictions (run_id TEXT, question_id TEXT, prediction TEXT)" + ) + # 743-1 三行(模拟 error/budget/finished),首行 prediction=NULL + conn.executemany( + "INSERT INTO predictions VALUES (?,?,?)", + [ + ("r1", "743-1", None), + ("r1", "743-1", None), + ("r1", "743-1", "C"), + ("r1", "q2", "A"), + ], + ) + conn.commit() + conn.close() + + dst = tmp_path / "dst.db" + from app.harness.store import extract_run_db + + extract_run_db(src, dst, "r1", dedupe_per_question=True) + + out = sqlite3.connect(dst) + rows = out.execute( + "SELECT question_id, prediction FROM predictions ORDER BY question_id" + ).fetchall() + out.close() + assert rows == [("743-1", None), ("q2", "A")], f"未按 rowid 首行去重: {rows}" + # --------------------------------------------------------------------------- # promote_to_seed