用 DuckDB 搭建本地优先的股票分析引擎——一个能卖的微型 SaaS

import duckdb
from pathlib import Path

class StockEngine:
    def __init__(self, db_path="stocks.duckdb"):
        self.con = duckdb.connect(db_path)
        self._init_schema()
    
    def _init_schema(self):
        self.con.execute("""
            CREATE TABLE IF NOT EXISTS kline (
                date DATE,
                code VARCHAR,
                open DOUBLE,
                high DOUBLE,
                low DOUBLE,
                close DOUBLE,
                volume BIGINT,
                turnover DOUBLE,
                industry VARCHAR
            )
        """)
        
    def load_csv(self, csv_path: str, if_exists: str = "replace"):
        """直接从 CSV 文件加载数据到 DuckDB"""
        if if_exists == "replace":
            self.con.execute("DROP TABLE IF EXISTS kline")
        self.con.execute(f"""
            COPY kline FROM '{csv_path}'
            (FORMAT CSV, HEADER, DELIMITER ',')
        """)
        print(f"✅ 数据加载完成")

    def calculate_ma(self, short_window: int = 20, long_window: int = 60):
        """计算多条均线，使用窗口函数实现"""
        sql = f"""
            SELECT 
                date,
                code,
                close,
                AVG(close) OVER (
                    PARTITION BY code 
                    ORDER BY date 
                    ROWS BETWEEN {short_window - 1} PRECEDING AND CURRENT ROW
                ) AS ma_short,
                AVG(close) OVER (
                    PARTITION BY code 
                    ORDER BY date 
                    ROWS BETWEEN {long_window - 1} PRECEDING AND CURRENT ROW
                ) AS ma_long,
                volume
            FROM kline
            ORDER BY code, date
        """
        return self.con.sql(sql).fetchdf()

    def generate_signals(self, short_window=20, long_window=60):
        """生成金叉/死叉交易信号"""
        df = self.calculate_ma(short_window, long_window)
        
        # 保存中间结果供后续 SQL 使用
        self.con.execute("DROP TABLE IF EXISTS ma")
        df.to_sql("ma", self.con)
        
        signals = self.con.sql("""
            WITH ma_with_lag AS (
                SELECT 
                    date,
                    code,
                    ma_short,
                    ma_long,
                    LAG(ma_short) OVER (PARTITION BY code ORDER BY date) AS prev_ma_short,
                    LAG(ma_long) OVER (PARTITION BY code ORDER BY date) AS prev_ma_long
                FROM ma
                WHERE ma_short IS NOT NULL AND ma_long IS NOT NULL
            )
            SELECT 
                date,
                code,
                CASE 
                    WHEN ma_short > ma_long AND prev_ma_short <= prev_ma_long THEN 'BUY'
                    WHEN ma_short < ma_long AND prev_ma_short >= prev_ma_long THEN 'SELL'
                    ELSE 'HOLD'
                END AS signal
            FROM ma_with_lag
        """).fetchdf()
        
        return signals

    def industry_performance(self, start_date=None, end_date=None):
        """按行业分组统计年化收益率和最大回撤"""
        sql = """
            WITH daily_returns AS (
                SELECT 
                    date,
                    code,
                    industry,
                    close / LAG(close) OVER (PARTITION BY code ORDER BY date) - 1 AS daily_ret
                FROM kline
            ),
            industry_daily AS (
                SELECT 
                    date,
                    industry,
                    AVG(daily_ret) AS ind_ret,
                    COUNT(*) AS stock_count
                FROM daily_returns
                WHERE daily_ret IS NOT NULL
                GROUP BY date, industry
            )
            SELECT 
                industry,
                AVG(ind_ret) * 252 AS annual_return,
                STDDEV(ind_ret) * SQRT(252) AS annual_volatility,
                (AVG(ind_ret) * 252) / (STDDEV(ind_ret) * SQRT(252)) AS sharpe_ratio
            FROM industry_daily
            GROUP BY industry
            ORDER BY sharpe_ratio DESC
        """
        return self.con.sql(sql).fetchdf()

    def backtest(self, strategy='ma_cross', params=None):
        """简单回测引擎"""
        signals = self.generate_signals(**(params or {}))
        signals.to_sql("signals", self.con, if_exists="replace")
        
        result = self.con.sql("""
            WITH signals_with_close AS (
                SELECT 
                    s.date,
                    s.code,
                    s.signal,
                    k.close,
                    k.industry
                FROM signals s
                JOIN kline k ON s.code = k.code AND s.date = k.date
            ),
            positions AS (
                SELECT 
                    *,
                    SUM(CASE WHEN signal = 'BUY' THEN 1 
                             WHEN signal = 'SELL' THEN -1 
                             ELSE 0 END) 
                        OVER (PARTITION BY code ORDER BY date) AS position
                FROM signals_with_close
            )
            SELECT 
                code,
                industry,
                AVG(CASE WHEN signal != 'HOLD' THEN ret END) AS avg_return_on_signal,
                COUNT(CASE WHEN signal != 'HOLD' THEN 1 END) AS trade_count
            FROM (
                SELECT *,
                    close / LAG(close) OVER (PARTITION BY code ORDER BY date) - 1 AS ret
                FROM positions
            ) sub
            GROUP BY code, industry
        """).fetchdf()
        
        return result