Parallelize exposure cache reads

prepare_exposure_cache.py

@@ -36,7 +36,9 @@ converts back to these raw tokens before reading this cache.
 from __future__ import annotations
 
 import argparse
+from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
 import json
+import os
 from pathlib import Path
 from typing import Iterable
 
@@ -162,6 +164,75 @@ 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 _quality_column(df: pd.DataFrame, name: str, n_rows: int) -> np.ndarray:
+    if name not in df:
+        return np.full(n_rows, np.nan, dtype=np.float32)
+    return df[name].to_numpy(dtype=np.float32, copy=True)
+
+
+def _process_exposure_task(task: dict) -> dict:
+    daily_file = Path(task["daily_path"])
+    monthly_file = Path(task["monthly_path"])
+    wanted = task["wanted"]
+    daily_cols = task["daily_cols"]
+    monthly_cols = task["monthly_cols"]
+
+    daily_read_cols = [
+        "eid",
+        "onset_date",
+        "token",
+        *_safe_columns(daily_file, daily_cols),
+        *_safe_columns(daily_file, ["n_days_nonmissing", "n_rh_days_nonmissing"]),
+    ]
+    monthly_read_cols = [
+        "eid",
+        "onset_date",
+        "token",
+        *_safe_columns(monthly_file, monthly_cols),
+        *_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]),
+    ]
+    daily_df = _read_matching_parquet_rows(daily_file, daily_read_cols, wanted)
+    monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted)
+    if daily_df.empty:
+        return {"positions": np.empty(0, dtype=np.int64)}
+
+    common_positions = np.intersect1d(
+        daily_df["position"].to_numpy(dtype=np.int64),
+        monthly_df["position"].to_numpy(dtype=np.int64),
+    )
+    if len(common_positions) == 0:
+        return {"positions": np.empty(0, dtype=np.int64)}
+
+    daily_df = daily_df.set_index("position").loc[common_positions].reset_index()
+    monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index()
+    n_match = len(common_positions)
+    quality = np.stack(
+        [
+            _quality_column(daily_df, "n_days_nonmissing", n_match),
+            _quality_column(daily_df, "n_rh_days_nonmissing", n_match),
+            _quality_column(monthly_df, "n_months_nonmissing", n_match),
+            _quality_column(monthly_df, "n_rh_months_nonmissing", n_match),
+        ],
+        axis=1,
+    )
+    return {
+        "positions": common_positions.astype(np.int64),
+        "daily": _reshape_window(
+            daily_df,
+            daily_cols,
+            DAILY_LENGTH,
+            len(DAILY_CHANNELS),
+        ),
+        "monthly": _reshape_window(
+            monthly_df,
+            monthly_cols,
+            MONTHLY_LENGTH,
+            len(MONTHLY_CHANNELS),
+        ),
+        "quality": quality,
+    }
+
+
 def _load_summary(
     exposure_dir: Path,
     summary_file: str,
@@ -257,6 +328,8 @@ def build_exposure_cache(
     data_prefix: str = "ukb",
     labels_file: str | Path = "labels.csv",
     summary_file: str = "summary.csv",
+    workers: int = 1,
+    max_in_flight: int = 0,
     overwrite: bool = False,
     show_progress: bool = True,
 ) -> int:
@@ -359,71 +432,83 @@ def build_exposure_cache(
     }
     write_offset = 0
 
-    iterator = tqdm(
+    tasks: list[dict] = []
-        summary.itertuples(index=False),
+    for row in summary.itertuples(index=False):
-        total=len(summary),
-        desc="Writing eid-sequence exposure cache",
-        unit="file",
-        disable=not show_progress,
-    )
-    for row in iterator:
-        daily_file = Path(row.daily_path)
-        monthly_file = Path(row.monthly_path)
         token = int(row.raw_token)
         wanted = wanted_by_token.get(token)
         if wanted is None or wanted.empty:
             continue
-
+        tasks.append(
-        daily_read_cols = [
+            {
-            "eid",
+                "daily_path": str(row.daily_path),
-            "onset_date",
+                "monthly_path": str(row.monthly_path),
-            "token",
+                "wanted": wanted,
-            *_safe_columns(daily_file, daily_cols),
+                "daily_cols": daily_cols,
-            *_safe_columns(daily_file, ["n_days_nonmissing", "n_rh_days_nonmissing"]),
+                "monthly_cols": monthly_cols,
-        ]
+            }
-        monthly_read_cols = [
-            "eid",
-            "onset_date",
-            "token",
-            *_safe_columns(monthly_file, monthly_cols),
-            *_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]),
-        ]
-        daily_df = _read_matching_parquet_rows(daily_file, daily_read_cols, wanted)
-        monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted)
-        if daily_df.empty:
-            continue
-
-        common_positions = np.intersect1d(
-            daily_df["position"].to_numpy(dtype=np.int64),
-            monthly_df["position"].to_numpy(dtype=np.int64),
         )
