"""Enrich snapshot with historical earnings surprise features. Adds PIT-safe trailing quarterly surprise features derived from Oracle's earnings surprise endpoint. The current event's earnings surprise is already materialized separately as `earnings_surprise_pct`; this script adds the prior 12 quarter surprise vector plus compact summary statistics. Usage: uv run python3 scripts/enrich_earnings_history_features.py \ --input data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_canonical \ --output data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_canonical_eh """ from __future__ import annotations import argparse import concurrent.futures import datetime as dt import json import sys from pathlib import Path from typing import Any import pyarrow as pa import pyarrow.parquet as pq import requests sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) ORACLE_URL = "http://localhost:18001" SURPRISE_VECTOR_COLS = [f"sue_lag_{i}_pct" for i in range(1, 13)] SUMMARY_COLS = [ "sue_hist_mean_4q", "sue_hist_mean_8q", "sue_hist_mean_12q", "sue_hist_pos_rate_4q", "sue_hist_pos_rate_12q", "sue_hist_latest_pct", "sue_hist_streak_pos", ] FEATURES = SURPRISE_VECTOR_COLS + SUMMARY_COLS def _parse_date(value: Any) -> dt.date | None: if value is None: return None text = str(value) if not text: return None try: return dt.date.fromisoformat(text[:10]) except ValueError: return None def _fetch_surprise_quarters(ticker: str) -> list[dict[str, Any]]: resp = requests.get(f"{ORACLE_URL}/api/v1/earnings/surprise/{ticker}", timeout=30) if resp.status_code != 200: return [] data = resp.json() quarters = data.get("quarters") or [] return [q for q in quarters if q.get("reported_date")] def _match_current_quarter(quarters: list[dict[str, Any]], event_date: dt.date) -> int | None: best_idx: int | None = None best_delta = 999999 for idx, quarter in enumerate(quarters): reported_date = _parse_date(quarter.get("reported_date")) if reported_date is None: continue delta = abs((reported_date - event_date).days) if delta <= 5 and delta < best_delta: best_delta = delta best_idx = idx return best_idx def _mean(values: list[float]) -> float | None: if not values: return None return sum(values) / len(values) def _positive_rate(values: list[float]) -> float | None: if not values: return None return sum(1 for value in values if value > 0.0) / len(values) def _positive_streak(values: list[float]) -> float | None: if not values: return None streak = 0 for value in values: if value > 0.0: streak += 1 else: break return float(streak) def compute_history_features( ticker: str, event_date: str | dt.date | None, quarters: list[dict[str, Any]] | None = None, ) -> dict[str, float | None]: event_dt = _parse_date(event_date) result = {feature: None for feature in FEATURES} if event_dt is None: return result if quarters is None: quarters = _fetch_surprise_quarters(ticker) if not quarters: return result matched_idx = _match_current_quarter(quarters, event_dt) if matched_idx is None: return result prior_quarters = quarters[matched_idx + 1 : matched_idx + 13] prior_values: list[float] = [] for idx, quarter in enumerate(prior_quarters, start=1): surprise_pct = quarter.get("surprise_percentage") if surprise_pct is None: prior_values.append(float("nan")) continue value = float(surprise_pct) prior_values.append(value) result[f"sue_lag_{idx}_pct"] = value clean_values = [value for value in prior_values if value == value] clean_4 = clean_values[:4] clean_8 = clean_values[:8] clean_12 = clean_values[:12] result["sue_hist_mean_4q"] = _mean(clean_4) result["sue_hist_mean_8q"] = _mean(clean_8) result["sue_hist_mean_12q"] = _mean(clean_12) result["sue_hist_pos_rate_4q"] = _positive_rate(clean_4) result["sue_hist_pos_rate_12q"] = _positive_rate(clean_12) result["sue_hist_latest_pct"] = clean_values[0] if clean_values else None result["sue_hist_streak_pos"] = _positive_streak(clean_values) return result def enrich_split(input_path: Path, output_path: Path) -> None: table = pq.read_table(input_path) n = len(table) if n == 0 or "ticker" not in table.column_names or "event_date" not in table.column_names: print(" Empty or schema-less split, copying") pq.write_table(table, output_path) return existing_cols = set(table.column_names) if all(feature in existing_cols for feature in FEATURES): null_counts = {feature: table.column(feature).null_count for feature in FEATURES} if all(count < n * 0.5 for count in null_counts.values()): print(f" Already enriched ({null_counts}), copying") pq.write_table(table, output_path) return tickers = table.column("ticker").to_pylist() event_dates = table.column("event_date").to_pylist() event_types = table.column("event_type").to_pylist() if "event_type" in table.column_names else [None] * n results = {feature: [None] * n for feature in FEATURES} success = 0 earnings_tickers = sorted( { str(tickers[idx]) for idx in range(n) if event_types[idx] == "earnings_release" and tickers[idx] } ) surprise_cache: dict[str, list[dict[str, Any]]] = {} if earnings_tickers: print(f" Prefetching surprise history for {len(earnings_tickers)} tickers...") def _load_one(ticker: str) -> tuple[str, list[dict[str, Any]]]: try: return ticker, _fetch_surprise_quarters(ticker) except Exception: return ticker, [] with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool: for loaded, (ticker, quarters) in enumerate(pool.map(_load_one, earnings_tickers), start=1): surprise_cache[ticker] = quarters if loaded % 100 == 0 or loaded == len(earnings_tickers): print(f" surprise prefetch {loaded}/{len(earnings_tickers)}") for idx in range(n): ticker = str(tickers[idx]) if tickers[idx] is not None else "" event_date = event_dates[idx] event_type = str(event_types[idx] or "") if idx % 100 == 0: print(f" {idx}/{n} ({success} enriched)...") if event_type != "earnings_release" or not ticker: continue quarters = surprise_cache.get(ticker, []) try: features = compute_history_features(ticker, event_date, quarters) if any(features[feature] is not None for feature in FEATURES): success += 1 for feature in FEATURES: results[feature][idx] = features[feature] except Exception as exc: if idx < 5: print(f" Error {ticker} {event_date}: {exc}") print(f" Done: {success}/{n}") for feature in FEATURES: arr = pa.array(results[feature], type=pa.float64()) if feature in existing_cols: col_idx = table.column_names.index(feature) table = table.set_column(col_idx, feature, arr) else: table = table.append_column(feature, arr) pq.write_table(table, output_path) print(f" Written to {output_path}") def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--input", required=True) parser.add_argument("--output", required=True) args = parser.parse_args() input_dir = Path(args.input) output_dir = Path(args.output) output_dir.mkdir(parents=True, exist_ok=True) manifest_src = input_dir / "manifest.json" if manifest_src.exists(): manifest = json.loads(manifest_src.read_text()) source_snapshot_id = manifest.get("snapshot_id") manifest["snapshot_id"] = output_dir.name manifest["output_dir"] = str(output_dir.resolve()) manifest["created_at_utc"] = dt.datetime.now(dt.UTC).isoformat() manifest["enrichment_source_snapshot_id"] = source_snapshot_id manifest["export_enrichments"] = sorted(set((manifest.get("export_enrichments") or []) + FEATURES)) (output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2)) for split in ("train", "valid", "test"): input_path = input_dir / f"{split}.parquet" if not input_path.exists(): continue output_path = output_dir / f"{split}.parquet" print(f"Enriching {split}...") enrich_split(input_path, output_path) print("All done!") if __name__ == "__main__": main()