diff --git a/libs/export/snapshot_export.py b/libs/export/snapshot_export.py index 483624f..31366f5 100644 --- a/libs/export/snapshot_export.py +++ b/libs/export/snapshot_export.py @@ -15,7 +15,11 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from libs.common.logging import get_logger -from libs.common.time_utils import filing_time_bucket as classify_time_bucket, utc_now +from libs.common.time_utils import ( + filing_time_bucket as classify_time_bucket, + trading_days_between, + utc_now, +) logger = get_logger(__name__) @@ -111,8 +115,14 @@ def _rows_to_table(rows: list[dict[str, Any]]) -> pa.Table: """Convert list of dicts to a PyArrow Table.""" if not rows: return pa.table({}) - # Collect all keys - keys = list(rows[0].keys()) + # Collect ALL keys from ALL rows (not just the first) + keys: list[str] = [] + seen: set[str] = set() + for row in rows: + for k in row: + if k not in seen: + keys.append(k) + seen.add(k) arrays: dict[str, list[Any]] = {k: [] for k in keys} for row in rows: for k in keys: @@ -281,21 +291,51 @@ async def _enrich_macro_features(rows: list[dict[str, Any]]) -> None: def _enrich_prior_event_drift(rows: list[dict[str, Any]]) -> None: - """Add prior_event_fwd5d: the same ticker's most recent prior event's fwd_return_5d. + """Add a PIT-safe prior_event_fwd5d from the same ticker's most recent prior event. - This is a non-leaking cross-event momentum feature. By the time the current - event occurs, the prior event's 5-day return is fully realized. - Sorted by (ticker, event_date) and looks back one event per ticker. + The feature is only populated when the prior event's 5-trading-day forward + window is fully realized strictly before the current event date. If the most + recent prior event has not fully realized yet, the current row gets null. """ - sorted_rows = sorted(rows, key=lambda r: (r.get("ticker", ""), r.get("event_date", ""))) - prev_by_ticker: dict[str, float | None] = {} + sorted_rows = sorted( + rows, + key=lambda r: ( + str(r.get("ticker", "")), + str(r.get("event_date", "")), + str(r.get("entry_date", "")), + ), + ) + prev_by_ticker: dict[str, tuple[float | None, dt.date | None]] = {} + realized_on_cache: dict[dt.date, dt.date | None] = {} + + def _fwd5_realized_on(entry_date: dt.date | None) -> dt.date | None: + if entry_date is None: + return None + cached = realized_on_cache.get(entry_date) + if cached is not None or entry_date in realized_on_cache: + return cached + trading_days = trading_days_between(entry_date, entry_date + dt.timedelta(days=14)) + realized_on = trading_days[5] if len(trading_days) > 5 else None + realized_on_cache[entry_date] = realized_on + return realized_on for row in sorted_rows: - ticker = row.get("ticker", "") - row["prior_event_fwd5d"] = prev_by_ticker.get(ticker) + ticker = str(row.get("ticker", "")) + current_event_date = _parse_iso_date(row.get("event_date")) + prior_value: float | None = None + prior = prev_by_ticker.get(ticker) + if prior is not None and current_event_date is not None: + candidate_value, realized_on = prior + if realized_on is not None and current_event_date > realized_on: + prior_value = candidate_value + row["prior_event_fwd5d"] = prior_value + fwd5 = row.get("fwd_return_5d") - if fwd5 is not None: - prev_by_ticker[ticker] = fwd5 + entry_date = _parse_iso_date(row.get("entry_date")) + prev_by_ticker[ticker] = ( + float(fwd5) if fwd5 is not None else None, + _fwd5_realized_on(entry_date), + ) async def _backfill_market_fields(rows: list[dict[str, Any]]) -> None: