Use parquet index for exposure cache

dataset.py

@@ -1,6 +1,8 @@
 # dataset.py
 from __future__ import annotations
 
+import json
+from collections import OrderedDict
 from pathlib import Path
 from typing import Dict, Iterable, List, Literal, Optional, Tuple
 
@@ -22,14 +24,41 @@ from targets import (
 ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
 DAILY_EXPOSURE_SHAPE = (1826, 4)
 MONTHLY_EXPOSURE_SHAPE = (241, 2)
+DAILY_EXPOSURE_CHANNELS = ("tmean", "tmax", "tmin", "rhmean")
+MONTHLY_EXPOSURE_CHANNELS = ("tmean", "rhmean")
+
+
+def _daily_exposure_columns() -> list[str]:
+    cols: list[str] = []
+    for name in DAILY_EXPOSURE_CHANNELS:
+        cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_EXPOSURE_SHAPE[0]))
+    return cols
+
+
+def _monthly_exposure_columns() -> list[str]:
+    cols: list[str] = []
+    for name in MONTHLY_EXPOSURE_CHANNELS:
+        cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_EXPOSURE_SHAPE[0]))
+    return cols
 
 
 class ExposureCache:
     """Random-access view over files produced by prepare_exposure_cache.py."""
 
-    def __init__(self, cache_dir: str | Path):
+    def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 16):
         cache_dir = Path(cache_dir)
         self.cache_dir = cache_dir
+        manifest_path = cache_dir / "exposure_manifest.json"
+        self.manifest = (
+            json.loads(manifest_path.read_text(encoding="utf-8"))
+            if manifest_path.is_file()
+            else {}
+        )
+        self.storage = self.manifest.get("storage", "dense_npy")
+        self._row_group_cache_size = int(row_group_cache_size)
+        self._row_group_cache: OrderedDict[tuple[str, int, int], pd.DataFrame] = OrderedDict()
+        self._parquet_files: dict[tuple[str, int], object] = {}
+        self._parquet_columns: dict[tuple[str, int], list[str]] = {}
         eid_path = cache_dir / "exposure_eid.npy"
         token_path = cache_dir / "exposure_token.npy"
         onset_date_path = cache_dir / "exposure_onset_date.npy"
@@ -42,11 +71,34 @@ class ExposureCache:
         self.eids = np.load(eid_path, mmap_mode="r")
         self.raw_tokens = np.load(token_path, mmap_mode="r")
         self.onset_dates = np.load(onset_date_path, mmap_mode="r")
+        self.daily = None
+        self.monthly = None
+        if self.storage == "dense_npy":
             self.daily = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r")
             self.monthly = np.load(cache_dir / "exposure_monthly.npy", mmap_mode="r")
+        elif self.storage == "parquet_index":
+            self.daily_file_ids = np.load(cache_dir / "exposure_daily_file_id.npy", mmap_mode="r")
+            self.daily_row_groups = np.load(cache_dir / "exposure_daily_row_group.npy", mmap_mode="r")
+            self.daily_row_in_groups = np.load(
+                cache_dir / "exposure_daily_row_in_group.npy", mmap_mode="r"
+            )
+            self.monthly_file_ids = np.load(
+                cache_dir / "exposure_monthly_file_id.npy", mmap_mode="r"
+            )
+            self.monthly_row_groups = np.load(
+                cache_dir / "exposure_monthly_row_group.npy", mmap_mode="r"
+            )
+            self.monthly_row_in_groups = np.load(
+                cache_dir / "exposure_monthly_row_in_group.npy", mmap_mode="r"
+            )
+            self.daily_files = [Path(path) for path in self.manifest["daily_files"]]
+            self.monthly_files = [Path(path) for path in self.manifest["monthly_files"]]
+        else:
+            raise ValueError(f"Unknown exposure cache storage mode: {self.storage!r}")
         quality_path = cache_dir / "exposure_quality.npy"
         self.quality = np.load(quality_path, mmap_mode="r") if quality_path.is_file() else None
 
+        if self.storage == "dense_npy":
             if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE:
                 raise ValueError(
                     f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, "
@@ -58,13 +110,22 @@ class ExposureCache:
                     f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}"
                 )
         n_rows = len(self.eids)
