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 typing import TYPE_CHECKING2021import numpy as np22import pandas as pd23import pyarrow as pa24import pyarrow.dataset as ds25import pyarrow.parquet as pq26from tqdm import tqdm2728from src.utils.config import get_config29from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL30from src.utils.logging import get_logger3132if TYPE_CHECKING:33    from src.storage.database import DatabaseManager3435logger = get_logger("factors.scoring")363738def _vectors_from_column(col) -> np.ndarray:39    """Decode a (chunked) list<float|double> Arrow column into a 2-D float32 array without Python loops."""40    import pyarrow as pa41    if isinstance(col, pa.ChunkedArray):42        col = col.combine_chunks()43    if pa.types.is_fixed_size_list(col.type):44        width = col.type.list_size45        flat = col.values46    else:47        offsets = col.offsets.to_numpy()48        width = int(offsets[1] - offsets[0]) if len(offsets) > 1 else 049        flat = col.values50    arr = np.asarray(flat.to_numpy(zero_copy_only=False), dtype=np.float32)51    return arr.reshape(-1, width) if width else arr.reshape(len(col), -1)525354class ScoringEngine:5556    def __init__(57        self,58        db: DatabaseManager,59        batch_size: int | None = None,60    ) -> None:61        self.db = db62        cfg = get_config("factors").get("factors", {}).get("scoring", {})63        self._batch_size = batch_size or cfg.get("batch_size", 10_000)6465    def _load_chunk_metadata(self) -> pd.DataFrame:66        sql = """67            SELECT68                c.chunk_id,69                c.accession_number,70                f.cik,71                f.ticker,72                f.form_type,73                f.filing_date,74                EXTRACT(YEAR FROM f.filing_date)::INTEGER AS fiscal_year,75                CEIL(EXTRACT(MONTH FROM f.filing_date) / 3.0)::INTEGER AS fiscal_quarter,76                COALESCE(s.section_name, 'full_body') AS section_name77            FROM filing_chunks c78            JOIN filings f ON c.accession_number = f.accession_number79            LEFT JOIN filing_sections s ON c.section_id = s.section_id80        """81        logger.info("Loading chunk metadata from DuckDB...")82        meta = self.db.fetch_df(sql)83        logger.info("Loaded metadata for %d chunks", len(meta))84        return meta8586    def _check_normalization(self, sample_vectors: np.ndarray) -> bool:87        norms = np.linalg.norm(sample_vectors, axis=1)88        mean_norm = norms.mean()89        is_normalized = 0.99 <= mean_norm <= 1.0190        logger.info(91            "Normalization check: mean norm=%.4f, pre-normalized=%s",92            mean_norm, is_normalized,93        )94        return is_normalized9596    def compute_scores(97        self,98        anchor_ids: list[str],99        anchor_matrix: np.ndarray,100    ) -> pd.DataFrame:101        anchor_matrix = anchor_matrix.astype(np.float32)102        chunk_meta = self._load_chunk_metadata()103        chunk_meta_indexed = chunk_meta.set_index("chunk_id")104105        emb_dir = DATA_DIRS["parquet"] / "embeddings"106        dataset = ds.dataset(str(emb_dir), format="parquet", partitioning="hive")107108        scanner = dataset.scanner(109            columns=["source_id", "vector"],110            batch_size=self._batch_size,111        )112113        all_ids: list[list[str]] = []114        all_scores: list[np.ndarray] = []115        normalization_checked = False116        is_prenormalized = True117        total_processed = 0118119        for batch in tqdm(scanner.to_batches(), desc="Scoring chunks", unit="batch"):120            if batch.num_rows == 0:121                continue122            source_ids = batch.column("source_id").to_pylist()123            vectors = _vectors_from_column(batch.column("vector"))124125            if not normalization_checked:126                is_prenormalized = self._check_normalization(vectors[:min(100, len(vectors))])127                normalization_checked = True128129            if not is_prenormalized:130                norms = np.linalg.norm(vectors, axis=1, keepdims=True)131                norms[norms == 0] = 1.0132                vectors = vectors / norms133134            scores = vectors @ anchor_matrix.T135136            all_ids.append(source_ids)137            all_scores.append(scores.astype(np.float32))138            total_processed += len(source_ids)139140        logger.info("Computed scores for %d chunks", total_processed)141142        flat_ids = []143        for id_list in all_ids:144            flat_ids.extend(id_list)145        score_matrix = np.vstack(all_scores)146147        score_cols = {f"score_{aid}": score_matrix[:, i] for i, aid in enumerate(anchor_ids)}148        scores_df = pd.DataFrame({"chunk_id": flat_ids, **score_cols})149150        scores_df = scores_df.merge(chunk_meta, on="chunk_id", how="inner")151        dropped = total_processed - len(scores_df)152        if dropped > 0:153            logger.warning("%d chunks had no metadata match (dropped)", dropped)154155        logger.info("Final chunk_scores: %d rows x %d columns", len(scores_df), len(scores_df.columns))156        return scores_df157158    def save_chunk_scores(self, scores_df: pd.DataFrame) -> None:159        anchor_ids = [c.replace("score_", "") for c in scores_df.columns if c.startswith("score_")]160161        self.db.execute("DROP TABLE IF EXISTS chunk_scores")162163        score_cols_ddl = ",\n    ".join(164            f"score_{aid} DOUBLE" for aid in anchor_ids165        )166        ddl = f"""167            CREATE TABLE chunk_scores (168                chunk_id         VARCHAR PRIMARY KEY,169                accession_number VARCHAR,170                cik              VARCHAR,171                ticker           VARCHAR,172                form_type        VARCHAR,173                filing_date      DATE,174                fiscal_year      INTEGER,175                fiscal_quarter   INTEGER,176                section_name     VARCHAR,177                {score_cols_ddl},178                created_at       TIMESTAMP DEFAULT current_timestamp179            )180        """181        self.db.execute(ddl)182183        cols = [184            "chunk_id", "accession_number", "cik", "ticker", "form_type",185            "filing_date", "fiscal_year", "fiscal_quarter", "section_name",186        ] + [f"score_{aid}" for aid in anchor_ids]187        insert_df = scores_df[cols].copy()188189        col_list = ", ".join(cols)190        self.db.conn.execute(f"INSERT INTO chunk_scores ({col_list}) SELECT * FROM insert_df")191192        self.db.execute("CREATE INDEX IF NOT EXISTS idx_cs_accession ON chunk_scores(accession_number)")193        self.db.execute("CREATE INDEX IF NOT EXISTS idx_cs_ticker_year ON chunk_scores(ticker, fiscal_year)")194195        logger.info("Saved %d rows to DuckDB chunk_scores", len(insert_df))196197        out_dir = DATA_DIRS["parquet"] / "chunk_scores"198        out_dir.mkdir(parents=True, exist_ok=True)199        table = pa.Table.from_pandas(insert_df, preserve_index=False)200        pq.write_to_dataset(201            table,202            root_path=str(out_dir),203            partition_cols=["form_type"],204            compression=PARQUET_COMPRESSION,205            compression_level=PARQUET_COMPRESSION_LEVEL,206            existing_data_behavior="overwrite_or_ignore",207        )208        logger.info("Exported chunk_scores to parquet")209