PyHealth 2.0: several modules still target the 1.x dataset API #1201

Summary

While preparing a fix for #952 I audited the modules that still reference the
PyHealth 1.x dataset API. kg_emb is not an isolated case. It is the only
one whose import breaks, which is why it is the only one with an issue.
The others fail later, and quietly, mostly because the 1.x base classes now
resolve to deprecation stubs whose __init__ only emits a warning:

class BaseEHRDataset:
    """This class is deprecated and should not be used."""

    def __init__(self, *args, **kwargs):
        import warnings
        warnings.warn(
            "The BaseEHRDataset class is deprecated and will be removed in a future version.",
            DeprecationWarning,
        )

A subclass of such a stub constructs successfully, sets none of its attributes,
and fails somewhere downstream, often several calls later, with an error that
does not point at the cause. All of the modules below are still exported from
their packages, so users can import and instantiate them today.

This is a tracker, not a request for one large PR. I am opening it so the scope
is visible in one place; I am not proposing to fix all of it myself.

Last checked on 2026-10-06 against master at 5568155 (17 commits past
v2.0.2). The only later commit on master, e94ca1d (#1273), does not change
any of the code cited below. What changed since the first version is listed at
the end.

Affected modules

Group A: reachable runtime failures

Group B: stale references, no runtime impact

Resolved since the first version

Two API-level observations

These are not module bugs, but they shape how a migration would be written.

1. get_dataloader assumes the SampleDataset hierarchy. It calls
   dataset.set_shuffle(shuffle), a litdata.StreamingDataset method that
   InMemorySampleDataset also implements, so it cannot accept a plain
   torch.utils.data.Dataset even though its docstring presents it as a general
   helper. Any module migrating to a plain Dataset has to bypass it.
2. collate_fn_dict_with_padding dispatches on tuple length. Any value that
   is a 2-tuple is treated as a temporal (time, values) pair, so a task whose
   samples legitimately contain 2-tuples would be silently mis-collated. KG
   triples are 3-tuples and pass through unaffected, so this is a latent hazard
   rather than a current bug.

Suggested triage

Each module needs one of three decisions, and I do not think that decision is mine:

• migrate to the 2.0 BaseDataset / SampleDataset pipeline;
• remove, with a clear error pointing at the 2.0 replacement;
• keep as documented-deprecated, but fail loudly.

The current situation is, in my view, the worst of the three, since the stubs
neither work nor fail clearly. Making the five stub __init__ methods (the four
in pyhealth/datasets/__init__.py and Visit in pyhealth/data/__init__.py)
raise NotImplementedError with a migration pointer, instead of warning and
returning a half-built object, is a small change that would convert the
stub-backed Group A entries (mimicextract, shhs, isruc, cardiology) from
a confusing downstream error into an actionable one, and would surface the real
scope immediately.

One inconsistency could be settled at the same time: five 1.x names were retired
behind warning stubs, but SampleBaseDataset was removed outright with no stub
at all. That asymmetry is why kg_emb was the only module to fail at import,
and therefore the only one that got reported. Whichever policy is chosen, applying
it uniformly would make the remaining breakage visible instead of silent.

How this was verified

Read-only audit against master at 0a75f99: import-chain tracing plus
construction attempts on the concrete classes. Group A entries, as first listed,
were reproduced; Group B entries are static findings from usage-site greps,
excluding the definitions in pyhealth/datasets/__init__.py and
pyhealth/data/__init__.py. Individual tracebacks are available on request.

Re-check on 2026-10-06 against master at 5568155, on a local install of that
commit (Windows, Python 3.12, NumPy 2.2.6, PyTables added for the HDF5 read):
the kg_emb import still raises ImportError (SampleBaseDataset); with that
name aliased to SampleDataset in memory (plus a placeholder for pandarallel,
which kg_emb imports but 2.0 does not declare), SampleKGDataset(samples=[...])
raises ValueError from litdata (a list where a directory is expected);
MIMICExtractDataset.parse_basic_info() on a minimal synthetic
all_hourly_data.h5 raises AttributeError at its first parallel_apply call;
Event(...) called with the keyword arguments used at mimicextract.py l. 267,
377 and 487 raises TypeError (missing event_type); SHHSDataset and
ISRUCDataset construct, then process_EEG_data() raises AttributeError;
CardiologyDataset constructs, keeps root, and has no set_task; the three
legacy drug recommendation functions raise
TypeError: object of type 'Patient' has no len() on a 2.0 Patient;
VAE(...) raises TypeError from BaseModel.__init__; the fairness_utils
doctest passes.

Changes since the first version

• medcode/.../kg_emb: reworded. The first version said the constructor passed a list into SampleDataset; that happens only once SampleBaseDataset is renamed to SampleDataset.
• metrics/fairness_utils: fixed by Fix broken example scripts and fairness metrics import #1200, moved to "Resolved".
• tasks/drug_recommendation: reworded after Add DrugRecommendationOMOP task #1203 (legacy note in the API docs, DrugRecommendationOMOP).
• datasets/mimicextract: added the three Event(...) calls without event_type. The first version said parse_basic_info fails before any parallel_apply call; with PyTables installed, it fails at the first one.
• datasets/cardiology: linked to the review of fix: use context managers for label files in cardiology_detect #1241.
• models/vae: moved from Group B to Group A, since construction fails. The first version also said that the __main__ example builds a SampleSignalDataset; it only imports it.
• Stubs: five, not four (Visit was not counted). The triage paragraph now names the Group A entries that the change would affect.
• API observation 1: InMemorySampleDataset also implements set_shuffle. The "general helper" remark now refers to the docstring, since the dataset parameter is annotated litdata.StreamingDataset.
• Adjacent item from my 29 Aug comment below: np.in1d is called at kg_base.py lines 155 and 171 on master (the comment first said 167, the line on the fix(kg_emb): migrate SampleKGDataset off the removed 1.x dataset API #1202 branch at the time). It works under the pinned numpy~=2.2.0; checked present in NumPy 2.2.0, 2.3.0 and 2.3.5, and absent (AttributeError) in 2.4.0.

Related: #952, #1202, #1200, #1203, #1210, #1241, #1247.