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"""Export this node's DuckDB tables to data/export/<table>.parquet for the merge.1819Embeddings are already parquet files under data/parquet/embeddings/ and are20rsync'ed as-is by the orchestrator.21"""2223from __future__ import annotations2425import argparse26import sys27from pathlib import Path2829sys.path.insert(0, str(Path(__file__).resolve().parent.parent))3031from src.storage.database import DatabaseManager32from src.utils.constants import PROJECT_ROOT33from src.utils.logging import setup_logging3435logger = setup_logging("export_shard")3637TABLES = ["companies", "filings", "filing_sections", "filing_chunks", "embeddings_metadata"]383940def main() -> None:41    p = argparse.ArgumentParser()42    p.add_argument("--out", type=str, default=str(PROJECT_ROOT / "data" / "export"))43    p.add_argument("--tables", type=str, default=",".join(TABLES))44    args = p.parse_args()45    out = Path(args.out)46    out.mkdir(parents=True, exist_ok=True)47    db = DatabaseManager()48    for t in args.tables.split(","):49        t = t.strip()50        path = out / f"{t}.parquet"51        n = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]52        db.execute(f"COPY (SELECT * FROM {t}) TO '{path}' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 200000)")53        logger.info("exported %-20s %10d rows -> %s (%.1f MB)", t, n, path, path.stat().st_size / 1e6)54    db.close()55    (out / "EXPORT_DONE").write_text("ok\n")565758if __name__ == "__main__":59    main()60