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

sanguo_backtest/__init__.py

@@ -0,0 +1 @@
+# Sanguo Backtest Module

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()

tests/backtest/__init__.py

@@ -0,0 +1 @@
+# Backtest tests package

tests/backtest/conftest.py

@@ -0,0 +1,20 @@
+"""Pytest fixtures for backtest tests."""
+import pytest
+import tempfile
+import shutil
+from pathlib import Path
+
+
+@pytest.fixture
+def temp_db_path(tmp_path):
+    """Create a temporary database path."""
+    db_path = tmp_path / "test_backtest.db"
+    yield str(db_path)
+
+
+@pytest.fixture
+def temp_file_dir(tmp_path):
+    """Create a temporary file directory for parquet files."""
+    file_dir = tmp_path / "parquet_files"
+    file_dir.mkdir(parents=True, exist_ok=True)
+    yield str(file_dir)

tests/backtest/test_result_store.py

@@ -0,0 +1,127 @@
+"""Tests for sanguo_backtest.result_store module."""
+import pytest
+import pandas as pd
+from datetime import datetime
+from sanguo_backtest.result_store import (
+    BacktestResult,
+    save_result,
+    load_result,
+    list_results
+)
+
+
+def test_save_and_load_result(temp_db_path):
+    """Test saving and loading a backtest result without parquet files."""
+    # Create a result without equity_curve and trades (file_dir=None by default)
+    result = BacktestResult(
+        task_id="test_task_001",
+        type="cta",
+        status="completed",
+        strategy="DualThrustStrategy",
+        symbol="IF2401",
+        params={"window": 20, "entry_threshold": 0.5},
+        start="2024-01-01",
+        end="2024-03-31",
+        statistics={
+            "total_return": 0.15,
+            "sharpe_ratio": 1.2,
+            "max_drawdown": -0.08,
+            "win_rate": 0.55
+        }
+    )
+
+    # Save the result
+    result_id = save_result(result, temp_db_path)
+    assert result_id > 0
+    assert isinstance(result_id, int)
+
+    # Load the result
+    loaded_result = load_result(result_id, temp_db_path)
+
+    # Verify all fields match
+    assert loaded_result.task_id == result.task_id
+    assert loaded_result.type == result.type
+    assert loaded_result.status == result.status
+    assert loaded_result.strategy == result.strategy
+    assert loaded_result.symbol == result.symbol
+    assert loaded_result.params == result.params
+    assert loaded_result.start == result.start
+    assert loaded_result.end == result.end
+    assert loaded_result.statistics == result.statistics
+    assert loaded_result.equity_curve is None
+    assert loaded_result.trades is None
+    assert loaded_result.error_msg is None
+
+
+def test_list_results_filter(temp_db_path):
+    """Test listing results with type filter."""
+    # Create multiple results of different types
+    cta_result = BacktestResult(
+        task_id="cta_task_001",
+        type="cta",
+        status="completed",
+        strategy="DualThrustStrategy",
+        symbol="IF2401",
+        params={"window": 20},
+        start="2024-01-01",
+        end="2024-03-31",
+        statistics={"total_return": 0.15}
+    )
+
+    factor_result = BacktestResult(
+        task_id="factor_task_001",
+        type="factor",
+        status="completed",
+        strategy="Alpha158Strategy",
+        symbol="000001.SZ",
+        params={"factors": ["alpha001", "alpha002"]},
+        start="2024-01-01",
+        end="2024-03-31",
+        statistics={"total_return": 0.20}
+    )
+
+    # Save both results
+    cta_id = save_result(cta_result, temp_db_path)
+    factor_id = save_result(factor_result, temp_db_path)
+
+    # List all results
+    all_results = list_results(db_path=temp_db_path)
+    assert len(all_results) == 2
+
+    # List only CTA results
+    cta_results = list_results(type_filter="cta", db_path=temp_db_path)
+    assert len(cta_results) == 1
+    assert cta_results[0].type == "cta"
+    assert cta_results[0].task_id == "cta_task_001"
+
+    # List only factor results
+    factor_results = list_results(type_filter="factor", db_path=temp_db_path)
+    assert len(factor_results) == 1
+    assert factor_results[0].type == "factor"
+    assert factor_results[0].task_id == "factor_task_001"
+
+
+def test_failed_result_stores_error_msg(temp_db_path):
+    """Test that failed results store error messages."""
+    # Create a failed result with error message
+    failed_result = BacktestResult(
+        task_id="failed_task_001",
+        type="cta",
+        status="failed",
+        strategy="DualThrustStrategy",
+        symbol="IF2401",
+        params={"window": 20},
+        start="2024-01-01",
+        end="2024-03-31",
+        statistics={},  # Empty statistics for failed result
+        error_msg="Data loading failed: insufficient historical data"
+    )
+
+    # Save the failed result
+    result_id = save_result(failed_result, temp_db_path)
+
+    # Load and verify error message is preserved
+    loaded_result = load_result(result_id, temp_db_path)
+    assert loaded_result.status == "failed"
+    assert loaded_result.error_msg == "Data loading failed: insufficient historical data"
+    assert loaded_result.statistics == {}