-        if (
-            len(self.raw_tokens) != n_rows
-            or len(self.onset_dates) != n_rows
-            or self.daily.shape[0] != n_rows
-            or self.monthly.shape[0] != n_rows
-        ):
+        if len(self.raw_tokens) != n_rows or len(self.onset_dates) != n_rows:
             raise ValueError("Exposure cache metadata/daily/monthly row counts do not match")
+        if self.storage == "dense_npy":
+            if self.daily.shape[0] != n_rows or self.monthly.shape[0] != n_rows:
+                raise ValueError("Exposure cache metadata/daily/monthly row counts do not match")
+        else:
+            indexed_lengths = [
+                len(self.daily_file_ids),
+                len(self.daily_row_groups),
+                len(self.daily_row_in_groups),
+                len(self.monthly_file_ids),
+                len(self.monthly_row_groups),
+                len(self.monthly_row_in_groups),
+            ]
+            if any(length != n_rows for length in indexed_lengths):
+                raise ValueError("Exposure parquet index row counts do not match metadata")
 
         self._key_to_index: dict[tuple[int, int, int], int] | None = None
 
@@ -100,12 +161,83 @@ class ExposureCache:
     def daily_window(self, index: int) -> np.ndarray:
         if index < 0:
             return np.full(DAILY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
+        if self.storage == "dense_npy":
             return np.asarray(self.daily[index], dtype=np.float32)
+        return self._parquet_window("daily", index)
 
     def monthly_window(self, index: int) -> np.ndarray:
         if index < 0:
             return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
+        if self.storage == "dense_npy":
             return np.asarray(self.monthly[index], dtype=np.float32)
+        return self._parquet_window("monthly", index)
+
+    def _parquet_window(self, kind: Literal["daily", "monthly"], index: int) -> np.ndarray:
+        if kind == "daily":
+            file_id = int(self.daily_file_ids[index])
+            row_group = int(self.daily_row_groups[index])
+            row_in_group = int(self.daily_row_in_groups[index])
+            shape = DAILY_EXPOSURE_SHAPE
+            columns = _daily_exposure_columns()
+        else:
+            file_id = int(self.monthly_file_ids[index])
+            row_group = int(self.monthly_row_groups[index])
+            row_in_group = int(self.monthly_row_in_groups[index])
+            shape = MONTHLY_EXPOSURE_SHAPE
+            columns = _monthly_exposure_columns()
+
+        frame = self._read_parquet_row_group(kind, file_id, row_group, columns)
+        row = frame.iloc[row_in_group].reindex(columns)
+        n_channels = shape[1]
+        return (
+            row.to_numpy(dtype=np.float32, copy=True)
+            .reshape(n_channels, shape[0])
+            .transpose(1, 0)
+        )
+
+    def _read_parquet_row_group(
+        self,
+        kind: Literal["daily", "monthly"],
+        file_id: int,
+        row_group: int,
+        columns: list[str],
+    ) -> pd.DataFrame:
+        cache_key = (kind, file_id, row_group)
+        cached = self._row_group_cache.get(cache_key)
+        if cached is not None:
+            self._row_group_cache.move_to_end(cache_key)
+            return cached
+
+        try:
+            import pyarrow.parquet as pq
+        except ImportError as exc:
+            raise ImportError(
+                "Parquet exposure index loading requires pyarrow. Install requirements "
+                "or use a dense numpy exposure cache."
+            ) from exc
+
+        parquet_key = (kind, file_id)
+        parquet_file = self._parquet_files.get(parquet_key)
+        if parquet_file is None:
+            path = self.daily_files[file_id] if kind == "daily" else self.monthly_files[file_id]
+            parquet_file = pq.ParquetFile(path)
+            self._parquet_files[parquet_key] = parquet_file
+
+        available_columns = self._parquet_columns.get(parquet_key)
+        if available_columns is None:
+            available = set(parquet_file.schema.names)
+            available_columns = [col for col in columns if col in available]
+            self._parquet_columns[parquet_key] = available_columns
+
+        table = parquet_file.read_row_group(row_group, columns=available_columns)
+        frame = table.to_pandas()
+        if available_columns != columns:
+            frame = frame.reindex(columns=columns)
+        self._row_group_cache[cache_key] = frame
+        self._row_group_cache.move_to_end(cache_key)
+        while len(self._row_group_cache) > self._row_group_cache_size:
+            self._row_group_cache.popitem(last=False)
+        return frame
 
 
 def load_label_vocab(

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()
+    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:,} exposure rows to {args.output_dir}")
+        print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}")
 
 
 if __name__ == "__main__":