Optimize eid exposure cache build

prepare_exposure_cache.py

@@ -85,11 +85,11 @@ def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]:
     return [col for col in columns if col in schema_names]
 
 
-def _read_parquet_columns(path: Path, columns: list[str]) -> pd.DataFrame:
+def _read_matching_parquet_rows(
-    return pd.read_parquet(path, columns=columns)
+    path: Path,
-
+    columns: list[str],
-
+    wanted: pd.DataFrame,
-def _parquet_row_count(path: Path) -> int:
+) -> pd.DataFrame:
     try:
         import pyarrow.parquet as pq
     except ImportError as exc:
@@ -97,7 +97,60 @@ def _parquet_row_count(path: Path) -> int:
             "prepare_exposure_cache.py requires pyarrow. Install requirements "
             "or run `pip install pyarrow`."
         ) from exc
-    return int(pq.ParquetFile(path).metadata.num_rows)
+
+    parquet_file = pq.ParquetFile(path)
+    available = [col for col in columns if col in set(parquet_file.schema.names)]
+    key_cols = ["eid", "onset_date", "token"]
+    missing_keys = [col for col in key_cols if col not in available]
+    if missing_keys:
+        raise ValueError(f"{path} is missing key columns: {missing_keys}")
+
+    wanted_keys = wanted[["eid", "onset_date", "token", "position"]].copy()
+    wanted_keys["eid"] = wanted_keys["eid"].astype(np.int64)
+    wanted_keys["token"] = wanted_keys["token"].astype(np.int64)
+    wanted_keys["onset_date"] = pd.to_datetime(
+        wanted_keys["onset_date"],
+        errors="coerce",
+    ).dt.normalize()
+
+    chunks: list[pd.DataFrame] = []
+    for row_group_idx in range(parquet_file.num_row_groups):
+        key_frame = parquet_file.read_row_group(
+            row_group_idx,
+            columns=key_cols,
+        ).to_pandas()
+        if key_frame.empty:
+            continue
+        key_frame = key_frame.copy()
+        key_frame["_row_in_group"] = np.arange(len(key_frame), dtype=np.int64)
+        key_frame["eid"] = key_frame["eid"].astype(np.int64)
+        key_frame["token"] = key_frame["token"].astype(np.int64)
+        key_frame["onset_date"] = pd.to_datetime(
+            key_frame["onset_date"],
+            errors="coerce",
+        ).dt.normalize()
+        matches = key_frame.merge(
+            wanted_keys,
+            on=key_cols,
+            how="inner",
+            sort=False,
+        )
+        if matches.empty:
+            continue
+
+        row_values = parquet_file.read_row_group(
+            row_group_idx,
+            columns=available,
+        ).to_pandas()
+        selected = row_values.iloc[
+            matches["_row_in_group"].to_numpy(dtype=np.int64)
+        ].copy()
+        selected["position"] = matches["position"].to_numpy(dtype=np.int64)
+        chunks.append(selected)
+
+    if not chunks:
+        return pd.DataFrame(columns=[*columns, "position"])
+    return pd.concat(chunks, ignore_index=True)
 
 
 def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: int) -> np.ndarray:
@@ -124,35 +177,26 @@ def _load_summary(
     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))
