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# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617"""Merge the worker shards into the master database (run on the master node).1819Expected layout (pulled by the orchestrator):20    data/shards/<node>/export/{companies,filings,filing_sections,filing_chunks,embeddings_metadata}.parquet21    data/shards/<node>/embeddings/model_name=<model>/*.parquet2223Steps24 1. back up the master DuckDB, then INSERT OR REPLACE every shard table into it;25 2. move the shards' embedding files under data/parquet/embeddings/model_name=<model>/;26 3. rewrite legacy float64 embedding files as float32 (uniform schema);27 4. regenerate the single-file parquet exports used by SPID and build_shares28    (companies, filings, filing_sections, filing_chunks).29"""3031from __future__ import annotations3233import argparse34import shutil35import sys36from datetime import date37from pathlib import Path3839import pyarrow as pa40import pyarrow.parquet as pq4142sys.path.insert(0, str(Path(__file__).resolve().parent.parent))4344from src.storage.database import DatabaseManager45from src.utils.constants import DATA_DIRS, DUCKDB_PATH, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL46from src.utils.logging import setup_logging4748logger = setup_logging("merge_shards")4950TABLES = ["companies", "filings", "filing_sections", "filing_chunks", "embeddings_metadata"]515253PK = {"filings": "accession_number", "filing_sections": "section_id",54      "filing_chunks": "chunk_id", "embeddings_metadata": "embedding_id"}555657def merge_tables(db: DatabaseManager, shard_dirs: list[Path]) -> None:58    """Append the shards' rows that the master does not have yet (keys are globally unique).5960    ``companies`` is already complete on the master (same universe everywhere); only the61    ``shard_node`` tag is updated from the shard's filings.62    """63    for t in ["filings", "filing_sections", "filing_chunks", "embeddings_metadata"]:64        cols = [r[0] for r in db.execute(f"DESCRIBE {t}").fetchall()]65        pk = PK[t]66        for sd in shard_dirs:67            f = sd / "export" / f"{t}.parquet"68            if not f.exists():69                logger.warning("%s: missing %s", sd.name, f.name)70                continue71            file_cols = [c for c in pq.read_schema(str(f)).names if c in cols]72            col_list = ", ".join(file_cols)73            before = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]74            db.execute(75                f"INSERT INTO {t} ({col_list}) SELECT {col_list} FROM read_parquet('{f}') s "76                f"WHERE NOT EXISTS (SELECT 1 FROM {t} m WHERE m.{pk} = s.{pk})"77            )78            after = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]79            logger.info("%-20s <- %-8s +%d rows (now %d)", t, sd.name, after - before, after)80    for sd in shard_dirs:81        f = sd / "export" / "filings.parquet"82        if f.exists():83            db.execute(f"UPDATE companies SET shard_node = '{sd.name}' WHERE cik IN "84                       f"(SELECT DISTINCT cik FROM read_parquet('{f}'))")85    logger.info("companies.shard_node updated from the shards")868788def move_embeddings(shard_dirs: list[Path]) -> None:89    dest_root = DATA_DIRS["parquet"] / "embeddings"90    for sd in shard_dirs:91        src_root = sd / "embeddings"92        if not src_root.exists():93            continue94        n = 095        for part in src_root.glob("model_name=*"):96            dest = dest_root / part.name97            dest.mkdir(parents=True, exist_ok=True)98            for f in part.glob("*.parquet"):99                target = dest / f"{sd.name}-{f.name}"100                if not target.exists():101                    shutil.move(str(f), str(target))102                    n += 1103        logger.info("embeddings <- %-8s %d files", sd.name, n)104105106def normalise_legacy_embeddings() -> None:107    """Rewrite list<double> files as list<float> and drop the in-file model_name column."""108    root = DATA_DIRS["parquet"] / "embeddings"109    n = 0110    for f in root.glob("model_name=*/*.parquet"):111        schema = pq.read_schema(str(f))112        vec = schema.field("vector").type113        needs = pa.types.is_list(vec) and pa.types.is_float64(vec.value_type)114        needs = needs or ("model_name" in schema.names)115        if not needs:116            continue117        t = pq.read_table(str(f))118        if "model_name" in t.column_names:119            t = t.drop_columns(["model_name"])120        t = t.set_column(t.schema.get_field_index("vector"), "vector", t.column("vector").cast(pa.list_(pa.float32())))121        tmp = f.with_suffix(".tmp.parquet")122        pq.write_table(t, str(tmp), compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL)123        tmp.replace(f)124        n += 1125    logger.info("normalised %d legacy embedding files to float32", n)126127128def export_single_parquet(db: DatabaseManager) -> None:129    for t in ["companies", "filings", "filing_sections", "filing_chunks"]:130        out_dir = DATA_DIRS["parquet"] / t131        out_dir.mkdir(parents=True, exist_ok=True)132        # remove stale hive partitions from earlier runs so the directory is one clean file133        for stale in out_dir.glob("form_type=*"):134            shutil.rmtree(stale, ignore_errors=True)135        path = out_dir / f"{t}.parquet"136        db.execute(f"COPY (SELECT * FROM {t}) TO '{path}' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 200000)")137        logger.info("parquet %-18s -> %s (%.1f MB)", t, path, path.stat().st_size / 1e6)138139140def main() -> None:141    p = argparse.ArgumentParser()142    p.add_argument("--shards-dir", type=str, default=str(DATA_DIRS["raw"].parent.parent / "shards"))143    p.add_argument("--nodes", type=str, default=None, help="comma list; default = every sub-directory")144    p.add_argument("--no-backup", action="store_true")145    args = p.parse_args()146147    shards_root = Path(args.shards_dir)148    shard_dirs = sorted(d for d in shards_root.iterdir() if d.is_dir()) if shards_root.exists() else []149    if args.nodes:150        wanted = {n.strip() for n in args.nodes.split(",")}151        shard_dirs = [d for d in shard_dirs if d.name in wanted]152    logger.info("merging %d shard(s): %s", len(shard_dirs), [d.name for d in shard_dirs])153154    if not args.no_backup and DUCKDB_PATH.exists():155        bak = DUCKDB_PATH.with_name(f"{DUCKDB_PATH.stem}.pre-merge-{date.today():%Y%m%d}.duckdb")156        if not bak.exists():157            shutil.copy2(DUCKDB_PATH, bak)158            logger.info("backup -> %s", bak)159160    db = DatabaseManager()161    db.initialize()162    db.execute("SET preserve_insertion_order=false")163    merge_tables(db, shard_dirs)164    move_embeddings(shard_dirs)165    normalise_legacy_embeddings()166    export_single_parquet(db)167    for t, n in db.get_stats().items():168        logger.info("  %-25s %d rows", t, n)169    db.close()170    logger.info("merge complete")171172173if __name__ == "__main__":174    main()175