feat(backtest): result_store SQLite+parquet 结果存储

sanguo_backtest/result_store.py

@@ -0,0 +1,178 @@
+"""Backtest result storage using SQLite + parquet files."""
+import sqlite3
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+import pandas as pd
+
+
+@dataclass
+class BacktestResult:
+    """Backtest result data structure."""
+    task_id: str
+    type: str
+    status: str
+    strategy: str
+    symbol: str
+    params: dict
+    start: str
+    end: str
+    statistics: dict
+    equity_curve: Optional[pd.DataFrame] = None
+    trades: Optional[pd.DataFrame] = None
+    error_msg: Optional[str] = None
+
+
+# SQLite schema for backtest stats
+_SCHEMA = """CREATE TABLE IF NOT EXISTS backtest_stats (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    task_id TEXT,
+    type TEXT,
+    status TEXT,
+    strategy TEXT,
+    symbol TEXT,
+    params TEXT,
+    start TEXT,
+    end TEXT,
+    statistics TEXT,
+    equity_path TEXT,
+    trades_path TEXT,
+    error_msg TEXT,
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+);"""
+
+
+def _connect(db_path: str) -> sqlite3.Connection:
+    """Create database connection and initialize schema."""
+    conn = sqlite3.connect(db_path)
+    conn.executescript(_SCHEMA)
+    return conn
+
+
+def save_result(result: BacktestResult, db_path: str, file_dir: Optional[str] = None) -> int:
+    """
+    Save backtest result to database and optionally save DataFrame fields to parquet files.
+
+    Args:
+        result: BacktestResult object to save
+        db_path: Path to SQLite database file
+        file_dir: Optional directory to save parquet files (equity_curve and trades)
+
+    Returns:
+        int: The ID of the inserted record
+    """
+    conn = _connect(db_path)
+    try:
+        equity_path = None
+        trades_path = None
+
+        # Save DataFrames to parquet if file_dir is provided
+        if file_dir:
+            fdir = Path(file_dir)
+            fdir.mkdir(parents=True, exist_ok=True)
+
+            if result.equity_curve is not None and not result.equity_curve.empty:
+                equity_path = str(fdir / f"{result.task_id}_equity.parquet")
+                result.equity_curve.to_parquet(equity_path)
+
+            if result.trades is not None and not result.trades.empty:
+                trades_path = str(fdir / f"{result.task_id}_trades.parquet")
+                result.trades.to_parquet(trades_path)
+
+        # Insert record into database
+        cur = conn.execute(
+            """INSERT INTO backtest_stats
+               (task_id, type, status, strategy, symbol, params, start, end,
+                statistics, equity_path, trades_path, error_msg)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+            (
+                result.task_id,
+                result.type,
+                result.status,
+                result.strategy,
+                result.symbol,
+                json.dumps(result.params),
+                result.start,
+                result.end,
+                json.dumps(result.statistics),
+                equity_path,
+                trades_path,
+                result.error_msg
+            )
+        )
+        conn.commit()
+        return cur.lastrowid
+    finally:
+        conn.close()
+
+
+def load_result(rid: int, db_path: str) -> BacktestResult:
+    """
+    Load backtest result by ID from database.
+
+    Args:
+        rid: Result ID to load
+        db_path: Path to SQLite database file
+
+    Returns:
+        BacktestResult: Loaded result object
+
+    Raises:
+        KeyError: If result ID not found
+    """
+    conn = _connect(db_path)
+    try:
+        row = conn.execute("SELECT * FROM backtest_stats WHERE id=?", (rid,)).fetchone()
+        if not row:
+            raise KeyError(f"result {rid} not found")
+
+        # Get column names from table description
+        cols = [d[0] for d in conn.execute("SELECT * FROM backtest_stats LIMIT 0").description]
+        d = dict(zip(cols, row))
+
+        # Load parquet files if paths exist
+        equity = pd.read_parquet(d["equity_path"]) if d.get("equity_path") else None
+        trades = pd.read_parquet(d["trades_path"]) if d.get("trades_path") else None
+
+        return BacktestResult(
+            task_id=d["task_id"],
+            type=d["type"],
+            status=d["status"],
+            strategy=d["strategy"],
+            symbol=d["symbol"],
+            params=json.loads(d["params"]),
+            start=d["start"],
+            end=d["end"],
+            statistics=json.loads(d["statistics"]) if d["statistics"] else {},
+            equity_curve=equity,
+            trades=trades,
+            error_msg=d.get("error_msg")
+        )
+    finally:
+        conn.close()
+
+
+def list_results(type_filter: Optional[str] = None, db_path: str = "") -> list[BacktestResult]:
+    """
+    List all backtest results, optionally filtered by type.
+
+    Args:
+        type_filter: Optional filter for result type (e.g., 'cta', 'factor')
+        db_path: Path to SQLite database file
+
+    Returns:
+        list[BacktestResult]: List of loaded result objects
+    """
+    conn = _connect(db_path)
+    try:
+        query = "SELECT id FROM backtest_stats"
+        args = ()
+
+        if type_filter:
+            query = query + " WHERE type=?"
+            args = (type_filter,)
+
+        return [load_result(row[0], db_path) for row in conn.execute(query, args).fetchall()]
+    finally:
+        conn.close()