# ============================================================ # Project: EDGAR Research Database # Description: SEC EDGAR Research Infrastructure # # Authors: # - Manel Kammoun — manel.kammoun@uqo.ca # - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca # - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca # # Institution: # Universite du Quebec en Outaouais (UQO) # # License: # Research / Academic Use # ============================================================ from __future__ import annotations from typing import TYPE_CHECKING import numpy as np import pandas as pd import pyarrow as pa import pyarrow.parquet as pq from tqdm import tqdm from src.utils.config import get_config from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL from src.utils.logging import get_logger if TYPE_CHECKING: from src.storage.database import DatabaseManager logger = get_logger("factors.aggregation") class FactorAggregator: def __init__(self, db: DatabaseManager) -> None: self.db = db cfg = get_config("factors").get("factors", {}) agg_cfg = cfg.get("aggregation", {}) self._top_pct = agg_cfg.get("top_percentile", 0.95) self._enable_weights = agg_cfg.get("enable_section_weights", True) self._section_weights = cfg.get("section_weights", {}) self._anchor_weight_profile = cfg.get("anchor_weight_profile", {}) def _get_anchor_ids(self) -> list[str]: row = self.db.execute( "SELECT * FROM chunk_scores LIMIT 1" ).fetchdf() return [c.replace("score_", "") for c in row.columns if c.startswith("score_")] def _compute_weighted_mean( self, df: pd.DataFrame, score_col: str, weight_col: str, ) -> float: w = df[weight_col].values s = df[score_col].values total_w = w.sum() if total_w == 0: return float(s.mean()) return float(np.dot(s, w) / total_w) def aggregate_to_filing(self, anchor_ids: list[str] | None = None) -> pd.DataFrame: """Filing-level statistics computed inside DuckDB (scales to tens of millions of chunks).""" if anchor_ids is None: anchor_ids = self._get_anchor_ids() logger.info("Aggregating chunk scores to filing level (SQL)...") weight_sql = [] for aid in anchor_ids: profile = self._anchor_weight_profile.get(aid, "operational_measures") weights = self._section_weights.get(profile, {}) if self._enable_weights else {} if weights: cases = " ".join(f"WHEN '{name}' THEN {float(w)}" for name, w in weights.items()) weight_sql.append(f"CASE section_name {cases} ELSE 1.0 END AS w_{aid}") else: weight_sql.append(f"1.0 AS w_{aid}") stats = [] for aid in anchor_ids: col = f"score_{aid}" stats.append(f"AVG({col}) AS {aid}_mean") stats.append(f"MAX({col}) AS {aid}_max") stats.append(f"COALESCE(STDDEV_POP({col}), 0.0) AS {aid}_std") stats.append( f"AVG(CASE WHEN {col} >= q_{aid} THEN {col} END) AS {aid}_top5pct" ) if self._enable_weights: stats.append(f"SUM({col} * w_{aid}) / NULLIF(SUM(w_{aid}), 0) AS {aid}_wmean") quantiles = ", ".join( f"QUANTILE_CONT(score_{aid}, {self._top_pct}) OVER (PARTITION BY accession_number) AS q_{aid}" for aid in anchor_ids ) sql = f""" WITH base AS ( SELECT *, {", ".join(weight_sql)}, {quantiles} FROM chunk_scores ) SELECT accession_number, ANY_VALUE(cik) AS cik, ANY_VALUE(ticker) AS ticker, ANY_VALUE(form_type) AS form_type, ANY_VALUE(filing_date) AS filing_date, ANY_VALUE(fiscal_year) AS fiscal_year, ANY_VALUE(fiscal_quarter) AS fiscal_quarter, COUNT(*) AS chunk_count, {", ".join(stats)} FROM base GROUP BY accession_number """ filing_df = self.db.fetch_df(sql) logger.info("Filing-level scores: %d rows", len(filing_df)) return filing_df def save_filing_scores(self, filing_df: pd.DataFrame) -> None: self.db.execute("DROP TABLE IF EXISTS filing_scores") self.db.conn.execute( "CREATE TABLE filing_scores AS SELECT * FROM filing_df" ) self.db.execute( "CREATE INDEX IF NOT EXISTS idx_fs_ticker_year ON filing_scores(ticker, fiscal_year)" ) logger.info("Saved %d rows to DuckDB filing_scores", len(filing_df)) out_dir = DATA_DIRS["parquet"] / "filing_scores" out_dir.mkdir(parents=True, exist_ok=True) table = pa.Table.from_pandas(filing_df, preserve_index=False) pq.write_to_dataset( table, root_path=str(out_dir), partition_cols=["form_type"], compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL, existing_data_behavior="overwrite_or_ignore", ) logger.info("Exported filing_scores to parquet") def aggregate_to_quarterly(self, anchor_ids: list[str] | None = None) -> pd.DataFrame: if anchor_ids is None: anchor_ids = self._get_anchor_ids() logger.info("Aggregating filing scores to quarterly panel...") stat_suffixes = ["mean", "max", "std", "top5pct"] if self._enable_weights: stat_suffixes.append("wmean") agg_exprs = { "filing_count": ("accession_number", "count"), "total_chunks": ("chunk_count", "sum"), } for aid in anchor_ids: for stat in stat_suffixes: col = f"{aid}_{stat}" if stat in ("max",): agg_exprs[col] = (col, "max") else: agg_exprs[col] = (col, "mean") filing_df = self.db.fetch_df("SELECT * FROM filing_scores") quarterly = filing_df.groupby( ["cik", "ticker", "fiscal_year", "fiscal_quarter"], as_index=False, ).agg(**agg_exprs) quarterly = quarterly.sort_values(["ticker", "fiscal_year", "fiscal_quarter"]) logger.info("Quarterly panel: %d rows", len(quarterly)) return quarterly def save_quarterly_panel(self, quarterly_df: pd.DataFrame) -> None: self.db.execute("DROP TABLE IF EXISTS quarterly_panel") self.db.conn.execute( "CREATE TABLE quarterly_panel AS SELECT * FROM quarterly_df" ) logger.info("Saved %d rows to DuckDB quarterly_panel", len(quarterly_df)) out_dir = DATA_DIRS["parquet"] / "quarterly_panel" out_dir.mkdir(parents=True, exist_ok=True) table = pa.Table.from_pandas(quarterly_df, preserve_index=False) pq.write_table( table, str(out_dir / "quarterly_panel.parquet"), compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL, ) logger.info("Exported quarterly_panel to parquet") def aggregate_to_annual(self, anchor_ids: list[str] | None = None) -> pd.DataFrame: if anchor_ids is None: anchor_ids = self._get_anchor_ids() logger.info("Aggregating filing scores to annual panel...") stat_suffixes = ["mean", "max", "std", "top5pct"] if self._enable_weights: stat_suffixes.append("wmean") agg_exprs = { "filing_count": ("accession_number", "count"), "total_chunks": ("chunk_count", "sum"), } for aid in anchor_ids: for stat in stat_suffixes: col = f"{aid}_{stat}" if stat in ("max",): agg_exprs[col] = (col, "max") else: agg_exprs[col] = (col, "mean") filing_df = self.db.fetch_df("SELECT * FROM filing_scores") annual = filing_df.groupby( ["cik", "ticker", "fiscal_year"], as_index=False, ).agg(**agg_exprs) annual = annual.sort_values(["ticker", "fiscal_year"]) logger.info("Annual panel: %d rows", len(annual)) return annual def save_annual_panel(self, annual_df: pd.DataFrame) -> None: self.db.execute("DROP TABLE IF EXISTS annual_panel") self.db.conn.execute( "CREATE TABLE annual_panel AS SELECT * FROM annual_df" ) logger.info("Saved %d rows to DuckDB annual_panel", len(annual_df)) out_dir = DATA_DIRS["parquet"] / "annual_panel" out_dir.mkdir(parents=True, exist_ok=True) table = pa.Table.from_pandas(annual_df, preserve_index=False) pq.write_table( table, str(out_dir / "annual_panel.parquet"), compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL, ) logger.info("Exported annual_panel to parquet")