ML4H merge: tranche 1

pyhealth/data/data.py

@@ -155,7 +155,7 @@ def _filter_by_time_range_fast(self, df: pl.DataFrame, start: Optional[datetime]
             start_idx = np.searchsorted(ts_col, np.datetime64(start, "ms"), side="left")
         if end is not None:
             end_idx = np.searchsorted(ts_col, np.datetime64(end, "ms"), side="right")
-        return df.slice(start_idx, end_idx - start_idx)
+        return df.slice(start_idx, max(0, end_idx - start_idx))
 
     def _filter_by_event_type_regular(self, df: pl.DataFrame, event_type: Optional[str]) -> pl.DataFrame:
         """Regular filtering by event type. Time complexity: O(n)."""

pyhealth/datasets/base_dataset.py

@@ -3,7 +3,7 @@
 import pickle
 from abc import ABC
 from pathlib import Path
-from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable
+from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable, Union
 import functools
 import operator
 from urllib.parse import urlparse, urlunparse
@@ -71,6 +71,17 @@ def clean_path(path: str) -> str:
         return str(Path(path).expanduser().resolve())
 
 
+def resolve_table_path(root: str, file_path: str) -> str:
+    """Resolve a table file_path against dataset root.
+
+    Absolute paths and URLs are kept as-is so a generated file that cannot
+    live in a read-only data root (PhysioNet) can still be loaded from cache.
+    """
+    if is_url(file_path) or os.path.isabs(file_path):
+        return clean_path(file_path)
+    return clean_path(f"{root}/{file_path}")
+
+
 def path_exists(path: str) -> bool:
     """
     Check if a path exists.
@@ -84,7 +95,13 @@ def path_exists(path: str) -> bool:
         except requests.RequestException:
             return False
     else:
-        return Path(path).exists()
+        try:
+            return Path(path).exists()
+        except OSError:
+            # Treat unreadable paths (e.g. stale/corrupted filesystem
+            # entries that raise I/O errors on stat) as non-existent so
+            # callers can fall back to an alternate extension.
+            return False
 
 
 def _csv_tsv_gz_path(path: str) -> str:
@@ -314,16 +331,7 @@ class BaseDataset(ABC):
         dataset_name (str): Name of the dataset.
         config (dict): Configuration loaded from a YAML file.
         global_event_df (pl.LazyFrame): The global event data frame.
-        dev (bool): Whether to enable dev mode (limit to 1000 patients).
-
-    Examples:
-        >>> from pyhealth.datasets import BaseDataset
-        >>> dataset = BaseDataset(
-        ...     root="/path/to/source",
-        ...     tables=["patients", "diagnoses"],
-        ...     config_path="/path/to/config.yaml",
-        ... )
-        >>> dataset.stats()
+        dev (Union[bool, int]): Whether to enable dev mode. If True, limit to 1000 patients. If an int, limit to that many patients.
     """
 
     def __init__(
@@ -334,7 +342,7 @@ def __init__(
         config_path: Optional[str] = None,
         cache_dir: str | Path | None = None,
         num_workers: int = 1,
-        dev: bool = False,
+        dev: Union[bool, int] = False,
     ):
         """Initializes the BaseDataset.
 
@@ -351,7 +359,7 @@ def __init__(
                 - **str** or **Path**: Used as the root cache directory path. A UUID
                   is appended to the provided path to capture dataset configuration.
             num_workers (int): Number of worker processes for parallel operations.
-            dev (bool): Whether to run in dev mode (limits to 1000 patients).
+            dev (Union[bool, int]): Whether to run in dev mode. If True, limits to 1000 patients. If an int, limits to that many patients.
         """
         if len(set(tables)) != len(tables):
             logger.warning("Duplicate table names in tables list. Removing duplicates.")
@@ -571,30 +579,52 @@ def _event_transform(self, output_dir: Path) -> None:
         compute_ok = False
         try:
             df = self.load_data()
-            with DaskCluster(
-                n_workers=self.num_workers,
-                threads_per_worker=1,
-                processes=not in_notebook(),
-                # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory
-                local_directory=str(self.create_tmpdir()),
-            ) as cluster:
-                with DaskClient(cluster) as client:
-                    if self.dev:
-                        logger.info("Dev mode enabled: limiting to 1000 patients")
-                        patients = df["patient_id"].unique().head(1000).tolist()
-                        filter = df["patient_id"].isin(patients)
-                        df = df[filter]
-
-                    logger.info(f"Caching event dataframe to {output_dir}...")
-                    collection = df.sort_values("patient_id").to_parquet(
-                        output_dir,
-                        write_index=False,
-                        compute=False,
-                    )
-                    handle = client.compute(collection)
-                    dask_progress(handle)
-                    handle.result()  # type: ignore
-                    compute_ok = True  # Data is fully written to disk
+            disable_distributed = os.environ.get(
+                "PYHEALTH_DISABLE_DASK_DISTRIBUTED", "0"
+            ) == "1"
+
+            if disable_distributed:
+                logger.info(
+                    "PYHEALTH_DISABLE_DASK_DISTRIBUTED=1 detected; using local dask scheduler."
+                )
+                if self.dev:
+                    n = 1000 if self.dev is True else int(self.dev)
+                    logger.info(f"Dev mode enabled: limiting to {n} patients")
+                    patients = df["patient_id"].unique().head(n, compute=True).tolist()
+                    patient_filter = df["patient_id"].isin(patients)
+                    df = df[patient_filter]
+
+                logger.info(f"Caching event dataframe to {output_dir}...")
+                df.sort_values("patient_id").to_parquet(
+                    output_dir,
+                    write_index=False,
+                    compute=True,
+                )
+            else:
+                with DaskCluster(
+                    n_workers=self.num_workers,
+                    threads_per_worker=1,
+                    processes=not in_notebook(),
+                    # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory
+                    local_directory=str(self.create_tmpdir()),
+                ) as cluster:
+                    with DaskClient(cluster) as client:
+                        if self.dev:
+                            logger.info(f"Dev mode enabled: limiting to {1000 if self.dev is True else int(self.dev)} patients")
+                            patients = df["patient_id"].unique().head(1000 if self.dev is True else int(self.dev)).tolist()
+                            filter = df["patient_id"].isin(patients)
+                            df = df[filter]
+
+                        logger.info(f"Caching event dataframe to {output_dir}...")
+                        collection = df.sort_values("patient_id").to_parquet(
+                            output_dir,
+                            write_index=False,
+                            compute=False,
+                        )
+                        handle = client.compute(collection)
+                        dask_progress(handle)
+                        handle.result()  # type: ignore
+                        compute_ok = True  # Data is fully written to disk
         except TimeoutError:
             if compute_ok:
                 # Cluster shutdown timed out after successful compute — data is intact
@@ -667,17 +697,15 @@ def load_table(self, table_name: str) -> dd.DataFrame:
 
         Raises:
             ValueError: If the table is not found in the config.
-            FileNotFoundError: If the source file (CSV/TSV or Parquet) for the
-                table or join is not found.
+            FileNotFoundError: If the CSV file for the table or join is not found.
         """
         assert self.config is not None, "Config must be provided to load tables"
 
         if table_name not in self.config.tables:
             raise ValueError(f"Table {table_name} not found in config")
 
         table_cfg = self.config.tables[table_name]
-        csv_path = f"{self.root}/{table_cfg.file_path}"
-        csv_path = clean_path(csv_path)
+        csv_path = resolve_table_path(self.root, table_cfg.file_path)
 
         logger.info(f"Scanning table: {table_name} from {csv_path}")
         df = self._scan_table(csv_path)
@@ -696,8 +724,7 @@ def load_table(self, table_name: str) -> dd.DataFrame:
 
         # Handle joins
         for join_cfg in table_cfg.join:
-            other_csv_path = f"{self.root}/{join_cfg.file_path}"
-            other_csv_path = clean_path(other_csv_path)
+            other_csv_path = resolve_table_path(self.root, join_cfg.file_path)
             logger.info(f"Joining with table: {other_csv_path}")
             join_df = self._scan_table(other_csv_path)
             join_df = join_df.rename(columns=str.lower)
@@ -723,21 +750,14 @@ def load_table(self, table_name: str) -> dd.DataFrame:
                 timestamp_series: dd.Series = functools.reduce(
                     operator.add, (df[col].astype("string") for col in timestamp_col)
                 )
-                timestamp_series = dd.to_datetime(
-                    timestamp_series,
-                    format=timestamp_format,
-                    errors="raise",
-                )
-            elif pd.api.types.is_datetime64_any_dtype(df[timestamp_col].dtype):
-                # Typed sources (e.g. Parquet) already carry native timestamps:
-                # skip the string round-trip and only normalize the unit below.
-                timestamp_series: dd.Series = df[timestamp_col]
             else:
-                timestamp_series = dd.to_datetime(
-                    df[timestamp_col].astype("string"),
-                    format=timestamp_format,
-                    errors="raise",
-                )
+                timestamp_series: dd.Series = df[timestamp_col].astype("string")
+
+            timestamp_series: dd.Series = dd.to_datetime(
+                timestamp_series,
+                format=timestamp_format,
+                errors="raise",
+            )
             df: dd.DataFrame = df.assign(
                 timestamp=timestamp_series.astype("datetime64[ms]")
             )
@@ -1165,4 +1185,4 @@ def _main_guard(self, func_name: str):
                 f"{func_name} method accessed from a non-main process. This may lead to unexpected behavior.\n"
                 + "Consider use __name__ == '__main__' guard when using multiprocessing."
             )
-            exit(1)
+            exit(1)

pyhealth/datasets/collate.py

@@ -17,14 +17,51 @@
 from typing import Any
 
 import torch
-from torch.nn.utils.rnn import pad_sequence
+import torch.nn.functional as F
+
+
+def _pad_stack(tensors: list[torch.Tensor]) -> torch.Tensor:
+    """Right-pad same-rank tensors to the per-dimension max, then stack.
+
+    ``pad_sequence`` only pads dimension 0 and requires every trailing dimension
+    to already match. Tokenized notes are ``(n_notes, seq_len)`` and, once the
+    text processor pads to the longest note in a sample rather than to a fixed
+    ``max_length``, BOTH dimensions vary across samples.
+    """
+    if len({t.dim() for t in tensors}) != 1:
+        raise ValueError("cannot pad tensors of differing rank")
+    target = [max(t.shape[d] for t in tensors) for d in range(tensors[0].dim())]
+    padded = []
+    for t in tensors:
+        spec: list[int] = []
+        for d in range(t.dim() - 1, -1, -1):
+            spec.extend([0, target[d] - t.shape[d]])
+        padded.append(F.pad(t, spec) if any(spec) else t)
+    return torch.stack(padded)
 
 
 def _stack_or_pad(tensors: list[torch.Tensor]) -> torch.Tensor:
-    """Stack if all shapes match; pad along dim-0 otherwise."""
+    """Stack if all shapes match; pad every ragged dimension otherwise."""
     if all(t.shape == tensors[0].shape for t in tensors):
         return torch.stack(tensors)
-    return pad_sequence(tensors, batch_first=True)
+    return _pad_stack(tensors)
+
+
+def _pad_mask(tensors: list[torch.Tensor]) -> torch.Tensor:
+    """Event-level validity for the tensor :func:`_stack_or_pad` just built.
+
+    Batch padding is created here and nowhere else, so it has to be recorded
+    here. A padded slot carries value 0.0 and time 0.0, which is
+    indistinguishable from a real measurement taken at admission time, so a
+    model given no mask treats padding as data.
+
+    This is deliberately NOT called ``mask``. A field may carry its own
+    ``{field}_mask`` meaning "was this value observed", which is a different
+    question from "is this slot real".
+    """
+    lengths = torch.tensor([t.shape[0] for t in tensors])
+    width = int(lengths.max())
+    return torch.arange(width)[None, :] < lengths[:, None]
 
 
 def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]:
@@ -62,6 +99,8 @@ def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]:
                     sub_result[sub_key] = [None] * len(sub_vals)
                 elif isinstance(sub_vals[0], torch.Tensor):
                     sub_result[sub_key] = _stack_or_pad(sub_vals)
+                    if sub_key == "time":
+                        sub_result["pad_mask"] = _pad_mask(sub_vals)
                 else:
                     sub_result[sub_key] = sub_vals
             result[key] = sub_result