diff --git a/prepare_exposure_cache.py b/prepare_exposure_cache.py index b7acedb..453cc76 100644 --- a/prepare_exposure_cache.py +++ b/prepare_exposure_cache.py @@ -51,12 +51,15 @@ model. from __future__ import annotations import argparse +from concurrent.futures import ProcessPoolExecutor, as_completed import json +import os from pathlib import Path from typing import Iterable import numpy as np import pandas as pd +from tqdm.auto import tqdm DAILY_LENGTH = 1826 @@ -111,6 +114,17 @@ def _read_parquet_columns(path: Path, columns: list[str]) -> pd.DataFrame: return pd.read_parquet(path, columns=columns) +def _parquet_row_count(path: Path) -> int: + try: + import pyarrow.parquet as pq + except ImportError as exc: + raise ImportError( + "prepare_exposure_cache.py requires pyarrow. Install requirements " + "or run `pip install pyarrow`." + ) from exc + return int(pq.ParquetFile(path).metadata.num_rows) + + def _row_group_positions(path: Path) -> tuple[np.ndarray, np.ndarray]: """Return row_group and row-in-group vectors for every parquet row.""" try: @@ -138,10 +152,96 @@ def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1) -def _count_rows(summary: pd.DataFrame) -> int: - if "n_cases" in summary.columns: - return int(summary["n_cases"].sum()) - return int(sum(pd.read_parquet(path, columns=["eid"]).shape[0] for path in summary["daily_path"])) +def _load_summary( + exposure_dir: Path, + summary_file: str, + *, + show_progress: bool, +) -> pd.DataFrame: + summary_path = exposure_dir / summary_file + if not summary_path.is_file(): + raise FileNotFoundError(f"summary.csv not found: {summary_path}") + + summary = pd.read_csv(summary_path) + required = {"label_code", "daily_file", "monthly_file"} + missing = required - set(summary.columns) + if missing: + raise ValueError(f"{summary_path} is missing columns: {sorted(missing)}") + + summary = summary.copy() + summary["daily_path"] = summary["daily_file"].map(lambda name: exposure_dir / str(name)) + summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name)) + + counts: list[int] = [] + iterator = summary.itertuples(index=False) + iterator = tqdm( + iterator, + total=len(summary), + desc="Counting exposure rows", + unit="file", + disable=not show_progress, + ) + for row in iterator: + daily_file = Path(row.daily_path) + monthly_file = Path(row.monthly_path) + if not daily_file.is_file(): + raise FileNotFoundError(f"Missing daily parquet: {daily_file}") + if not monthly_file.is_file(): + raise FileNotFoundError(f"Missing monthly parquet: {monthly_file}") + daily_count = _parquet_row_count(daily_file) + monthly_count = _parquet_row_count(monthly_file) + if daily_count != monthly_count: + raise ValueError( + f"Daily/monthly row count mismatch for {row.label_code}: " + f"{daily_count} vs {monthly_count}" + ) + counts.append(daily_count) + + summary["n_rows"] = counts + summary["offset"] = np.cumsum([0, *counts[:-1]], dtype=np.int64) + return summary + + +def _process_index_file_pair(task: tuple[int, str, str, str]) -> dict: + file_id, label_code, daily_path, monthly_path = task + daily_file = Path(daily_path) + monthly_file = Path(monthly_path) + + daily_df = _read_parquet_columns(daily_file, ["eid", "onset_date", "token"]) + monthly_df = _read_parquet_columns(monthly_file, ["eid", "onset_date", "token"]) + if len(daily_df) != len(monthly_df): + raise ValueError( + f"Daily/monthly row count mismatch for {label_code}: " + f"{len(daily_df)} vs {len(monthly_df)}" + ) + + daily_rg, daily_row = _row_group_positions(daily_file) + monthly_rg_all, monthly_row_all = _row_group_positions(monthly_file) + n = len(daily_df) + if len(daily_rg) != n or len(monthly_rg_all) != n: + raise ValueError(f"Parquet row-group metadata row count mismatch for {label_code}") + + daily_index = pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]]) + monthly_index = pd.MultiIndex.from_frame(monthly_df[["eid", "onset_date", "token"]]) + monthly_pos = monthly_index.get_indexer(daily_index) + if np.any(monthly_pos < 0): + raise ValueError(f"Monthly parquet is missing daily exposure keys for {label_code}") + + return { + "file_id": int(file_id), + "label_code": label_code, + "n_rows": int(n), + "eid": daily_df["eid"].to_numpy(dtype=np.int64), + "token": daily_df["token"].to_numpy(dtype=np.int32), + "onset_date": pd.to_datetime( + daily_df["onset_date"], + errors="coerce", + ).to_numpy(dtype="datetime64[D]"), + "daily_row_group": daily_rg, + "daily_row_in_group": daily_row, + "monthly_row_group": monthly_rg_all[monthly_pos], + "monthly_row_in_group": monthly_row_all[monthly_pos], + } def build_exposure_index( @@ -150,12 +250,11 @@ def build_exposure_index( output_dir: str | Path, summary_file: str = "summary.csv", overwrite: bool = False, + workers: int = 1, + show_progress: bool = True, ) -> int: exposure_dir = Path(exposure_dir) output_dir = Path(output_dir) - summary_path = exposure_dir / summary_file - if not summary_path.is_file(): - raise FileNotFoundError(f"summary.csv not found: {summary_path}") output_dir.mkdir(parents=True, exist_ok=True) output_paths = [ @@ -175,16 +274,12 @@ def build_exposure_index( f"{output_dir} already contains exposure index files; pass --overwrite" ) - summary = pd.read_csv(summary_path) - required = {"label_code", "daily_file", "monthly_file"} - missing = required - set(summary.columns) - if missing: - raise ValueError(f"{summary_path} is missing columns: {sorted(missing)}") - summary = summary.copy() - summary["daily_path"] = summary["daily_file"].map(lambda name: exposure_dir / str(name)) - summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name)) - - n_rows = _count_rows(summary) + summary = _load_summary( + exposure_dir, + summary_file, + show_progress=show_progress, + ) + n_rows = int(summary["n_rows"].sum()) eids_mm = np.lib.format.open_memmap( output_dir / "exposure_eid.npy", mode="w+", dtype=np.int64, shape=(n_rows,) ) @@ -234,65 +329,63 @@ def build_exposure_index( shape=(n_rows,), ) - daily_files: list[str] = [] - monthly_files: list[str] = [] - offset = 0 - for file_id, row in enumerate(summary.itertuples(index=False)): - daily_file = Path(row.daily_path) - monthly_file = Path(row.monthly_path) - if not daily_file.is_file(): - raise FileNotFoundError(f"Missing daily parquet: {daily_file}") - if not monthly_file.is_file(): - raise FileNotFoundError(f"Missing monthly parquet: {monthly_file}") + tasks = [ + ( + int(file_id), + str(row.label_code), + str(Path(row.daily_path)), + str(Path(row.monthly_path)), + ) + for file_id, row in enumerate(summary.itertuples(index=False)) + ] + workers = max(1, int(workers)) - daily_df = _read_parquet_columns(daily_file, ["eid", "onset_date", "token"]) - monthly_df = _read_parquet_columns(monthly_file, ["eid", "onset_date", "token"]) - if len(daily_df) != len(monthly_df): - raise ValueError( - f"Daily/monthly row count mismatch for {row.label_code}: " - f"{len(daily_df)} vs {len(monthly_df)}" + def write_result(result: dict) -> None: + file_id = int(result["file_id"]) + row = summary.iloc[file_id] + offset = int(row.offset) + expected_n = int(row.n_rows) + n = int(result["n_rows"]) + if n != expected_n: + raise RuntimeError( + f"Expected {expected_n} rows for {result['label_code']} " + f"from metadata but indexed {n}" ) - daily_rg, daily_row = _row_group_positions(daily_file) - monthly_rg_all, monthly_row_all = _row_group_positions(monthly_file) - n = len(daily_df) - if len(daily_rg) != n or len(monthly_rg_all) != n: - raise ValueError(f"Parquet row-group metadata row count mismatch for {row.label_code}") - - daily_index = pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]]) - monthly_index = pd.MultiIndex.from_frame(monthly_df[["eid", "onset_date", "token"]]) - monthly_pos = monthly_index.get_indexer(daily_index) - if np.any(monthly_pos < 0): - raise ValueError( - f"Monthly parquet is missing daily exposure keys for {row.label_code}" - ) - monthly_rg = monthly_rg_all[monthly_pos] - monthly_row = monthly_row_all[monthly_pos] - end = offset + n if end > n_rows: raise RuntimeError("Exposure index row count exceeded preallocated size") - eids_mm[offset:end] = daily_df["eid"].to_numpy(dtype=np.int64) - tokens_mm[offset:end] = daily_df["token"].to_numpy(dtype=np.int32) - onset_dates_mm[offset:end] = pd.to_datetime( - daily_df["onset_date"], - errors="coerce", - ).to_numpy(dtype="datetime64[D]") + eids_mm[offset:end] = result["eid"] + tokens_mm[offset:end] = result["token"] + onset_dates_mm[offset:end] = result["onset_date"] daily_file_id_mm[offset:end] = file_id - daily_row_group_mm[offset:end] = daily_rg - daily_row_in_group_mm[offset:end] = daily_row + daily_row_group_mm[offset:end] = result["daily_row_group"] + daily_row_in_group_mm[offset:end] = result["daily_row_in_group"] monthly_file_id_mm[offset:end] = file_id - monthly_row_group_mm[offset:end] = monthly_rg - monthly_row_in_group_mm[offset:end] = monthly_row - daily_files.append(str(daily_file.resolve())) - monthly_files.append(str(monthly_file.resolve())) - offset = end + monthly_row_group_mm[offset:end] = result["monthly_row_group"] + monthly_row_in_group_mm[offset:end] = result["monthly_row_in_group"] - if offset != n_rows: - raise RuntimeError( - f"Expected {n_rows} rows from summary but indexed {offset}. " - "Regenerate summary.csv or remove n_cases before building." - ) + if workers == 1: + iterator = map(_process_index_file_pair, tasks) + for result in tqdm( + iterator, + total=len(tasks), + desc="Indexing exposure parquet", + unit="file", + disable=not show_progress, + ): + write_result(result) + else: + with ProcessPoolExecutor(max_workers=workers) as executor: + futures = [executor.submit(_process_index_file_pair, task) for task in tasks] + for future in tqdm( + as_completed(futures), + total=len(futures), + desc=f"Indexing exposure parquet ({workers} workers)", + unit="file", + disable=not show_progress, + ): + write_result(future.result()) for memmap in ( eids_mm, @@ -313,8 +406,12 @@ def build_exposure_index( "n_rows": int(n_rows), "alignment_key": "(eid, raw_token, onset_date - date_of_birth)", "requires_basic_info_column": "date_of_birth", - "daily_files": daily_files, - "monthly_files": monthly_files, + "daily_files": [ + str(Path(path).resolve()) for path in summary["daily_path"].tolist() + ], + "monthly_files": [ + str(Path(path).resolve()) for path in summary["monthly_path"].tolist() + ], "daily_shape_per_row": [DAILY_LENGTH, len(DAILY_CHANNELS)], "daily_channels": list(DAILY_CHANNELS), "monthly_shape_per_row": [MONTHLY_LENGTH, len(MONTHLY_CHANNELS)], @@ -334,13 +431,10 @@ def build_exposure_cache( output_dir: str | Path, summary_file: str = "summary.csv", overwrite: bool = False, + show_progress: bool = True, ) -> int: exposure_dir = Path(exposure_dir) output_dir = Path(output_dir) - summary_path = exposure_dir / summary_file - if not summary_path.is_file(): - raise FileNotFoundError(f"summary.csv not found: {summary_path}") - output_dir.mkdir(parents=True, exist_ok=True) keys_path = output_dir / "exposure_keys.npy" eid_path = output_dir / "exposure_eid.npy" @@ -365,16 +459,12 @@ def build_exposure_cache( f"{output_dir} already contains exposure cache files; pass --overwrite" ) - summary = pd.read_csv(summary_path) - required = {"label_code", "daily_file", "monthly_file"} - missing = required - set(summary.columns) - if missing: - raise ValueError(f"{summary_path} is missing columns: {sorted(missing)}") - summary = summary.copy() - summary["daily_path"] = summary["daily_file"].map(lambda name: exposure_dir / str(name)) - summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name)) - - n_rows = _count_rows(summary) + summary = _load_summary( + exposure_dir, + summary_file, + show_progress=show_progress, + ) + n_rows = int(summary["n_rows"].sum()) keys = np.lib.format.open_memmap(keys_path, mode="w+", dtype=np.uint64, shape=(n_rows,)) eids_mm = np.lib.format.open_memmap(eid_path, mode="w+", dtype=np.int64, shape=(n_rows,)) tokens_mm = np.lib.format.open_memmap(token_path, mode="w+", dtype=np.int32, shape=(n_rows,)) @@ -407,7 +497,14 @@ def build_exposure_cache( monthly_cols = _monthly_columns() offset = 0 - for row in summary.itertuples(index=False): + rows = tqdm( + summary.itertuples(index=False), + total=len(summary), + desc="Materializing dense exposure cache", + unit="file", + disable=not show_progress, + ) + for row in rows: daily_file = Path(row.daily_path) monthly_file = Path(row.monthly_path) if not daily_file.is_file(): @@ -486,7 +583,7 @@ def build_exposure_cache( keys = np.lib.format.open_memmap(keys_path, mode="r+", dtype=np.uint64, shape=(offset,)) raise RuntimeError( f"Expected {n_rows} rows from summary but wrote {offset}. " - "Regenerate summary.csv or remove n_cases before building." + "Check parquet metadata and regenerate summary.csv before building." ) manifest = { @@ -521,14 +618,31 @@ def main() -> None: "all exposure windows into numpy memmaps." ), ) + parser.add_argument( + "--workers", + type=int, + default=max(1, min(8, (os.cpu_count() or 1))), + help=( + "Number of worker processes for --mode index. Dense mode remains " + "single-writer to avoid concurrent writes to the same memmap." + ), + ) + parser.add_argument( + "--no-progress", + action="store_true", + help="Disable tqdm progress bars.", + ) parser.add_argument("--overwrite", action="store_true") args = parser.parse_args() + show_progress = not args.no_progress if args.mode == "index": n_rows = build_exposure_index( exposure_dir=args.exposure_dir, output_dir=args.output_dir, summary_file=args.summary_file, overwrite=args.overwrite, + workers=args.workers, + show_progress=show_progress, ) print(f"Wrote {n_rows:,} exposure row pointers to {args.output_dir}") else: @@ -537,6 +651,7 @@ def main() -> None: output_dir=args.output_dir, summary_file=args.summary_file, overwrite=args.overwrite, + show_progress=show_progress, ) print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}")