-        if len(common_positions) == 0:
-            continue
 
-        daily_df = daily_df.set_index("position").loc[common_positions].reset_index()
+    workers = max(1, int(workers))
-        monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index()
+    max_in_flight = int(max_in_flight)
-        positions = common_positions.astype(np.int64)
+
+    def write_result(result: dict) -> None:
+        nonlocal write_offset
+        positions = result["positions"]
+        if len(positions) == 0:
+            return
         n_match = len(positions)
         end_offset = write_offset + n_match
-        daily_mm[write_offset:end_offset] = _reshape_window(
+        daily_mm[write_offset:end_offset] = result["daily"]
-            daily_df,
+        monthly_mm[write_offset:end_offset] = result["monthly"]
-            daily_cols,
+        quality_mm[write_offset:end_offset] = result["quality"]
-            DAILY_LENGTH,
-            len(DAILY_CHANNELS),
-        )
-        monthly_mm[write_offset:end_offset] = _reshape_window(
-            monthly_df,
-            monthly_cols,
-            MONTHLY_LENGTH,
-            len(MONTHLY_CHANNELS),
-        )
-        quality_mm[write_offset:end_offset, 0] = daily_df.get("n_days_nonmissing", np.nan)
-        quality_mm[write_offset:end_offset, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
-        quality_mm[write_offset:end_offset, 2] = monthly_df.get("n_months_nonmissing", np.nan)
-        quality_mm[write_offset:end_offset, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
         row_index_mm[positions] = np.arange(write_offset, end_offset, dtype=np.int64)
         write_offset = end_offset
 
+    if workers == 1:
+        iterator = tqdm(
+            map(_process_exposure_task, tasks),
+            total=len(tasks),
+            desc="Reading exposure parquet and writing cache",
+            unit="file",
+            disable=not show_progress,
+        )
+        for result in iterator:
+            write_result(result)
+    else:
+        with ProcessPoolExecutor(max_workers=workers) as executor:
+            task_iter = iter(tasks)
+            iterator = tqdm(
+                total=len(tasks),
+                desc=f"Reading exposure parquet ({workers} workers)",
+                unit="file",
+                disable=not show_progress,
+            )
+            if max_in_flight <= 0:
+                in_flight = {
+                    executor.submit(_process_exposure_task, task)
+                    for task in task_iter
+                }
+                while in_flight:
+                    done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED)
+                    for future in done:
+                        write_result(future.result())
+                        iterator.update(1)
+            else:
+                in_flight = {
+                    executor.submit(_process_exposure_task, task)
+                    for task in [next(task_iter, None) for _ in range(max_in_flight)]
+                    if task is not None
+                }
+                while in_flight:
+                    done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED)
+                    for future in done:
+                        write_result(future.result())
+                        iterator.update(1)
+                        task = next(task_iter, None)
+                        if task is not None:
+                            in_flight.add(executor.submit(_process_exposure_task, task))
+            iterator.close()
+
     row_index_mm.flush()
     daily_mm.flush()
     monthly_mm.flush()
@@ -460,6 +545,24 @@ def main() -> None:
     parser.add_argument("--data-prefix", default="ukb")
     parser.add_argument("--labels-file", default="labels.csv")
     parser.add_argument("--summary-file", default="summary.csv")
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=max(1, os.cpu_count() or 1),
+        help=(
+            "Worker processes for parquet reading and window extraction. "
+            "The main process remains the only writer to the output memmaps."
+        ),
+    )
+    parser.add_argument(
+        "--max-in-flight",
+        type=int,
+        default=0,
+        help=(
+            "Maximum submitted parquet tasks waiting/running at once. "
+            "Use 0 to submit all tasks, which is the default for high-memory servers."
+        ),
+    )
     parser.add_argument(
         "--no-progress",
         action="store_true",
@@ -474,6 +577,8 @@ def main() -> None:
         data_prefix=args.data_prefix,
         labels_file=args.labels_file,
         summary_file=args.summary_file,
+        workers=args.workers,
+        max_in_flight=args.max_in_flight,
         overwrite=args.overwrite,
         show_progress=not args.no_progress,
     )