spb/uqo-research-db
Public
1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Universite du Quebec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819from typing import TYPE_CHECKING2021import numpy as np22import pandas as pd23import pyarrow as pa24import pyarrow.parquet as pq25from tqdm import tqdm2627from src.utils.config import get_config28from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL29from src.utils.logging import get_logger3031if TYPE_CHECKING:32 from src.storage.database import DatabaseManager3334logger = get_logger("factors.aggregation")353637class FactorAggregator:3839 def __init__(self, db: DatabaseManager) -> None:40 self.db = db41 cfg = get_config("factors").get("factors", {})42 agg_cfg = cfg.get("aggregation", {})43 self._top_pct = agg_cfg.get("top_percentile", 0.95)44 self._enable_weights = agg_cfg.get("enable_section_weights", True)45 self._section_weights = cfg.get("section_weights", {})46 self._anchor_weight_profile = cfg.get("anchor_weight_profile", {})4748 def _get_anchor_ids(self) -> list[str]:49 row = self.db.execute(50 "SELECT * FROM chunk_scores LIMIT 1"51 ).fetchdf()52 return [c.replace("score_", "") for c in row.columns if c.startswith("score_")]5354 def _compute_weighted_mean(55 self,56 df: pd.DataFrame,57 score_col: str,58 weight_col: str,59 ) -> float:60 w = df[weight_col].values61 s = df[score_col].values62 total_w = w.sum()63 if total_w == 0:64 return float(s.mean())65 return float(np.dot(s, w) / total_w)6667 def aggregate_to_filing(self, anchor_ids: list[str] | None = None) -> pd.DataFrame:68 """Filing-level statistics computed inside DuckDB (scales to tens of millions of chunks)."""69 if anchor_ids is None:70 anchor_ids = self._get_anchor_ids()71 logger.info("Aggregating chunk scores to filing level (SQL)...")7273 weight_sql = []74 for aid in anchor_ids:75 profile = self._anchor_weight_profile.get(aid, "operational_measures")76 weights = self._section_weights.get(profile, {}) if self._enable_weights else {}77 if weights:78 cases = " ".join(f"WHEN '{name}' THEN {float(w)}" for name, w in weights.items())79 weight_sql.append(f"CASE section_name {cases} ELSE 1.0 END AS w_{aid}")80 else:81 weight_sql.append(f"1.0 AS w_{aid}")8283 stats = []84 for aid in anchor_ids:85 col = f"score_{aid}"86 stats.append(f"AVG({col}) AS {aid}_mean")87 stats.append(f"MAX({col}) AS {aid}_max")88 stats.append(f"COALESCE(STDDEV_POP({col}), 0.0) AS {aid}_std")89 stats.append(90 f"AVG(CASE WHEN {col} >= q_{aid} THEN {col} END) AS {aid}_top5pct"91 )92 if self._enable_weights:93 stats.append(f"SUM({col} * w_{aid}) / NULLIF(SUM(w_{aid}), 0) AS {aid}_wmean")9495 quantiles = ", ".join(96 f"QUANTILE_CONT(score_{aid}, {self._top_pct}) OVER (PARTITION BY accession_number) AS q_{aid}"97 for aid in anchor_ids98 )99 sql = f"""100 WITH base AS (101 SELECT *, {", ".join(weight_sql)}, {quantiles}102 FROM chunk_scores103 )104 SELECT accession_number,105 ANY_VALUE(cik) AS cik, ANY_VALUE(ticker) AS ticker, ANY_VALUE(form_type) AS form_type,106 ANY_VALUE(filing_date) AS filing_date, ANY_VALUE(fiscal_year) AS fiscal_year,107 ANY_VALUE(fiscal_quarter) AS fiscal_quarter,108 COUNT(*) AS chunk_count,109 {", ".join(stats)}110 FROM base111 GROUP BY accession_number112 """113 filing_df = self.db.fetch_df(sql)114 logger.info("Filing-level scores: %d rows", len(filing_df))115 return filing_df116117 def save_filing_scores(self, filing_df: pd.DataFrame) -> None:118 self.db.execute("DROP TABLE IF EXISTS filing_scores")119 self.db.conn.execute(120 "CREATE TABLE filing_scores AS SELECT * FROM filing_df"121 )122 self.db.execute(123 "CREATE INDEX IF NOT EXISTS idx_fs_ticker_year ON filing_scores(ticker, fiscal_year)"124 )125 logger.info("Saved %d rows to DuckDB filing_scores", len(filing_df))126127 out_dir = DATA_DIRS["parquet"] / "filing_scores"128 out_dir.mkdir(parents=True, exist_ok=True)129 table = pa.Table.from_pandas(filing_df, preserve_index=False)130 pq.write_to_dataset(131 table,132 root_path=str(out_dir),133 partition_cols=["form_type"],134 compression=PARQUET_COMPRESSION,135 compression_level=PARQUET_COMPRESSION_LEVEL,136 existing_data_behavior="overwrite_or_ignore",137 )138 logger.info("Exported filing_scores to parquet")139140 def aggregate_to_quarterly(self, anchor_ids: list[str] | None = None) -> pd.DataFrame:141 if anchor_ids is None:142 anchor_ids = self._get_anchor_ids()143144 logger.info("Aggregating filing scores to quarterly panel...")145146 stat_suffixes = ["mean", "max", "std", "top5pct"]147 if self._enable_weights:148 stat_suffixes.append("wmean")149150 agg_exprs = {151 "filing_count": ("accession_number", "count"),152 "total_chunks": ("chunk_count", "sum"),153 }154 for aid in anchor_ids:155 for stat in stat_suffixes:156 col = f"{aid}_{stat}"157 if stat in ("max",):158 agg_exprs[col] = (col, "max")159 else:160 agg_exprs[col] = (col, "mean")161162 filing_df = self.db.fetch_df("SELECT * FROM filing_scores")163 quarterly = filing_df.groupby(164 ["cik", "ticker", "fiscal_year", "fiscal_quarter"],165 as_index=False,166 ).agg(**agg_exprs)167168 quarterly = quarterly.sort_values(["ticker", "fiscal_year", "fiscal_quarter"])169 logger.info("Quarterly panel: %d rows", len(quarterly))170 return quarterly171172 def save_quarterly_panel(self, quarterly_df: pd.DataFrame) -> None:173 self.db.execute("DROP TABLE IF EXISTS quarterly_panel")174 self.db.conn.execute(175 "CREATE TABLE quarterly_panel AS SELECT * FROM quarterly_df"176 )177 logger.info("Saved %d rows to DuckDB quarterly_panel", len(quarterly_df))178179 out_dir = DATA_DIRS["parquet"] / "quarterly_panel"180 out_dir.mkdir(parents=True, exist_ok=True)181 table = pa.Table.from_pandas(quarterly_df, preserve_index=False)182 pq.write_table(183 table,184 str(out_dir / "quarterly_panel.parquet"),185 compression=PARQUET_COMPRESSION,186 compression_level=PARQUET_COMPRESSION_LEVEL,187 )188 logger.info("Exported quarterly_panel to parquet")189190 def aggregate_to_annual(self, anchor_ids: list[str] | None = None) -> pd.DataFrame:191 if anchor_ids is None:192 anchor_ids = self._get_anchor_ids()193194 logger.info("Aggregating filing scores to annual panel...")195196 stat_suffixes = ["mean", "max", "std", "top5pct"]197 if self._enable_weights:198 stat_suffixes.append("wmean")199200 agg_exprs = {201 "filing_count": ("accession_number", "count"),202 "total_chunks": ("chunk_count", "sum"),203 }204 for aid in anchor_ids:205 for stat in stat_suffixes:206 col = f"{aid}_{stat}"207 if stat in ("max",):208 agg_exprs[col] = (col, "max")209 else:210 agg_exprs[col] = (col, "mean")211212 filing_df = self.db.fetch_df("SELECT * FROM filing_scores")213 annual = filing_df.groupby(214 ["cik", "ticker", "fiscal_year"],215 as_index=False,216 ).agg(**agg_exprs)217218 annual = annual.sort_values(["ticker", "fiscal_year"])219 logger.info("Annual panel: %d rows", len(annual))220 return annual221222 def save_annual_panel(self, annual_df: pd.DataFrame) -> None:223 self.db.execute("DROP TABLE IF EXISTS annual_panel")224 self.db.conn.execute(225 "CREATE TABLE annual_panel AS SELECT * FROM annual_df"226 )227 logger.info("Saved %d rows to DuckDB annual_panel", len(annual_df))228229 out_dir = DATA_DIRS["parquet"] / "annual_panel"230 out_dir.mkdir(parents=True, exist_ok=True)231 table = pa.Table.from_pandas(annual_df, preserve_index=False)232 pq.write_table(233 table,234 str(out_dir / "annual_panel.parquet"),235 compression=PARQUET_COMPRESSION,236 compression_level=PARQUET_COMPRESSION_LEVEL,237 )238 logger.info("Exported annual_panel to parquet")239