|
|
"""Point-in-time snapshot data loader for the backtester.
|
|
|
|
|
|
The public interface is:
|
|
|
store = SnapshotStore.load(snapshot_dir, split_name, oracle_url, db_dsn)
|
|
|
rows = store.get_candidates_for_date(date)
|
|
|
bar = store.get_bar(symbol, date)
|
|
|
macro = store.get_macro_for_date(date)
|
|
|
|
|
|
For tests, inject data directly:
|
|
|
store = SnapshotStore(candidates_by_exec_date=..., bars_by_symbol_date=..., macro_by_date=...)
|
|
|
"""
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import asyncio
|
|
|
import datetime as dt
|
|
|
from pathlib import Path
|
|
|
from typing import Any
|
|
|
|
|
|
import pyarrow.parquet as pq
|
|
|
|
|
|
from libs.common.logging import get_logger
|
|
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
|
|
|
class SnapshotStore:
|
|
|
"""In-memory cache of point-in-time backtest data."""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]],
|
|
|
bars_by_symbol_date: dict[str, dict[dt.date, dict[str, Any]]],
|
|
|
macro_by_date: dict[dt.date, dict[str, Any]] | None = None,
|
|
|
) -> None:
|
|
|
self._candidates = candidates_by_exec_date
|
|
|
self._bars = bars_by_symbol_date
|
|
|
self._macro = macro_by_date or {}
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
|
# Public query interface
|
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
def get_candidates_for_date(self, date: dt.date) -> list[dict[str, Any]]:
|
|
|
"""Return candidates where execution_date == date. No look-ahead."""
|
|
|
return list(self._candidates.get(date, []))
|
|
|
|
|
|
def get_bar(self, symbol: str, date: dt.date) -> dict[str, Any] | None:
|
|
|
"""Return OHLCV bar for symbol on date, or None if missing."""
|
|
|
sym_bars = self._bars.get(symbol)
|
|
|
if sym_bars is None:
|
|
|
return None
|
|
|
return sym_bars.get(date)
|
|
|
|
|
|
def get_macro_for_date(self, date: dt.date) -> dict[str, Any]:
|
|
|
"""Return macro observations for date (empty dict if none)."""
|
|
|
return dict(self._macro.get(date, {}))
|
|
|
|
|
|
def all_execution_dates(self) -> list[dt.date]:
|
|
|
"""Sorted list of dates that have at least one candidate."""
|
|
|
return sorted(self._candidates.keys())
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
|
# Factory: load from Parquet + DB + Oracle
|
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
@classmethod
|
|
|
def load(
|
|
|
cls,
|
|
|
snapshot_dir: str | Path,
|
|
|
split_name: str,
|
|
|
oracle_url: str,
|
|
|
db_dsn: str,
|
|
|
) -> "SnapshotStore":
|
|
|
"""Synchronous factory. Internally uses asyncio.run() to prefetch data.
|
|
|
|
|
|
Raises RuntimeError if called from within a running event loop.
|
|
|
"""
|
|
|
try:
|
|
|
loop = asyncio.get_running_loop()
|
|
|
except RuntimeError:
|
|
|
loop = None
|
|
|
|
|
|
if loop is not None and loop.is_running():
|
|
|
raise RuntimeError(
|
|
|
"SnapshotStore.load() cannot be called from a running event loop. "
|
|
|
"Use SnapshotStore._async_load() directly in async contexts."
|
|
|
)
|
|
|
|
|
|
data = asyncio.run(cls._async_load(Path(snapshot_dir), split_name, oracle_url, db_dsn))
|
|
|
return cls(**data)
|
|
|
|
|
|
@classmethod
|
|
|
async def _async_load(
|
|
|
cls,
|
|
|
snapshot_dir: Path,
|
|
|
split_name: str,
|
|
|
oracle_url: str,
|
|
|
db_dsn: str,
|
|
|
) -> dict[str, Any]:
|
|
|
"""Async data loading pipeline. Returns kwargs dict for __init__."""
|
|
|
parquet_path = snapshot_dir / f"{split_name}.parquet"
|
|
|
if not parquet_path.exists():
|
|
|
raise FileNotFoundError(f"Parquet file not found: {parquet_path}")
|
|
|
|
|
|
# Step 1: Read Parquet
|
|
|
logger.info("snapshot_store_reading_parquet", path=str(parquet_path))
|
|
|
table = pq.read_table(str(parquet_path))
|
|
|
rows: list[dict[str, Any]] = table.to_pydict()
|
|
|
# Convert column-oriented dict to list of row dicts
|
|
|
num_rows = table.num_rows
|
|
|
col_names = list(rows.keys())
|
|
|
row_list: list[dict[str, Any]] = [
|
|
|
{col: rows[col][i] for col in col_names} for i in range(num_rows)
|
|
|
]
|
|
|
logger.info("snapshot_store_rows_loaded", count=num_rows)
|
|
|
|
|
|
# Collect event_ids for DB lookup
|
|
|
event_ids = [str(r.get("event_id", "")) for r in row_list]
|
|
|
|
|
|
# Step 2–6: DB + Oracle enrichment
|
|
|
event_meta = await cls._fetch_event_metadata(event_ids, db_dsn)
|
|
|
unique_symbols = list({m.get("ticker", "") for m in event_meta.values() if m.get("ticker")})
|
|
|
|
|
|
date_range = cls._compute_date_range(row_list)
|
|
|
bars_by_symbol, avg_dvol = await cls._fetch_price_data(
|
|
|
unique_symbols, date_range, oracle_url
|
|
|
)
|
|
|
sectors = await cls._fetch_sectors(unique_symbols, oracle_url)
|
|
|
macro_by_date = await cls._fetch_macro(date_range, db_dsn)
|
|
|
|
|
|
# Step 7: Build candidates_by_exec_date
|
|
|
candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]] = {}
|
|
|
for row in row_list:
|
|
|
eid = str(row.get("event_id", ""))
|
|
|
meta = event_meta.get(eid, {})
|
|
|
ticker = meta.get("ticker")
|
|
|
if not ticker:
|
|
|
logger.debug("snapshot_store_skip_no_ticker", event_id=eid)
|
|
|
continue
|
|
|
|
|
|
# Map entry_date → execution_date at this boundary
|
|
|
raw_exec = row.get("entry_date") or row.get("execution_date")
|
|
|
if raw_exec is None:
|
|
|
continue
|
|
|
if isinstance(raw_exec, str):
|
|
|
exec_date = dt.date.fromisoformat(raw_exec)
|
|
|
elif isinstance(raw_exec, dt.date):
|
|
|
exec_date = raw_exec
|
|
|
else:
|
|
|
continue
|
|
|
|
|
|
enriched = dict(row)
|
|
|
enriched["execution_date"] = exec_date
|
|
|
enriched["symbol"] = ticker
|
|
|
enriched["issuer_id"] = meta.get("issuer_id")
|
|
|
enriched["event_type"] = meta.get("event_type", "")
|
|
|
enriched["event_timestamp"] = meta.get("event_timestamp")
|
|
|
enriched["avg_dollar_volume"] = avg_dvol.get(ticker, 0.0)
|
|
|
enriched["sector"] = sectors.get(ticker, "UNKNOWN")
|
|
|
|
|
|
candidates_by_exec_date.setdefault(exec_date, []).append(enriched)
|
|
|
|
|
|
# Build bars_by_symbol_date: symbol -> date -> bar dict
|
|
|
bars_by_symbol_date: dict[str, dict[dt.date, dict[str, Any]]] = {}
|
|
|
for sym, date_bars in bars_by_symbol.items():
|
|
|
bars_by_symbol_date[sym] = date_bars
|
|
|
|
|
|
logger.info(
|
|
|
"snapshot_store_built",
|
|
|
exec_dates=len(candidates_by_exec_date),
|
|
|
symbols=len(bars_by_symbol_date),
|
|
|
)
|
|
|
return {
|
|
|
"candidates_by_exec_date": candidates_by_exec_date,
|
|
|
"bars_by_symbol_date": bars_by_symbol_date,
|
|
|
"macro_by_date": macro_by_date,
|
|
|
}
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
|
# Internal async helpers
|
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
@staticmethod
|
|
|
async def _fetch_event_metadata(
|
|
|
event_ids: list[str],
|
|
|
db_dsn: str,
|
|
|
) -> dict[str, dict[str, Any]]:
|
|
|
"""Batch-query Event + SymbolMaster for event metadata."""
|
|
|
if not event_ids:
|
|
|
return {}
|
|
|
try:
|
|
|
from sqlalchemy import select
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
|
|
from libs.db.models import Event, SymbolMaster
|
|
|
|
|
|
engine = create_async_engine(db_dsn, echo=False)
|
|
|
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
result: dict[str, dict[str, Any]] = {}
|
|
|
async with async_session() as session:
|
|
|
# Single IN clause for all event_ids
|
|
|
stmt = (
|
|
|
select(Event, SymbolMaster)
|
|
|
.outerjoin(SymbolMaster, Event.symbol_id == SymbolMaster.symbol_id)
|
|
|
.where(Event.event_id.in_(event_ids))
|
|
|
)
|
|
|
rows = (await session.execute(stmt)).all()
|
|
|
for event, sym in rows:
|
|
|
result[event.event_id] = {
|
|
|
"issuer_id": event.issuer_id,
|
|
|
"event_type": event.event_type,
|
|
|
"event_timestamp": event.filed_at_utc,
|
|
|
"ticker": sym.ticker if sym else None,
|
|
|
}
|
|
|
await engine.dispose()
|
|
|
return result
|
|
|
except Exception as exc:
|
|
|
logger.warning("snapshot_store_db_fetch_failed", error=str(exc))
|
|
|
return {}
|
|
|
|
|
|
@staticmethod
|
|
|
async def _fetch_price_data(
|
|
|
symbols: list[str],
|
|
|
date_range: tuple[dt.date, dt.date] | None,
|
|
|
oracle_url: str,
|
|
|
) -> tuple[dict[str, dict[dt.date, dict[str, Any]]], dict[str, float]]:
|
|
|
"""Fetch daily OHLCV bars and compute avg_dollar_volume per symbol."""
|
|
|
if not symbols or date_range is None:
|
|
|
return {}, {}
|
|
|
try:
|
|
|
from libs.oracle_client.client import OracleClient
|
|
|
from libs.oracle_client.price import PriceService
|
|
|
|
|
|
start_str = date_range[0].isoformat()
|
|
|
end_str = date_range[1].isoformat()
|
|
|
|
|
|
bars_by_symbol: dict[str, dict[dt.date, dict[str, Any]]] = {}
|
|
|
avg_dvol: dict[str, float] = {}
|
|
|
|
|
|
async with OracleClient(base_url=oracle_url) as client:
|
|
|
svc = PriceService(client)
|
|
|
for sym in symbols:
|
|
|
try:
|
|
|
resp = await svc.get_daily_bars(sym, start=start_str, end=end_str)
|
|
|
date_bars: dict[dt.date, dict[str, Any]] = {}
|
|
|
dollar_vols: list[float] = []
|
|
|
for bar in resp.bars:
|
|
|
d = dt.date.fromisoformat(bar.date)
|
|
|
b = {
|
|
|
"date": d,
|
|
|
"open": bar.open,
|
|
|
"high": bar.high,
|
|
|
"low": bar.low,
|
|
|
"close": bar.close,
|
|
|
"volume": bar.volume,
|
|
|
}
|
|
|
date_bars[d] = b
|
|
|
dollar_vols.append(bar.close * bar.volume)
|
|
|
bars_by_symbol[sym] = date_bars
|
|
|
# 20-day mean of dollar volume
|
|
|
if dollar_vols:
|
|
|
last_20 = dollar_vols[-20:]
|
|
|
avg_dvol[sym] = sum(last_20) / len(last_20)
|
|
|
else:
|
|
|
avg_dvol[sym] = 0.0
|
|
|
except Exception as sym_exc:
|
|
|
logger.warning(
|
|
|
"snapshot_store_price_fetch_failed", symbol=sym, error=str(sym_exc)
|
|
|
)
|
|
|
bars_by_symbol[sym] = {}
|
|
|
avg_dvol[sym] = 0.0
|
|
|
return bars_by_symbol, avg_dvol
|
|
|
except Exception as exc:
|
|
|
logger.warning("snapshot_store_oracle_failed", error=str(exc))
|
|
|
return {}, {}
|
|
|
|
|
|
@staticmethod
|
|
|
async def _fetch_sectors(
|
|
|
symbols: list[str],
|
|
|
oracle_url: str,
|
|
|
) -> dict[str, str]:
|
|
|
"""Fetch company sector for each symbol. Default 'UNKNOWN' if unavailable."""
|
|
|
if not symbols:
|
|
|
return {}
|
|
|
result: dict[str, str] = {}
|
|
|
try:
|
|
|
from libs.oracle_client.client import OracleClient
|
|
|
from libs.oracle_client.financial import FinancialService
|
|
|
|
|
|
async with OracleClient(base_url=oracle_url) as client:
|
|
|
for sym in symbols:
|
|
|
try:
|
|
|
info = await client.get(f"/api/v1/company/{sym}")
|
|
|
sector = info.get("sector") if isinstance(info, dict) else None
|
|
|
result[sym] = str(sector) if sector else "UNKNOWN"
|
|
|
except Exception:
|
|
|
result[sym] = "UNKNOWN"
|
|
|
except Exception as exc:
|
|
|
logger.warning("snapshot_store_sector_fetch_failed", error=str(exc))
|
|
|
# Default all remaining to UNKNOWN
|
|
|
for sym in symbols:
|
|
|
result.setdefault(sym, "UNKNOWN")
|
|
|
return result
|
|
|
|
|
|
@staticmethod
|
|
|
async def _fetch_macro(
|
|
|
date_range: tuple[dt.date, dt.date] | None,
|
|
|
db_dsn: str,
|
|
|
) -> dict[dt.date, dict[str, Any]]:
|
|
|
"""Load MacroObservation regime data from DB."""
|
|
|
if date_range is None:
|
|
|
return {}
|
|
|
try:
|
|
|
from sqlalchemy import select
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
|
|
from libs.db.models import MacroObservation
|
|
|
|
|
|
engine = create_async_engine(db_dsn, echo=False)
|
|
|
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
result: dict[dt.date, dict[str, Any]] = {}
|
|
|
async with async_session() as session:
|
|
|
stmt = select(MacroObservation).where(
|
|
|
MacroObservation.observation_date >= date_range[0],
|
|
|
MacroObservation.observation_date <= date_range[1],
|
|
|
)
|
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
|
for obs in rows:
|
|
|
d = obs.observation_date
|
|
|
result.setdefault(d, {})[obs.series_id] = (
|
|
|
float(obs.value) if obs.value is not None else None
|
|
|
)
|
|
|
await engine.dispose()
|
|
|
return result
|
|
|
except Exception as exc:
|
|
|
logger.warning("snapshot_store_macro_fetch_failed", error=str(exc))
|
|
|
return {}
|
|
|
|
|
|
@staticmethod
|
|
|
def _compute_date_range(
|
|
|
rows: list[dict[str, Any]],
|
|
|
) -> tuple[dt.date, dt.date] | None:
|
|
|
"""Compute (min_date, max_date) from entry_date/execution_date column."""
|
|
|
dates: list[dt.date] = []
|
|
|
for r in rows:
|
|
|
raw = r.get("entry_date") or r.get("execution_date")
|
|
|
if raw is None:
|
|
|
continue
|
|
|
if isinstance(raw, str):
|
|
|
try:
|
|
|
dates.append(dt.date.fromisoformat(raw))
|
|
|
except ValueError:
|
|
|
pass
|
|
|
elif isinstance(raw, dt.date):
|
|
|
dates.append(raw)
|
|
|
if not dates:
|
|
|
return None
|
|
|
return min(dates), max(dates)
|