Python+PostgreSQL：滚动/扩展窗口计算与多线程的数据调用最优方案

Hey there, let's break down how to tackle this efficiently—since you're dealing with 100 stocks (5050 ordered pairs) and millions of rows, optimizing both data retrieval and computation is key. Here's a structured approach:

1. Start with Smart Data Retrieval #

First, minimize the data you pull from PostgreSQL. Fetch only the columns you need (timestamp + 100 stock metrics) and ensure your timestamp column is indexed (critical for fast ordering and window operations).

Use SQLAlchemy to load the data directly into a pandas DataFrame—this is far more efficient than row-by-row fetching:

from sqlalchemy import create_engine
import pandas as pd

# Initialize your connection (replace with your credentials)
engine = create_engine('postgresql+psycopg2://user:pass@host:port/db')

# Fetch aligned time series data
query = """
SELECT timestamp, stock_1, stock_2, ..., stock_100
FROM stock_data
ORDER BY timestamp ASC;
"""

# Load into DataFrame, set timestamp as index
df = pd.read_sql_query(query, engine, index_col='timestamp')

# Clean missing values (adjust based on your needs: drop, fill, etc.)
df = df.dropna(subset=df.columns)

2. Optimize Regression Calculations with Covariance Matrices #

Running 5050 separate regressions per window is slow. Instead, precompute rolling/expanding covariance and variance matrices—this lets you derive all pairwise betas in vectorized operations (way faster than OLS per pair).

Rolling Windows Example (e.g., 60-day window) #

window_size = 60

# Compute rolling covariance matrix for all stocks
rolling_cov = df.rolling(window=window_size).cov()

# Extract rolling variances (diagonal of the covariance matrix)
rolling_var = df.rolling(window=window_size).var()

# Calculate betas for all ordered pairs
beta_results = {}
for stock_x in df.columns:
    for stock_y in df.columns:
        # Beta of y regressed on x: cov(y,x)/var(x)
        beta_y_on_x = rolling_cov.loc[:, stock_y, stock_x] / rolling_var.loc[:, stock_x]
        # Beta of x regressed on y: cov(x,y)/var(y)
        beta_x_on_y = rolling_cov.loc[:, stock_x, stock_y] / rolling_var.loc[:, stock_y]
        
        beta_results[(stock_y, stock_x)] = beta_y_on_x
        beta_results[(stock_x, stock_y)] = beta_x_on_y

Expanding Windows Example #

Swap out rolling for expanding:

expanding_cov = df.expanding().cov()
expanding_var = df.expanding().var()

# Same beta calculation logic as above, using expanding_cov/expanding_var

3. Multithreading/Multiprocessing When Needed #

If you need full OLS regressions (e.g., with intercepts or additional controls) instead of just betas, vectorized covariance won't cut it. For this, use multiprocessing (to bypass Python's GIL for CPU-bound work) to parallelize pair processing.

Parallel OLS Regression Example #

import numpy as np
from concurrent.futures import ProcessPoolExecutor

# Generate all ordered pairs (y, x) for regression
all_pairs = [(y_col, x_col) for y_col in df.columns for x_col in df.columns]

def run_ols_regression(pair, window_type="rolling", window_size=60):
    y_col, x_col = pair
    y = df[y_col]
    x = df[x_col]
    
    # Define window logic
    if window_type == "rolling":
        window = x.rolling(window=window_size)
    elif window_type == "expanding":
        window = x.expanding()
    else:
        raise ValueError("Window type must be 'rolling' or 'expanding'")
    
    # Helper to run OLS on a single window
    def ols_window(window_idx):
        # Align x and y for the window
        window_data = pd.concat([x.loc[window_idx], y.loc[window_idx]], axis=1).dropna()
        if len(window_data) < 2:
            return np.nan
        
        # Add intercept term
        X = np.column_stack((np.ones(len(window_data)), window_data[x_col]))
        # Solve for coefficients
        coeffs = np.linalg.lstsq(X, window_data[y_col], rcond=None)[0]
        return coeffs[1]  # Return slope (beta)
    
    # Apply OLS to each window
    beta_series = window.apply(lambda w: ols_window(w.index), raw=False)
    return (pair, beta_series)

# Run in parallel (adjust num_workers based on your CPU cores)
num_workers = 4
with ProcessPoolExecutor(max_workers=num_workers) as executor:
    futures = [executor.submit(run_ols_regression, pair) for pair in all_pairs]
    results = [future.result() for future in futures]

# Convert results to a dictionary for easy access
beta_results = {pair: series for pair, series in results}

4. Critical Optimizations for Large Datasets #

• Chunked Processing: If your data won't fit in memory, process it in time-based chunks. Fetch a chunk, compute regressions, write results back to PostgreSQL, then move to the next chunk.
• DB-Side Precomputation: Use PostgreSQL's window functions to calculate rolling means, variances, and covariances directly in the DB—this reduces the amount of data you need to transfer to Python.
• Cache Intermediate Data: Save fetched stock time series to a parquet file (using df.to_parquet()) if you're running multiple iterations, to avoid re-fetching from the DB every time.
• Indexing: Double-check that your timestamp column in PostgreSQL has an index—this will drastically speed up ordering and window queries.

5. Key Tradeoffs to Keep in Mind #

• Covariance Matrix vs. OLS: The covariance approach is way faster but only gives you beta coefficients. Use full OLS only if you need intercepts, p-values, or other stats.
• Multiprocessing Overhead: Multiprocessing adds overhead for data serialization. It's worth it for CPU-heavy tasks, but the covariance matrix method is more efficient if it meets your needs.
• DB Load: Avoid spamming concurrent DB connections. Reuse your SQLAlchemy engine connection for all fetch operations.

内容的提问来源于stack exchange，提问作者eternity1