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# ============================================================1617from __future__ import annotations1819import uuid20from pathlib import Path21from typing import TYPE_CHECKING2223import pyarrow as pa24import pyarrow.dataset as ds25import pyarrow.parquet as pq26import pandas as pd2728from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL29from src.utils.logging import get_logger3031if TYPE_CHECKING:32 from src.storage.database import DatabaseManager3334logger = get_logger("storage.parquet")353637class ParquetManager:3839 def __init__(self, base_dir: Path | str | None = None) -> None:40 self.base_dir = Path(base_dir) if base_dir else DATA_DIRS["parquet"]41 self.base_dir.mkdir(parents=True, exist_ok=True)4243 def _table_dir(self, table_name: str) -> Path:44 path = self.base_dir / table_name45 path.mkdir(parents=True, exist_ok=True)46 return path4748 def _write_dataset(49 self,50 table: pa.Table,51 table_name: str,52 partition_by: list[str] | None = None,53 ) -> None:54 out_dir = self._table_dir(table_name)55 pq.write_to_dataset(56 table,57 root_path=str(out_dir),58 partition_cols=partition_by,59 compression=PARQUET_COMPRESSION,60 compression_level=PARQUET_COMPRESSION_LEVEL,61 existing_data_behavior="overwrite_or_ignore",62 )63 logger.info(64 "Wrote %d rows to %s (partitions=%s)",65 table.num_rows,66 out_dir,67 partition_by,68 )6970 def write_filings(71 self,72 filings: list[dict],73 partition_by: list[str] | None = None,74 ) -> None:75 if partition_by is None:76 partition_by = ["form_type", "fiscal_year"]77 table = pa.Table.from_pylist(filings)78 self._write_dataset(table, "filings", partition_by)7980 def write_sections(81 self,82 sections: list[dict],83 partition_by: list[str] | None = None,84 ) -> None:85 if partition_by is None:86 partition_by = ["form_type", "fiscal_year"]87 table = pa.Table.from_pylist(sections)88 self._write_dataset(table, "filing_sections", partition_by)8990 def write_chunks(91 self,92 chunks: list[dict],93 partition_by: list[str] | None = None,94 ) -> None:95 if partition_by is None:96 partition_by = ["form_type", "fiscal_year"]97 table = pa.Table.from_pylist(chunks)98 self._write_dataset(table, "filing_chunks", partition_by)99100 def write_embeddings(self, embeddings: list[dict], model_name: str) -> None:101 for emb in embeddings:102 emb.setdefault("model_name", model_name)103 table = pa.Table.from_pylist(embeddings)104 self._write_dataset(table, "embeddings", partition_by=["model_name"])105106 def write_embeddings_table(self, table: pa.Table, model_name: str) -> Path:107 """Write one parquet file per call (no tiny-file explosion for large corpora)."""108 out_dir = self._table_dir("embeddings") / f"model_name={model_name}"109 out_dir.mkdir(parents=True, exist_ok=True)110 path = out_dir / f"emb-{uuid.uuid4().hex}.parquet"111 pq.write_table(table.drop_columns(["model_name"]) if "model_name" in table.column_names else table,112 str(path), compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL)113 return path114115 def read_table(self, table_name: str, filters: list | None = None) -> pa.Table:116 table_dir = self._table_dir(table_name)117 dataset = ds.dataset(str(table_dir), format="parquet", partitioning="hive")118 if filters:119 return dataset.to_table(filter=ds.field(*filters) if len(filters) == 3 else None)120 return dataset.to_table()121122 def read_df(self, table_name: str, filters: list | None = None) -> pd.DataFrame:123 return self.read_table(table_name, filters).to_pandas()124125 def export_to_parquet(self, db: DatabaseManager, table_name: str) -> None:126 df = db.fetch_df(f"SELECT * FROM {table_name}")127 table = pa.Table.from_pandas(df)128 out_dir = self._table_dir(table_name)129 pq.write_table(130 table,131 str(out_dir / f"{table_name}.parquet"),132 compression=PARQUET_COMPRESSION,133 compression_level=PARQUET_COMPRESSION_LEVEL,134 )135 logger.info("Exported table '%s' (%d rows) to parquet", table_name, len(df))136137 def import_from_parquet(self, db: DatabaseManager, table_name: str) -> None:138 table_dir = self._table_dir(table_name)139 dataset = ds.dataset(str(table_dir), format="parquet", partitioning="hive")140 arrow_table = dataset.to_table()141 df = arrow_table.to_pandas()142 db.conn.execute(f"INSERT OR REPLACE INTO {table_name} SELECT * FROM df")143 logger.info("Imported %d rows from parquet into table '%s'", len(df), table_name)144145 def get_dataset_info(self, table_name: str) -> dict:146 table_dir = self._table_dir(table_name)147 parquet_files = list(table_dir.rglob("*.parquet"))148 total_size = sum(f.stat().st_size for f in parquet_files)149150 row_count = 0151 if parquet_files:152 dataset = ds.dataset(str(table_dir), format="parquet", partitioning="hive")153 row_count = dataset.count_rows()154155 return {156 "table_name": table_name,157 "row_count": row_count,158 "file_count": len(parquet_files),159 "total_size_bytes": total_size,160 "total_size_mb": round(total_size / (1024 * 1024), 2),161 "path": str(table_dir),162 }163