Use parquet index for exposure cache

prepare_exposure_cache.py

@@ -1,11 +1,31 @@
-"""Build a random-access exposure cache from disease-level parquet files.
+"""Build a random-access exposure index/cache from disease-level parquet files.
 
 The README-described exposure dataset is stored as one daily and one monthly
 parquet file per disease. That layout is good for disease-specific analysis but
 too expensive for mini-batch training, where we need exposure windows aligned
 to arbitrary event sequences.
 
-This script converts those parquet files into a compact directory:
+By default this script builds a lightweight parquet index. It does not copy the
+daily/monthly exposure windows; it only records which source parquet file,
+row-group, and row each exposure event lives in. Dataset loading then reads the
+original parquet row groups on demand.
+
+The default index directory contains:
+
+    exposure_eid.npy                int64 eid per exposure row
+    exposure_token.npy              int32 raw disease token per exposure row
+    exposure_onset_date.npy         datetime64[D] onset date per exposure row
+    exposure_daily_file_id.npy      int32 source daily file id per row
+    exposure_daily_row_group.npy    int32 source daily row group per row
+    exposure_daily_row_in_group.npy int32 row offset inside daily row group
+    exposure_monthly_file_id.npy    int32 source monthly file id per row
+    exposure_monthly_row_group.npy  int32 source monthly row group per row
+    exposure_monthly_row_in_group.npy
+                                    int32 row offset inside monthly row group
+    exposure_manifest.json          metadata and source parquet paths
+
+For faster but much larger training storage, ``--mode dense`` materializes a
+full dense numpy cache:
 
     exposure_keys.npy       uint64 legacy keys, key = (eid << 16) | raw_token
     exposure_eid.npy        int64 eid per exposure row
@@ -91,6 +111,28 @@ def _read_parquet_columns(path: Path, columns: list[str]) -> pd.DataFrame:
     return pd.read_parquet(path, columns=columns)
 
 
+def _row_group_positions(path: Path) -> tuple[np.ndarray, np.ndarray]:
+    """Return row_group and row-in-group vectors for every parquet row."""
+    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
+
+    parquet_file = pq.ParquetFile(path)
+    row_groups: list[np.ndarray] = []
+    row_offsets: list[np.ndarray] = []
+    for row_group_idx in range(parquet_file.num_row_groups):
+        n = parquet_file.metadata.row_group(row_group_idx).num_rows
+        row_groups.append(np.full(n, row_group_idx, dtype=np.int32))
+        row_offsets.append(np.arange(n, dtype=np.int32))
+    if not row_groups:
+        return np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
+    return np.concatenate(row_groups), np.concatenate(row_offsets)
+
+
 def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: int) -> np.ndarray:
     arr = df.reindex(columns=cols).to_numpy(dtype=np.float32, copy=True)
     return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1)
@@ -102,6 +144,190 @@ def _count_rows(summary: pd.DataFrame) -> int:
     return int(sum(pd.read_parquet(path, columns=["eid"]).shape[0] for path in summary["daily_path"]))
 
 
