Parallelize exposure index preparation

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}")