spb/uqo-research-db

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 pathlib import Path20from typing import TYPE_CHECKING2122import numpy as np23import pandas as pd24import pyarrow as pa25import pyarrow.parquet as pq2627from src.utils.config import get_config28from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL, PROJECT_ROOT29from src.utils.logging import get_logger3031if TYPE_CHECKING:32    from src.storage.database import DatabaseManager3334logger = get_logger("factors.panels")353637class PanelExporter:3839    def __init__(self, db: DatabaseManager) -> None:40        self.db = db41        cfg = get_config("factors").get("factors", {}).get("export", {})42        self._formats = cfg.get("formats", ["parquet", "csv"])43        self._output_dir = PROJECT_ROOT / cfg.get("output_dir", "outputs/factors")4445    def build_annual_panel(46        self,47        anchor_ids: list[str],48        drift_df: pd.DataFrame | None = None,49        novelty_df: pd.DataFrame | None = None,50    ) -> pd.DataFrame:51        logger.info("Building comprehensive annual panel...")5253        annual = self.db.fetch_df("SELECT * FROM annual_panel")5455        companies = self.db.fetch_df(56            "SELECT cik, ticker, company_name, sector, sub_industry FROM companies"57        )58        annual = annual.merge(59            companies[["ticker", "company_name", "sector", "sub_industry"]],60            on="ticker",61            how="left",62        )6364        from src.factors.longitudinal import LongitudinalBuilder65        longi = LongitudinalBuilder(self.db)66        annual = longi.add_yoy_changes(annual, anchor_ids)67        annual = longi.add_acceleration(annual, anchor_ids)68        annual = longi.add_rolling_stats(annual, anchor_ids)6970        if drift_df is not None and not drift_df.empty:71            annual = annual.merge(drift_df, on=["ticker", "fiscal_year"], how="left")72        else:73            annual["embedding_drift"] = np.nan7475        if novelty_df is not None and not novelty_df.empty:76            annual = annual.merge(novelty_df, on=["ticker", "fiscal_year"], how="left")77        else:78            annual["novelty_score"] = np.nan7980        annual["log_total_chunks"] = np.log1p(annual["total_chunks"])8182        id_cols = [83            "cik", "ticker", "company_name", "sector", "sub_industry", "fiscal_year",84            "filing_count", "total_chunks", "log_total_chunks",85        ]86        score_cols = sorted([c for c in annual.columns if any(87            c.startswith(f"{aid}_") for aid in anchor_ids88        )])89        feature_cols = ["embedding_drift", "novelty_score"]90        existing = [c for c in id_cols + score_cols + feature_cols if c in annual.columns]91        annual = annual[existing]9293        annual = annual.sort_values(["ticker", "fiscal_year"]).reset_index(drop=True)94        logger.info(95            "Annual panel: %d rows x %d columns (%d firms, %d years)",96            len(annual), len(annual.columns),97            annual["ticker"].nunique(),98            annual["fiscal_year"].nunique(),99        )100        return annual101102    def build_quarterly_panel(self, anchor_ids: list[str]) -> pd.DataFrame:103        quarterly = self.db.fetch_df("SELECT * FROM quarterly_panel")104        companies = self.db.fetch_df(105            "SELECT ticker, company_name, sector, sub_industry FROM companies"106        )107        quarterly = quarterly.merge(companies, on="ticker", how="left")108        quarterly = quarterly.sort_values(109            ["ticker", "fiscal_year", "fiscal_quarter"]110        ).reset_index(drop=True)111        logger.info("Quarterly panel: %d rows x %d columns", len(quarterly), len(quarterly.columns))112        return quarterly113114    def save_final_panel(self, panel: pd.DataFrame, name: str) -> None:115        self.db.execute(f"DROP TABLE IF EXISTS {name}_final")116        self.db.conn.execute(f"CREATE TABLE {name}_final AS SELECT * FROM panel")117        logger.info("Saved %s_final (%d rows) to DuckDB", name, len(panel))118119        out_dir = DATA_DIRS["parquet"] / f"{name}_final"120        out_dir.mkdir(parents=True, exist_ok=True)121        table = pa.Table.from_pandas(panel, preserve_index=False)122        pq.write_table(123            table,124            str(out_dir / f"{name}_final.parquet"),125            compression=PARQUET_COMPRESSION,126            compression_level=PARQUET_COMPRESSION_LEVEL,127        )128129        self._output_dir.mkdir(parents=True, exist_ok=True)130        if "csv" in self._formats:131            csv_path = self._output_dir / f"{name}_panel.csv"132            panel.to_csv(csv_path, index=False)133            logger.info("Exported %s to CSV: %s", name, csv_path)134135        if "parquet" in self._formats:136            pq_path = self._output_dir / f"{name}_panel.parquet"137            pq.write_table(138                table,139                str(pq_path),140                compression=PARQUET_COMPRESSION,141                compression_level=PARQUET_COMPRESSION_LEVEL,142            )143            logger.info("Exported %s to parquet: %s", name, pq_path)144145    def export_all(146        self,147        annual: pd.DataFrame,148        quarterly: pd.DataFrame,149    ) -> None:150        self.save_final_panel(annual, "annual")151        self.save_final_panel(quarterly, "quarterly")152        logger.info("All panels exported")153