Optimize exposure index training IO

dataset.py

@@ -45,7 +45,7 @@ def _monthly_exposure_columns() -> list[str]:
 class ExposureCache:
     """Random-access view over files produced by prepare_exposure_cache.py."""
 
-    def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 16):
+    def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 4):
         cache_dir = Path(cache_dir)
         self.cache_dir = cache_dir
         manifest_path = cache_dir / "exposure_manifest.json"
@@ -129,6 +129,21 @@ class ExposureCache:
 
         self._key_to_index: dict[tuple[int, int, int], int] | None = None
 
+    def locality_key(self, indices: np.ndarray) -> tuple[int, int]:
+        """Return a stable parquet locality key for sampler-side batching."""
+        indices = np.asarray(indices, dtype=np.int64)
+        valid = indices[indices >= 0]
+        if len(valid) == 0:
+            return (2**31 - 1, 2**31 - 1)
+        if self.storage != "parquet_index":
+            return (0, int(valid[0] // 1024))
+        file_ids = np.asarray(self.daily_file_ids[valid], dtype=np.int64)
+        row_groups = np.asarray(self.daily_row_groups[valid], dtype=np.int64)
+        groups = np.stack([file_ids, row_groups], axis=1)
+        unique_groups, counts = np.unique(groups, axis=0, return_counts=True)
+        best = unique_groups[int(np.argmax(counts))]
+        return (int(best[0]), int(best[1]))
+
     def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None:
         keys: dict[tuple[int, int, int], int] = {}
         eids = np.asarray(self.eids, dtype=np.int64)
@@ -159,18 +174,67 @@ class ExposureCache:
         return out
 
     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)
+        return self.daily_windows(np.asarray([index], dtype=np.int64))[0]
 
     def monthly_window(self, index: int) -> np.ndarray:
-        if index < 0:
-            return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
+        return self.monthly_windows(np.asarray([index], dtype=np.int64))[0]
+
+    def daily_windows(self, indices: np.ndarray) -> np.ndarray:
+        return self._windows("daily", indices)
+
+    def monthly_windows(self, indices: np.ndarray) -> np.ndarray:
+        return self._windows("monthly", indices)
+
+    def _windows(
+        self,
+        kind: Literal["daily", "monthly"],
+        indices: np.ndarray,
+    ) -> np.ndarray:
+        indices = np.asarray(indices, dtype=np.int64)
+        shape = DAILY_EXPOSURE_SHAPE if kind == "daily" else MONTHLY_EXPOSURE_SHAPE
+        out = np.full((len(indices), shape[0], shape[1]), np.nan, dtype=np.float32)
+        valid_pos = np.nonzero(indices >= 0)[0]
+        if len(valid_pos) == 0:
+            return out
+
+        valid_indices = indices[valid_pos]
         if self.storage == "dense_npy":
-            return np.asarray(self.monthly[index], dtype=np.float32)
-        return self._parquet_window("monthly", index)
+            source = self.daily if kind == "daily" else self.monthly
+            out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32)
+            return out
+
+        if kind == "daily":
+            file_ids = np.asarray(self.daily_file_ids[valid_indices], dtype=np.int64)
+            row_groups = np.asarray(self.daily_row_groups[valid_indices], dtype=np.int64)
+            row_in_groups = np.asarray(self.daily_row_in_groups[valid_indices], dtype=np.int64)
+            columns = _daily_exposure_columns()
+        else:
+            file_ids = np.asarray(self.monthly_file_ids[valid_indices], dtype=np.int64)
+            row_groups = np.asarray(self.monthly_row_groups[valid_indices], dtype=np.int64)
+            row_in_groups = np.asarray(
+                self.monthly_row_in_groups[valid_indices],
+                dtype=np.int64,
+            )
+            columns = _monthly_exposure_columns()
+
+        group_keys = np.stack([file_ids, row_groups], axis=1)
+        unique_groups, inverse = np.unique(group_keys, axis=0, return_inverse=True)
+        for group_idx, (file_id, row_group) in enumerate(unique_groups):
+            group_pos = np.nonzero(inverse == group_idx)[0]
+            frame = self._read_parquet_row_group(
+                kind,
+                int(file_id),
+                int(row_group),
+                columns,
+            )
+            row_values = frame.iloc[row_in_groups[group_pos]].reindex(columns=columns)
+            values = (
+                row_values.to_numpy(dtype=np.float32, copy=True)
+                .reshape(len(group_pos), shape[1], shape[0])
+                .transpose(0, 2, 1)
+            )
+            out[valid_pos[group_pos]] = values
+        return out
 
     def _parquet_window(self, kind: Literal["daily", "monthly"], index: int) -> np.ndarray:
         if kind == "daily":
@@ -308,13 +372,17 @@ class _ExpoBaseDataset(Dataset):
         include_no_event_in_uts_target: bool = False,
         exposure_cache_dir: str | Path | None = None,
         mask_onset_exposure: bool = False,
+        exposure_row_group_cache_size: int = 4,
     ) -> None:
         self.data_prefix = data_prefix
         self.labels_file = labels_file
         self.no_event_interval_years = float(no_event_interval_years)
         self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
         self.exposure_cache = (
-            ExposureCache(exposure_cache_dir)
+            ExposureCache(
+                exposure_cache_dir,
+                row_group_cache_size=exposure_row_group_cache_size,
+            )
             if exposure_cache_dir is not None
             else None
         )
@@ -450,14 +518,14 @@ class _ExpoBaseDataset(Dataset):
         if self.exposure_cache is None:
             raise RuntimeError("Exposure cache is not enabled")
 
-        daily = np.stack(
-            [self.exposure_cache.daily_window(int(idx)) for idx in exposure_index],
-            axis=0,
-        ).astype(np.float32, copy=True)
-        monthly = np.stack(
-            [self.exposure_cache.monthly_window(int(idx)) for idx in exposure_index],
-            axis=0,
-        ).astype(np.float32, copy=True)
+        daily = self.exposure_cache.daily_windows(exposure_index).astype(
+            np.float32,
+            copy=False,
+        )
+        monthly = self.exposure_cache.monthly_windows(exposure_index).astype(
+            np.float32,
+            copy=False,
+        )
 
         if self.mask_onset_exposure:
             daily[:, 0, :] = np.nan
@@ -485,6 +553,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
         include_no_event_in_uts_target: bool = False,
         exposure_cache_dir: str | Path | None = None,
         mask_onset_exposure: bool = False,
+        exposure_row_group_cache_size: int = 4,
     ) -> None:
         super().__init__(
             data_prefix=data_prefix,
@@ -493,6 +562,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
             mask_onset_exposure=mask_onset_exposure,
+            exposure_row_group_cache_size=exposure_row_group_cache_size,
         )
 
         self.samples: List[Dict] = []
@@ -581,6 +651,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         validation_query_seed: int = 42,
         exposure_cache_dir: str | Path | None = None,
         mask_onset_exposure: bool = False,
+        exposure_row_group_cache_size: int = 4,
     ) -> None:
         if split not in {"train", "valid", "test"}:
             raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -592,6 +663,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
             mask_onset_exposure=mask_onset_exposure,
+            exposure_row_group_cache_size=exposure_row_group_cache_size,
         )
 
         self.split = split