-
+    for row in summary.itertuples(index=False):
-    counts: list[int] = []
-    iterator = tqdm(
-        summary.itertuples(index=False),
-        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
     return summary
 
 
+def _load_label_token_map(labels_file: str | Path) -> dict[str, int]:
+    out: dict[str, int] = {}
+    with Path(labels_file).open(encoding="utf-8") as handle:
+        for idx, line in enumerate(handle):
+            parts = line.strip().split(" ")
+            if parts and parts[0]:
+                out[parts[0]] = idx + 2
+    return out
+
+
 def _load_sequence_rows(data_prefix: str) -> pd.DataFrame:
     event_data = np.load(f"{data_prefix}_event_data.npy")
     if event_data.ndim != 2 or event_data.shape[1] < 3:
@@ -207,6 +251,7 @@ def build_exposure_cache(
     exposure_dir: str | Path,
     output_dir: str | Path,
     data_prefix: str = "ukb",
+    labels_file: str | Path = "labels.csv",
     summary_file: str = "summary.csv",
     overwrite: bool = False,
     show_progress: bool = True,
@@ -232,15 +277,29 @@ def build_exposure_cache(
             f"{output_dir} already contains exposure cache files; pass --overwrite"
         )
 
+    sequence_rows = _load_sequence_rows(data_prefix)
+    n_rows = len(sequence_rows)
+    if n_rows == 0:
+        raise ValueError(f"{data_prefix}_event_data.npy contains no real disease events")
+
+    label_token_map = _load_label_token_map(labels_file)
     summary = _load_summary(
         exposure_dir,
         summary_file,
         show_progress=show_progress,
     )
-    sequence_rows = _load_sequence_rows(data_prefix)
+    summary["raw_token"] = summary["label_code"].map(label_token_map)
-    n_rows = len(sequence_rows)
+    needed_tokens = set(sequence_rows["token"].astype(np.int64).unique().tolist())
-    if n_rows == 0:
+    summary = summary[
-        raise ValueError(f"{data_prefix}_event_data.npy contains no real disease events")
+        summary["raw_token"].notna()
+        & summary["raw_token"].astype(np.int64).isin(needed_tokens)
+    ].copy()
+    summary["raw_token"] = summary["raw_token"].astype(np.int64)
+    if summary.empty:
+        raise ValueError(
+            "No exposure summary rows match disease tokens in "
+            f"{data_prefix}_event_data.npy. Check --summary-file and --labels-file."
+        )
 
     eid_path = output_dir / "exposure_eid.npy"
     token_path = output_dir / "exposure_token.npy"
@@ -300,6 +359,10 @@ def build_exposure_cache(
     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
 
         daily_read_cols = [
             "eid",
@@ -315,64 +378,37 @@ def build_exposure_cache(
             *_safe_columns(monthly_file, monthly_cols),
             *_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]),
         ]
-        daily_df = _read_parquet_columns(daily_file, daily_read_cols)
+        daily_df = _read_matching_parquet_rows(daily_file, daily_read_cols, wanted)
-        monthly_df = _read_parquet_columns(monthly_file, monthly_read_cols)
+        monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted)
-        if len(daily_df) != len(monthly_df):
+        if daily_df.empty:
-            raise ValueError(
-                f"Daily/monthly row count mismatch for {row.label_code}: "
-                f"{len(daily_df)} vs {len(monthly_df)}"
-            )
-
-        daily_df = daily_df.copy()
-        monthly_df = monthly_df.copy()
-        daily_df["_source_row"] = np.arange(len(daily_df), dtype=np.int64)
-        daily_df["onset_date"] = pd.to_datetime(
-            daily_df["onset_date"],
-            errors="coerce",
-        ).dt.normalize()
-        monthly_df["onset_date"] = pd.to_datetime(
-            monthly_df["onset_date"],
-            errors="coerce",
-        ).dt.normalize()
-        monthly_df = monthly_df.set_index(["eid", "onset_date", "token"]).reindex(
-            pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
-        ).reset_index()
-
-        tokens = daily_df["token"].dropna().astype(np.int64).unique()
-        wanted = pd.concat(
-            [wanted_by_token[int(token)] for token in tokens if int(token) in wanted_by_token],
-            ignore_index=True,
-        ) if len(tokens) else pd.DataFrame()
-        if wanted.empty:
             continue
 
-        matches = daily_df[["eid", "onset_date", "token", "_source_row"]].merge(
+        common_positions = np.intersect1d(
-            wanted[["eid", "onset_date", "token", "position"]],
+            daily_df["position"].to_numpy(dtype=np.int64),
-            on=["eid", "onset_date", "token"],
+            monthly_df["position"].to_numpy(dtype=np.int64),
-            how="inner",
-            sort=False,
         )
-        if matches.empty:
+        if len(common_positions) == 0:
             continue
 
-        source_rows = matches["_source_row"].to_numpy(dtype=np.int64)
+        daily_df = daily_df.set_index("position").loc[common_positions].reset_index()
-        positions = matches["position"].to_numpy(dtype=np.int64)
+        monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index()
+        positions = common_positions.astype(np.int64)
         daily_mm[positions] = _reshape_window(
-            daily_df.iloc[source_rows],
+            daily_df,
             daily_cols,
             DAILY_LENGTH,
             len(DAILY_CHANNELS),
         )
         monthly_mm[positions] = _reshape_window(
-            monthly_df.iloc[source_rows],
+            monthly_df,
             monthly_cols,
             MONTHLY_LENGTH,
             len(MONTHLY_CHANNELS),
         )
-        quality_mm[positions, 0] = daily_df.iloc[source_rows].get("n_days_nonmissing", np.nan)
+        quality_mm[positions, 0] = daily_df.get("n_days_nonmissing", np.nan)
-        quality_mm[positions, 1] = daily_df.iloc[source_rows].get("n_rh_days_nonmissing", np.nan)
+        quality_mm[positions, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
-        quality_mm[positions, 2] = monthly_df.iloc[source_rows].get("n_months_nonmissing", np.nan)
+        quality_mm[positions, 2] = monthly_df.get("n_months_nonmissing", np.nan)
-        quality_mm[positions, 3] = monthly_df.iloc[source_rows].get("n_rh_months_nonmissing", np.nan)
+        quality_mm[positions, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
         matched[positions] = True
 
     daily_mm.flush()
@@ -383,6 +419,7 @@ def build_exposure_cache(
         "storage": "eid_sequence_npy",
         "source_dir": str(exposure_dir.resolve()),
         "data_prefix": data_prefix,
+        "labels_file": str(Path(labels_file).resolve()),
         "n_rows": int(n_rows),
         "matched_rows": int(matched.sum()),
         "missing_rows": int((~matched).sum()),
@@ -404,6 +441,7 @@ def main() -> None:
     parser.add_argument("--exposure-dir", required=True)
     parser.add_argument("--output-dir", default="ukb_exposure_cache")
     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(
         "--no-progress",
@@ -417,6 +455,7 @@ def main() -> None:
         exposure_dir=args.exposure_dir,
         output_dir=args.output_dir,
         data_prefix=args.data_prefix,
+        labels_file=args.labels_file,
         summary_file=args.summary_file,
         overwrite=args.overwrite,
         show_progress=not args.no_progress,