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,29 +71,61 @@ 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 = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r")
-        self.monthly = np.load(cache_dir / "exposure_monthly.npy", 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.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]}, "
-                f"{DAILY_EXPOSURE_SHAPE[1]}), got {self.daily.shape}"
-            )
-        if self.monthly.ndim != 3 or self.monthly.shape[1:] != MONTHLY_EXPOSURE_SHAPE:
-            raise ValueError(
-                f"exposure_monthly.npy must have shape (N, {MONTHLY_EXPOSURE_SHAPE[0]}, "
-                f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}"
-            )
+        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]}, "
+                    f"{DAILY_EXPOSURE_SHAPE[1]}), got {self.daily.shape}"
+                )
+            if self.monthly.ndim != 3 or self.monthly.shape[1:] != MONTHLY_EXPOSURE_SHAPE:
+                raise ValueError(
+                    f"exposure_monthly.npy must have shape (N, {MONTHLY_EXPOSURE_SHAPE[0]}, "
+                    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)
-        return np.asarray(self.daily[index], 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)
-        return np.asarray(self.monthly[index], 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(