spb/uqo-research-db
Public
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"]515253def merge_tables(db: DatabaseManager, shard_dirs: list[Path]) -> None:54 for t in TABLES:55 cols = [r[0] for r in db.execute(f"DESCRIBE {t}").fetchall()]56 for sd in shard_dirs:57 f = sd / "export" / f"{t}.parquet"58 if not f.exists():59 logger.warning("%s: missing %s", sd.name, f.name)60 continue61 file_cols = [c for c in pq.read_schema(str(f)).names if c in cols]62 col_list = ", ".join(file_cols)63 before = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]64 db.execute(f"INSERT OR REPLACE INTO {t} ({col_list}) SELECT {col_list} FROM read_parquet('{f}')")65 after = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]66 logger.info("%-20s <- %-8s +%d rows (now %d)", t, sd.name, after - before, after)676869def move_embeddings(shard_dirs: list[Path]) -> None:70 dest_root = DATA_DIRS["parquet"] / "embeddings"71 for sd in shard_dirs:72 src_root = sd / "embeddings"73 if not src_root.exists():74 continue75 n = 076 for part in src_root.glob("model_name=*"):77 dest = dest_root / part.name78 dest.mkdir(parents=True, exist_ok=True)79 for f in part.glob("*.parquet"):80 target = dest / f"{sd.name}-{f.name}"81 if not target.exists():82 shutil.move(str(f), str(target))83 n += 184 logger.info("embeddings <- %-8s %d files", sd.name, n)858687def normalise_legacy_embeddings() -> None:88 """Rewrite list<double> files as list<float> and drop the in-file model_name column."""89 root = DATA_DIRS["parquet"] / "embeddings"90 n = 091 for f in root.glob("model_name=*/*.parquet"):92 schema = pq.read_schema(str(f))93 vec = schema.field("vector").type94 needs = pa.types.is_list(vec) and pa.types.is_float64(vec.value_type)95 needs = needs or ("model_name" in schema.names)96 if not needs:97 continue98 t = pq.read_table(str(f))99 if "model_name" in t.column_names:100 t = t.drop_columns(["model_name"])101 t = t.set_column(t.schema.get_field_index("vector"), "vector", t.column("vector").cast(pa.list_(pa.float32())))102 tmp = f.with_suffix(".tmp.parquet")103 pq.write_table(t, str(tmp), compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL)104 tmp.replace(f)105 n += 1106 logger.info("normalised %d legacy embedding files to float32", n)107108109def export_single_parquet(db: DatabaseManager) -> None:110 for t in ["companies", "filings", "filing_sections", "filing_chunks"]:111 out_dir = DATA_DIRS["parquet"] / t112 out_dir.mkdir(parents=True, exist_ok=True)113 # remove stale hive partitions from earlier runs so the directory is one clean file114 for stale in out_dir.glob("form_type=*"):115 shutil.rmtree(stale, ignore_errors=True)116 path = out_dir / f"{t}.parquet"117 db.execute(f"COPY (SELECT * FROM {t}) TO '{path}' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 200000)")118 logger.info("parquet %-18s -> %s (%.1f MB)", t, path, path.stat().st_size / 1e6)119120121def main() -> None:122 p = argparse.ArgumentParser()123 p.add_argument("--shards-dir", type=str, default=str(DATA_DIRS["raw"].parent.parent / "shards"))124 p.add_argument("--nodes", type=str, default=None, help="comma list; default = every sub-directory")125 p.add_argument("--no-backup", action="store_true")126 args = p.parse_args()127128 shards_root = Path(args.shards_dir)129 shard_dirs = sorted(d for d in shards_root.iterdir() if d.is_dir()) if shards_root.exists() else []130 if args.nodes:131 wanted = {n.strip() for n in args.nodes.split(",")}132 shard_dirs = [d for d in shard_dirs if d.name in wanted]133 logger.info("merging %d shard(s): %s", len(shard_dirs), [d.name for d in shard_dirs])134135 if not args.no_backup and DUCKDB_PATH.exists():136 bak = DUCKDB_PATH.with_name(f"{DUCKDB_PATH.stem}.pre-merge-{date.today():%Y%m%d}.duckdb")137 if not bak.exists():138 shutil.copy2(DUCKDB_PATH, bak)139 logger.info("backup -> %s", bak)140141 db = DatabaseManager()142 db.initialize()143 db.execute("SET preserve_insertion_order=false")144 merge_tables(db, shard_dirs)145 move_embeddings(shard_dirs)146 normalise_legacy_embeddings()147 export_single_parquet(db)148 for t, n in db.get_stats().items():149 logger.info(" %-25s %d rows", t, n)150 db.close()151 logger.info("merge complete")152153154if __name__ == "__main__":155 main()156