+def build_exposure_index(
+    *,
+    exposure_dir: str | Path,
+    output_dir: str | Path,
+    summary_file: str = "summary.csv",
+    overwrite: bool = False,
+) -> 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 = [
+        output_dir / "exposure_eid.npy",
+        output_dir / "exposure_token.npy",
+        output_dir / "exposure_onset_date.npy",
+        output_dir / "exposure_daily_file_id.npy",
+        output_dir / "exposure_daily_row_group.npy",
+        output_dir / "exposure_daily_row_in_group.npy",
+        output_dir / "exposure_monthly_file_id.npy",
+        output_dir / "exposure_monthly_row_group.npy",
+        output_dir / "exposure_monthly_row_in_group.npy",
+        output_dir / "exposure_manifest.json",
+    ]
+    if any(path.exists() for path in output_paths) and not overwrite:
+        raise FileExistsError(
+            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)
+    eids_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_eid.npy", mode="w+", dtype=np.int64, shape=(n_rows,)
+    )
+    tokens_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_token.npy", mode="w+", dtype=np.int32, shape=(n_rows,)
+    )
+    onset_dates_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_onset_date.npy",
+        mode="w+",
+        dtype="datetime64[D]",
+        shape=(n_rows,),
+    )
+    daily_file_id_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_daily_file_id.npy",
+        mode="w+",
+        dtype=np.int32,
+        shape=(n_rows,),
+    )
+    daily_row_group_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_daily_row_group.npy",
+        mode="w+",
+        dtype=np.int32,
+        shape=(n_rows,),
+    )
+    daily_row_in_group_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_daily_row_in_group.npy",
+        mode="w+",
+        dtype=np.int32,
+        shape=(n_rows,),
+    )
+    monthly_file_id_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_monthly_file_id.npy",
+        mode="w+",
+        dtype=np.int32,
+        shape=(n_rows,),
+    )
+    monthly_row_group_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_monthly_row_group.npy",
+        mode="w+",
+        dtype=np.int32,
+        shape=(n_rows,),
+    )
+    monthly_row_in_group_mm = np.lib.format.open_memmap(
+        output_dir / "exposure_monthly_row_in_group.npy",
+        mode="w+",
+        dtype=np.int32,
+        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}")
+
+        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)}"
+            )
+        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]")
+        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
+        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
+
+    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."
+        )
+
+    for memmap in (
+        eids_mm,
+        tokens_mm,
+        onset_dates_mm,
+        daily_file_id_mm,
+        daily_row_group_mm,
+        daily_row_in_group_mm,
+        monthly_file_id_mm,
+        monthly_row_group_mm,
+        monthly_row_in_group_mm,
+    ):
+        memmap.flush()
+
+    manifest = {
+        "storage": "parquet_index",
+        "source_dir": str(exposure_dir.resolve()),
+        "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_shape_per_row": [DAILY_LENGTH, len(DAILY_CHANNELS)],
+        "daily_channels": list(DAILY_CHANNELS),
+        "monthly_shape_per_row": [MONTHLY_LENGTH, len(MONTHLY_CHANNELS)],
+        "monthly_channels": list(MONTHLY_CHANNELS),
+        "raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2",
+    }
+    (output_dir / "exposure_manifest.json").write_text(
+        json.dumps(manifest, indent=2),
+        encoding="utf-8",
+    )
+    return int(n_rows)
+
+
 def build_exposure_cache(
     *,
     exposure_dir: str | Path,
@@ -264,6 +490,7 @@ def build_exposure_cache(
         )
 
     manifest = {
+        "storage": "dense_npy",
         "source_dir": str(exposure_dir),
         "n_rows": int(n_rows),
         "legacy_key": "(eid << 16) | raw_token",
@@ -285,15 +512,33 @@ def main() -> None:
     parser.add_argument("--exposure-dir", required=True)
     parser.add_argument("--output-dir", default="ukb_exposure_cache")
     parser.add_argument("--summary-file", default="summary.csv")
+    parser.add_argument(
+        "--mode",
+        choices=("index", "dense"),
+        default="index",
+        help=(
+            "index writes only lightweight parquet row pointers; dense copies "
+            "all exposure windows into numpy memmaps."
+        ),
+    )
     parser.add_argument("--overwrite", action="store_true")
     args = parser.parse_args()
-    n_rows = build_exposure_cache(
-        exposure_dir=args.exposure_dir,
-        output_dir=args.output_dir,
-        summary_file=args.summary_file,
-        overwrite=args.overwrite,
-    )
-    print(f"Wrote {n_rows:,} exposure rows to {args.output_dir}")
+    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,
+        )
+        print(f"Wrote {n_rows:,} exposure row pointers to {args.output_dir}")
+    else:
+        n_rows = build_exposure_cache(
+            exposure_dir=args.exposure_dir,
+            output_dir=args.output_dir,
+            summary_file=args.summary_file,
+            overwrite=args.overwrite,
+        )
+        print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}")
 
 
 if __name__ == "__main__":