diff --git a/docs/api/datasets/pyhealth.datasets.splitter.rst b/docs/api/datasets/pyhealth.datasets.splitter.rst
index d9e96447d..899b8ba14 100644
--- a/docs/api/datasets/pyhealth.datasets.splitter.rst
+++ b/docs/api/datasets/pyhealth.datasets.splitter.rst
@@ -3,6 +3,80 @@
 
 Several data splitting function for `pyhealth.datasets` module to obtain training / validation / test sets.
 
+Fitting processors on the training split
+----------------------------------------
+
+The ``split_by_*`` functions split an already-processed ``SampleDataset``: its
+processors (code vocabularies, and any statistics a processor learns) were fitted
+on all samples, validation and test patients included. To fit them on the
+training patients only, let ``set_task`` do the split:
+
+.. code-block:: python
+
+    from pyhealth.datasets import PatientSplit
+
+    train, val, test = dataset.set_task(
+        task, split=PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42)
+    )
+    train.fit_split   # {'kind': 'patient', 'ratios': [0.7, 0.1, 0.2], 'seed': 42}
+
+- Every processor is fitted on the training patients; validation and test samples
+  are processed with those frozen processors, so a code seen only in test maps to
+  ``<unk>``, as it would at deployment.
+- Patients are sorted by ID before the seeded shuffle, so a ``PatientSplit`` gives
+  the same patients on every run and machine. It takes two or more ratios summing
+  to 1, training part first: ``(train, test)``, ``(train, val, test)``,
+  ``(train, val, cal, test)`` and so on, returned in that order.
+- Samples are streamed while fitting: the training patients are selected from the
+  index of samples per patient that processing already keeps, and each processor
+  reads the same sample stream, skipping other patients. Memory does not grow with
+  the size of the samples. ``examples/benchmark_perf/benchmark_split_set_task.py``
+  measures it at scale on synthetic data.
+- The result is cached per split: a rerun with the same ``PatientSplit`` reuses it,
+  another seed or ratio builds its own.
+- Processors that learn statistics from the data (``learns_statistics = True``)
+  log a warning when fitted on all samples, since that leaks validation/test
+  information. Vocabulary-only processors do not warn.
+
+The default, ``set_task(task)`` without ``split``, is unchanged: it fits on all
+samples and returns one dataset.
+
+Writing your own split
+----------------------
+
+``split=`` accepts any subclass of ``Split``. It implements two methods:
+
+- ``split_indices(patient_to_index)`` returns one array of sample indices per part,
+  training part first. ``patient_to_index`` maps each patient ID to the indices of
+  its samples. ``set_task`` fits the processors on part 0 and returns every part as
+  given: parts may overlap or leave samples out.
+- ``to_dict()`` returns a JSON-serialisable description. It goes into the cache key,
+  so it must change whenever the parts would, and is saved as ``fit_split``.
+
+For example, holding out a fixed set of patients:
+
+.. code-block:: python
+
+    import numpy as np
+    from pyhealth.datasets import Split
+
+    class HoldoutSplit(Split):
+        def __init__(self, test_patients):
+            self.test_patients = sorted(test_patients)
+
+        def split_indices(self, patient_to_index):
+            test = set(self.test_patients)
+            parts = ([], [])
+            for pid, indices in patient_to_index.items():
+                parts[pid in test].extend(indices)
+            return [np.array(sorted(p), dtype=np.int64) for p in parts]
+
+        def to_dict(self):
+            return {"kind": "holdout", "test_patients": self.test_patients}
+
+    train, test = dataset.set_task(task, split=HoldoutSplit(["patient-a", "patient-b"]))
+
+
 .. automodule:: pyhealth.datasets.splitter
     :members:
     :undoc-members:
diff --git a/docs/api/processors.rst b/docs/api/processors.rst
index e4262fa4c..ea83659b4 100644
--- a/docs/api/processors.rst
+++ b/docs/api/processors.rst
@@ -276,6 +276,34 @@ For maximum control, you can pass pre-configured processor instances directly in
             return samples
 
 
+Reusing Fitted Processors
+-------------------------
+
+``set_task()`` and ``create_sample_dataset()`` accept already-fitted processors
+through ``input_processors`` and ``output_processors``. A supplied processor is
+used as it is and never refitted; every schema field you do not supply gets a
+processor fitted on the samples as usual. A common use is fitting on training
+samples only and reusing those processors for validation and test data, so test
+patients do not shape preprocessing:
+
+.. code-block:: python
+
+    from pyhealth.datasets import create_sample_dataset
+
+    train = create_sample_dataset(train_samples, input_schema, output_schema)
+    test = create_sample_dataset(
+        test_samples, input_schema, output_schema,
+        input_processors=train.input_processors,
+        output_processors=train.output_processors,
+    )
+
+In PyHealth 2.0.2 and earlier, supplying a processor for some fields fitted
+none for the rest of that schema: those fields reached the model as raw Python
+values, without an error. See ``examples/reuse_train_processors.py``.
+
+For datasets built with ``set_task``, ``set_task(task, split=PatientSplit(...))``
+does this in one step and streams the samples; see :doc:`datasets/pyhealth.datasets.splitter`.
+
 Processor String Keys
 ---------------------
 
diff --git a/examples/benchmark_perf/benchmark_split_set_task.py b/examples/benchmark_perf/benchmark_split_set_task.py
new file mode 100644
index 000000000..d4496b758
--- /dev/null
+++ b/examples/benchmark_perf/benchmark_split_set_task.py
@@ -0,0 +1,235 @@
+"""Benchmark set_task with and without split=PatientSplit(...) at scale.
+
+Generates synthetic MIMIC-III-format tables (no real patient data) for N
+patients, then runs MIMIC3Dataset + MortalityPredictionMIMIC3 twice in separate
+processes, each with a fresh cache:
+
+  nosplit : dataset.set_task(task)                         (current behaviour)
+  split   : dataset.set_task(task, split=PatientSplit(...)) (fit on train only)
+
+For each run it reports wall time per stage and the peak resident memory of the
+process plus its workers, sampled every 0.1 s. Use it to check that fitting on
+the training split streams the samples (no extra memory) on your machine.
+
+Usage:
+    python examples/benchmark_perf/benchmark_split_set_task.py --patients 50000
+    python examples/benchmark_perf/benchmark_split_set_task.py --patients 300000 --num-workers 4
+
+Data and caches go under --workdir (default: a temporary folder, deleted at the
+end; pass a path to keep them). Generation is reused when the folder exists.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import psutil
+
+
+# ---------------------------------------------------------------- data ----
+def generate_mimic3(root: Path, n_patients: int, seed: int = 0) -> dict:
+    """Writes synthetic MIMIC-III CSVs with the columns PyHealth reads."""
+    rng = np.random.default_rng(seed)
+    root.mkdir(parents=True, exist_ok=True)
+    subject = np.arange(10_000, 10_000 + n_patients)
+
+    pd.DataFrame(
+        {
+            "row_id": np.arange(n_patients),
+            "subject_id": subject,
+            "gender": rng.choice(["F", "M"], n_patients),
+            "dob": "2080-01-01 00:00:00",
+            "dod": "",
+            "dod_hosp": "",
+            "dod_ssn": "",
+            "expire_flag": 0,
+        }
+    ).to_csv(root / "PATIENTS.csv", index=False)
+
+    n_adm = rng.integers(1, 7, n_patients)  # 1-6 admissions per patient
+    adm_subject = np.repeat(subject, n_adm)
+    n_a = adm_subject.size
+    hadm = np.arange(100_000, 100_000 + n_a)
+    order = np.concatenate([np.arange(k) for k in n_adm])
+    admit = pd.Timestamp("2100-01-01") + pd.to_timedelta(order * 60 + rng.integers(0, 30, n_a), unit="D")
+    disch = admit + pd.to_timedelta(rng.integers(1, 15, n_a), unit="D")
+    pd.DataFrame(
+        {
+            "row_id": np.arange(n_a), "subject_id": adm_subject, "hadm_id": hadm,
+            "admittime": admit.strftime("%Y-%m-%d %H:%M:%S"),
+            "dischtime": disch.strftime("%Y-%m-%d %H:%M:%S"),
+            "deathtime": "", "admission_type": "EMERGENCY",
+            "admission_location": "EMERGENCY ROOM ADMIT", "discharge_location": "HOME",
+            "insurance": "Medicare", "language": "ENGL", "religion": "", "marital_status": "",
+            "ethnicity": "WHITE", "edregtime": "", "edouttime": "", "diagnosis": "",
+            "hospital_expire_flag": (rng.random(n_a) < 0.03).astype(int),
+            "has_chartevents_data": 1,
+        }
+    ).to_csv(root / "ADMISSIONS.csv", index=False)
+
+    pd.DataFrame(
+        {
+            "subject_id": adm_subject, "intime": admit.strftime("%Y-%m-%d %H:%M:%S"),
+            "icustay_id": np.arange(200_000, 200_000 + n_a), "first_careunit": "MICU",
+            "dbsource": "carevue", "last_careunit": "MICU",
+            "outtime": disch.strftime("%Y-%m-%d %H:%M:%S"),
+        }
+    ).to_csv(root / "ICUSTAYS.csv", index=False)
+
+    def codes_table(name, lo, hi, vocab, extra):
+        k = rng.integers(lo, hi + 1, n_a)
+        frame = pd.DataFrame(
+            {
+                "row_id": np.arange(k.sum()),
+                "subject_id": np.repeat(adm_subject, k),
+                "hadm_id": np.repeat(hadm, k),
+                **extra(k),
+            }
+        )
+        frame.to_csv(root / name, index=False)
+        return int(k.sum())
+
+    icd = np.array([f"{c:05d}" for c in rng.choice(100_000, 6_000, replace=False)])
+    proc = np.array([f"{c:04d}" for c in rng.choice(10_000, 2_000, replace=False)])
+    ndc = np.array([f"{c:011d}" for c in rng.choice(10**11, 3_000, replace=False)])
+    n_diag = codes_table(
+        "DIAGNOSES_ICD.csv", 3, 15, icd,
+        lambda k: {"seq_num": np.concatenate([np.arange(1, j + 1) for j in k]),
+                   "icd9_code": rng.choice(icd, k.sum())},
+    )
+    n_proc = codes_table(
+        "PROCEDURES_ICD.csv", 1, 5, proc,
+        lambda k: {"seq_num": np.concatenate([np.arange(1, j + 1) for j in k]),
+                   "icd9_code": rng.choice(proc, k.sum())},
+    )
+
+    def presc_cols(k):
+        start = (np.repeat(admit, k)).strftime("%Y-%m-%d %H:%M:%S")
+        return {
+            "icustay_id": "", "startdate": start, "enddate": start, "drug_type": "MAIN",
+            "drug": "drug", "drug_name_poe": "", "drug_name_generic": "", "formulary_drug_cd": "",
+            "gsn": "", "ndc": rng.choice(ndc, k.sum()), "prod_strength": "", "dose_val_rx": "1",
+            "dose_unit_rx": "mg", "form_val_disp": "1", "form_unit_disp": "TAB", "route": "PO",
+        }
+
+    n_presc = codes_table("PRESCRIPTIONS.csv", 2, 15, ndc, presc_cols)
+    return {"patients": n_patients, "admissions": int(n_a), "diagnoses": n_diag,
+            "procedures": n_proc, "prescriptions": n_presc}
+
+
+# ------------------------------------------------------------- memory ----
+class PeakMemory:
+    """Peak RSS of this process plus its children, sampled in a thread."""
+
+    def __init__(self, interval: float = 0.1):
+        self.interval, self.peak, self._stop = interval, 0, threading.Event()
+        self._proc = psutil.Process()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+
+    def _rss(self) -> int:
+        total = self._proc.memory_info().rss
+        for child in self._proc.children(recursive=True):
+            try:
+                total += child.memory_info().rss
+            except psutil.Error:
+                pass
+        return total
+
+    def _run(self):
+        while not self._stop.is_set():
+            self.peak = max(self.peak, self._rss())
+            time.sleep(self.interval)
+
+    def __enter__(self):
+        self.peak = self._rss()
+        self._thread.start()
+        return self
+
+    def __exit__(self, *exc):
+        self._stop.set()
+        self._thread.join()
+        self.peak = max(self.peak, self._rss())
+
+
+# --------------------------------------------------------------- child ----
+def run_child(root: str, cache: str, mode: str, num_workers: int) -> dict:
+    from pyhealth.datasets import MIMIC3Dataset, PatientSplit
+    from pyhealth.tasks import MortalityPredictionMIMIC3
+
+    out = {"mode": mode, "baseline_rss_mib": psutil.Process().memory_info().rss / 2**20}
+    t0 = time.perf_counter()
+    with PeakMemory() as mem:
+        dataset = MIMIC3Dataset(
+            root=root, tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
+            cache_dir=cache, num_workers=num_workers,
+        )
+        _ = dataset.global_event_df
+    out["events_s"], out["events_peak_mib"] = time.perf_counter() - t0, mem.peak / 2**20
+
+    split = PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42) if mode == "split" else None
+    t1 = time.perf_counter()
+    with PeakMemory() as mem:
+        result = dataset.set_task(MortalityPredictionMIMIC3(), num_workers=num_workers, split=split)
+    out["set_task_s"], out["set_task_peak_mib"] = time.perf_counter() - t1, mem.peak / 2**20
+
+    parts = result if isinstance(result, tuple) else (result,)
+    out["samples"] = [len(p) for p in parts]
+    out["patients_in_index"] = len(parts[0].patient_to_index) if mode == "nosplit" else None
+    schema = Path(parts[0].path) / "schema.pkl"
+    out["schema_pkl_mib"] = schema.stat().st_size / 2**20
+    split_file = Path(parts[0].path) / "split.npz"
+    out["split_npz_mib"] = split_file.stat().st_size / 2**20 if split_file.exists() else 0.0
+    return out
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--patients", type=int, default=20_000)
+    parser.add_argument("--num-workers", type=int, default=1)
+    parser.add_argument("--workdir", type=str, default=None)
+    parser.add_argument("--modes", type=str, default="nosplit,split")
+    parser.add_argument("--child", nargs=4, metavar=("ROOT", "CACHE", "MODE", "WORKERS"))
+    args = parser.parse_args()
+
+    if args.child:
+        root, cache, mode, workers = args.child
+        print("RESULT " + json.dumps(run_child(root, cache, mode, int(workers))), flush=True)
+        return
+
+    workdir = Path(args.workdir) if args.workdir else Path(tempfile.mkdtemp(prefix="pyhealth_split_bench_"))
+    data = workdir / f"mimic3_synthetic_{args.patients}"
+    try:
+        if not (data / "PRESCRIPTIONS.csv").exists():
+            t = time.perf_counter()
+            sizes = generate_mimic3(data, args.patients)
+            print(f"generated {sizes} in {time.perf_counter() - t:.0f} s", flush=True)
+        for mode in args.modes.split(","):
+            cache = workdir / f"cache_{args.patients}_{mode}"
+            shutil.rmtree(cache, ignore_errors=True)
+            proc = subprocess.run(
+                [sys.executable, __file__, "--child", str(data), str(cache), mode, str(args.num_workers)],
+                capture_output=True, text=True,
+            )
+            result = [l for l in proc.stdout.splitlines() if l.startswith("RESULT ")]
+            if proc.returncode or not result:
+                print(proc.stderr[-3000:], file=sys.stderr)
+                raise SystemExit(f"{mode} run failed (exit {proc.returncode})")
+            print(result[-1][7:], flush=True)
+    finally:
+        if not args.workdir:
+            shutil.rmtree(workdir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/reuse_train_processors.py b/examples/reuse_train_processors.py
new file mode 100644
index 000000000..bdb8bafee
--- /dev/null
+++ b/examples/reuse_train_processors.py
@@ -0,0 +1,52 @@
+"""Fit processors on training samples and reuse them for test samples.
+
+Fitting processors only on training data keeps test patients from shaping
+preprocessing (vocabularies, and any statistics a processor learns). Pass the
+training dataset's processors when building the test dataset: supplied
+processors are reused as they are, never refitted.
+
+Runs in a few seconds on CPU with synthetic data; no download.
+
+Usage:
+    python examples/reuse_train_processors.py
+"""
+
+from pyhealth.datasets import create_sample_dataset
+
+INPUT_SCHEMA = {"codes": "sequence", "age": "tensor"}
+OUTPUT_SCHEMA = {"label": "binary"}
+
+train_samples = [
+    {"patient_id": "p1", "codes": ["I10", "E11"], "age": [61.0], "label": 1},
+    {"patient_id": "p2", "codes": ["J45"], "age": [34.0], "label": 0},
+    {"patient_id": "p3", "codes": ["E11", "N18"], "age": [72.0], "label": 1},
+]
+test_samples = [
+    # "C50" never occurs in training, so it maps to <unk>, as at deployment
+    {"patient_id": "p9", "codes": ["I10", "C50"], "age": [58.0], "label": 0},
+]
+
+
+def main():
+    train = create_sample_dataset(
+        train_samples, INPUT_SCHEMA, OUTPUT_SCHEMA, dataset_name="train"
+    )
+    codes = train.input_processors["codes"]
+    print("training vocabulary:", sorted(codes.code_vocab))
+
+    # Reuse every fitted processor. (A field left out of input_processors would
+    # get a processor fitted on the samples passed here, i.e. on test data.)
+    test = create_sample_dataset(
+        test_samples,
+        INPUT_SCHEMA,
+        OUTPUT_SCHEMA,
+        input_processors=train.input_processors,
+        output_processors=train.output_processors,
+        dataset_name="test",
+    )
+    print("same codes processor:", test.input_processors["codes"] is codes)
+    print("test sample:", test[0])
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/split_aware_set_task_mimic3demo.py b/examples/split_aware_set_task_mimic3demo.py
new file mode 100644
index 000000000..8f5f5f0e8
--- /dev/null
+++ b/examples/split_aware_set_task_mimic3demo.py
@@ -0,0 +1,62 @@
+"""Fit processors on training patients only with set_task(task, split=...).
+
+Splitting after set_task() means the code vocabularies (and any statistics a
+processor learns) were fitted on every patient, test patients included.
+Passing a PatientSplit to set_task fits them on the training patients only and
+returns train / val / test datasets ready for a model.
+
+Uses the small MIMIC-III demo bundled in ``test-resources/``; no download.
+
+Usage (from the repository root):
+    python examples/split_aware_set_task_mimic3demo.py
+"""
+
+import tempfile
+from pathlib import Path
+
+from pyhealth.datasets import MIMIC3Dataset, PatientSplit, get_dataloader
+from pyhealth.models import RNN
+from pyhealth.tasks import MortalityPredictionMIMIC3
+from pyhealth.trainer import Trainer
+
+DEMO_ROOT = (
+    Path(__file__).resolve().parent.parent / "test-resources" / "core" / "mimic3demo"
+)
+
+
+def main():
+    with tempfile.TemporaryDirectory() as cache:
+        dataset = MIMIC3Dataset(
+            root=str(DEMO_ROOT),
+            tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
+            cache_dir=cache,
+        )
+        train, val, test = dataset.set_task(
+            MortalityPredictionMIMIC3(),
+            split=PatientSplit(ratios=(0.6, 0.2, 0.2), seed=0),
+        )
+        # Ratios split patients, not samples: patients have different numbers
+        # of admissions, so sample counts need not follow the ratios.
+        print("patients (train, val, test):",
+              *(len(part.patient_to_index) for part in (train, val, test)))
+        print("samples  (train, val, test):", len(train), len(val), len(test))
+        print("processors fitted on:", train.fit_split)
+
+        vocab = train.input_processors["conditions"].code_vocab
+        test_codes = {c for i in range(len(test)) for c in test[i]["conditions"].tolist()}
+        unknown = vocab["<unk>"]
+        print(f"training vocabulary: {len(vocab)} codes; "
+              f"test samples use <unk>: {unknown in test_codes}")
+
+        model = RNN(dataset=train)
+        trainer = Trainer(model=model, metrics=["roc_auc", "pr_auc"], enable_logging=False)
+        trainer.train(
+            get_dataloader(train, batch_size=16, shuffle=True),
+            val_dataloader=get_dataloader(val, batch_size=16),
+            epochs=1,
+        )
+        print("test:", trainer.evaluate(get_dataloader(test, batch_size=16)))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py
index 57a9956c2..7acd921e6 100644
--- a/pyhealth/datasets/__init__.py
+++ b/pyhealth/datasets/__init__.py
@@ -71,6 +71,8 @@ def __init__(self, *args, **kwargs):
 from .support2 import Support2Dataset
 from .tcga_prad import TCGAPRADDataset
 from .splitter import (
+    PatientSplit as PatientSplit,
+    Split as Split,
     sample_balanced,
     split_by_patient,
     split_by_patient_conformal,
diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py
index 8534778a0..9a16c668f 100644
--- a/pyhealth/datasets/base_dataset.py
+++ b/pyhealth/datasets/base_dataset.py
@@ -45,6 +45,7 @@
 from ..processors.base_processor import FeatureProcessor
 from .configs import load_yaml_config
 from .sample_dataset import SampleDataset, SampleBuilder
+from .splitter import Split
 from ..utils import set_env
 
 # Set logging level for distributed to ERROR to reduce verbosity
@@ -205,6 +206,14 @@ def _litdata_merge(cache_dir: Path) -> None:
     """
     Merges LitData binary writer index files in the given cache directory.
 
+    Workers write contiguous sample ranges, each with its own
+    ``{rank}.index.json``. ``patient_to_index`` and ``record_to_index`` are
+    built before processing and assume merged sample ``i`` is input sample
+    ``i``, so the per-worker indexes must be joined in numeric rank order.
+    ``BinaryWriter.merge`` does this from litdata 0.2.64 on; earlier versions
+    sort the file names as strings (0, 1, 10, 11, ..., 2), which reorders
+    samples with more than 10 workers.
+
     Args:
         cache_dir (Path): The cache directory containing LitData binary writer files.
     """
@@ -1106,7 +1115,8 @@ def set_task(
         num_workers: Optional[int] = None,
         input_processors: Optional[Dict[str, FeatureProcessor]] = None,
         output_processors: Optional[Dict[str, FeatureProcessor]] = None,
-    ) -> SampleDataset:
+        split: Split | None = None,
+    ) -> SampleDataset | tuple[SampleDataset, ...]:
         """Processes the base dataset to generate the task-specific sample dataset.
         The cache structure is as follows::
 
@@ -1114,6 +1124,7 @@ def set_task(
                 task_df.ld/                 # Intermediate task dataframe based on schema
                 samples_{proc_uuid}.ld/     # Final processed samples after applying processors
                     schema.pkl              # Saved SampleBuilder schema
+                    split.npz               # Sample indices per split part (with split= only)
                     *.bin                   # Processed sample files
 
         Args:
@@ -1125,9 +1136,25 @@ def set_task(
             output_processors (Optional[Dict[str, FeatureProcessor]]):
                 Pre-fitted output processors. If provided, these will be used
                 instead of creating new ones from task's output_schema. Defaults to None.
+            split (Optional[Split]): Split the samples, e.g. by patient with
+                ``PatientSplit``, and fit every processor on the first
+                (training) part only, so the other parts never shape
+                preprocessing. The parts are used as the split returns them.
+                Samples are streamed;
+                nothing is loaded into memory beyond the per-sample index that
+                processing already keeps. Defaults to None: fit on all samples
+                and return one dataset, as before.
 
         Returns:
-            SampleDataset: The generated sample dataset.
+            SampleDataset: The generated sample dataset, or, with ``split``, a
+            tuple with one dataset per part, training part first, whose
+            processors were fitted on the training part.
+
+        Examples:
+            >>> from pyhealth.datasets import PatientSplit
+            >>> train, val, test = dataset.set_task(  # doctest: +SKIP
+            ...     task, split=PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42)
+            ... )
 
         Raises:
             AssertionError: If no default task is found and task is None.
@@ -1180,6 +1207,8 @@ def set_task(
                     if output_processors
                     else None
                 ),
+                # Only present with a split, so existing unsplit caches keep their key.
+                **({"split": split.to_dict()} if split is not None else {}),
             },
             sort_keys=True,
             default=str,
@@ -1240,8 +1269,12 @@ def _is_valid_litdata_cache(path: Path) -> bool:
                         input_processors=input_processors,
                         output_processors=output_processors,
                     )
-                    builder.fit(dataset)
+                    builder.fit(dataset, split=split)
                     builder.save(str(samples_path / "schema.pkl"))
+                    if split is not None:
+                        # Written before the samples, whose index.json marks the
+                        # cache complete, so a valid cache always has its split.
+                        np.savez(samples_path / "split.npz", *builder.split_indices)
 
                     # Apply processors and save final samples to cache_dir
                     logger.info(f"Processing samples and saving to {samples_path}...")
@@ -1256,11 +1289,17 @@ def _is_valid_litdata_cache(path: Path) -> bool:
                 f"Found cached processed samples at {samples_path}, skipping processing."
             )
 
-        return SampleDataset(
+        samples = SampleDataset(
             path=str(samples_path),
             dataset_name=self.dataset_name,
             task_name=task.task_name,
         )
+        if split is None:
+            return samples
+        with np.load(samples_path / "split.npz") as parts:
+            return tuple(
+                samples.subset(parts[f"arr_{k}"]) for k in range(len(parts.files))
+            )
 
     def _main_guard(self, func_name: str):
         """Warn if method is accessed from a non-main process."""
diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py
index 6b9a751ee..dfd99b2a2 100644
--- a/pyhealth/datasets/sample_dataset.py
+++ b/pyhealth/datasets/sample_dataset.py
@@ -4,16 +4,43 @@
 import shutil
 import tempfile
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
+from typing import TYPE_CHECKING
 import inspect
 import random
 from bisect import bisect_right
 import litdata
 from litdata.utilities.train_test_split import deepcopy_dataset
 import copy
+import logging
+
+import numpy as np
 
 from ..processors import get_processor, IgnoreProcessor
 from ..processors.base_processor import FeatureProcessor
 
+if TYPE_CHECKING:  # splitter imports this module
+    from .splitter import Split
+
+logger = logging.getLogger(__name__)
+
+
+class _MaskedSamples:
+    """A re-iterable view of a sample stream that keeps only masked positions.
+
+    Each iteration walks the underlying stream in order and yields the samples
+    whose position is True in ``mask``. Nothing is collected, so fitting on a
+    split costs no more memory than fitting on all samples.
+    """
+
+    def __init__(self, samples: Iterable[dict[str, Any]], mask: np.ndarray):
+        self._samples = samples
+        self._mask = mask
+
+    def __iter__(self):
+        for keep, sample in zip(self._mask, self._samples):
+            if keep:
+                yield sample
+
 
 def _remap_index_mapping(
     mapping: dict[str, list[int]], indices: Sequence[int]
@@ -49,6 +76,29 @@ class SampleBuilder:
     After saving the schema, `litdata.optimize` can be used with `builder.transform`
     to serialize and chunk pickled sample items into a directory that can be
     loaded via SampleDataset.
+
+    Pre-fitted processors passed as `input_processors` / `output_processors` are
+    used as they are; every schema field without one gets a fitted processor.
+
+    Examples:
+        >>> from pyhealth.datasets.sample_dataset import SampleBuilder
+        >>> from pyhealth.processors import SequenceProcessor
+        >>> samples = [
+        ...     {"patient_id": "p1", "codes": ["a", "b"], "age": [40.0], "label": 1},
+        ...     {"patient_id": "p2", "codes": ["c"], "age": [55.0], "label": 0},
+        ... ]
+        >>> codes = SequenceProcessor()
+        >>> codes.fit(samples, "codes")
+        >>> builder = SampleBuilder(
+        ...     input_schema={"codes": "sequence", "age": "tensor"},
+        ...     output_schema={"label": "binary"},
+        ...     input_processors={"codes": codes},
+        ... )
+        >>> builder.fit(samples)
+        >>> builder.input_processors["codes"] is codes
+        True
+        >>> type(builder.input_processors["age"]).__name__
+        'TensorProcessor'
     """
 
     def __init__(
@@ -61,14 +111,16 @@ def __init__(
         self.input_schema = input_schema
         self.output_schema = output_schema
         self._input_processors = (
-            input_processors if input_processors is not None else {}
+            dict(input_processors) if input_processors is not None else {}
         )
         self._output_processors = (
-            output_processors if output_processors is not None else {}
+            dict(output_processors) if output_processors is not None else {}
         )
         self._patient_to_index: Dict[str, List[int]] = {}
         self._record_to_index: Dict[str, List[int]] = {}
         self._fitted = False
+        self._split = None
+        self._split_indices: list[np.ndarray] | None = None
 
     @property
     def input_processors(self) -> Dict[str, FeatureProcessor]:
@@ -86,6 +138,15 @@ def output_processors(self) -> Dict[str, FeatureProcessor]:
             )
         return self._output_processors
 
+    @property
+    def split_indices(self) -> list[np.ndarray] | None:
+        """Sample indices of each split part (train first), or None without a split."""
+        if not self._fitted:
+            raise RuntimeError(
+                "SampleBuilder.fit must be called before accessing split_indices."
+            )
+        return self._split_indices
+
     @property
     def patient_to_index(self) -> Dict[str, List[int]]:
         if not self._fitted:
@@ -130,6 +191,7 @@ def _get_processor_instance(self, processor_spec):
     def fit(
         self,
         samples: Iterable[Dict[str, Any]],
+        split: "Split | None" = None,
     ) -> None:
         """Fit processors and build mapping indices from an iterator of samples.
 
@@ -138,6 +200,13 @@ def fit(
                 sample should contain keys covering both the configured
                 `input_schema` and `output_schema`. These samples are not
                 required to be pickled; `fit` operates on in-memory dicts.
+                It is iterated several times and never collected, so a
+                streaming dataset works without loading it into memory.
+            split: Optional :class:`~pyhealth.datasets.Split`, such as
+                :class:`~pyhealth.datasets.PatientSplit`. When given,
+                processors are fitted on the first (training) part only, and
+                the sample indices of every part are kept in
+                ``split_indices``.
 
         Behavior:
             - Validates the samples contain all keys specified by the input
@@ -145,9 +214,9 @@ def fit(
             - Builds `patient_to_index` and `record_to_index` mappings by
               recording the sample indices associated with `patient_id` and
               `record_id`/`visit_id` fields.
-            - Instantiates and fits input/output processors from the provided
-              schemas (unless pre-fitted processors were supplied to the
-              constructor).
+            - Instantiates and fits a processor for every schema field that
+              has none. Pre-fitted processors supplied to the constructor are
+              used as they are; any remaining fields still get fitted ones.
         """
         # Validate the samples
         input_keys = set(self.input_schema.keys())
@@ -163,7 +232,9 @@ def fit(
         # Build index mappings
         self._patient_to_index = {}
         self._record_to_index = {}
+        n_samples = 0
         for i, sample in enumerate(samples):
+            n_samples = i + 1
             patient_id = sample.get("patient_id")
             if patient_id is not None:
                 self._patient_to_index.setdefault(patient_id, []).append(i)
@@ -171,17 +242,38 @@ def fit(
             if record_id is not None:
                 self._record_to_index.setdefault(record_id, []).append(i)
 
-        # Fit processors if they were not provided
-        if not self._input_processors:
-            for key, spec in self.input_schema.items():
-                processor = self._get_processor_instance(spec)
-                processor.fit(samples, key)
-                self._input_processors[key] = processor
-        if not self._output_processors:
-            for key, spec in self.output_schema.items():
-                processor = self._get_processor_instance(spec)
-                processor.fit(samples, key)
-                self._output_processors[key] = processor
+        # With a split, fit on the training patients only: a masked view of
+        # the same stream, so nothing is loaded into memory.
+        self._split = split
+        self._split_indices = None
+        fit_samples = samples
+        if split is not None:
+            self._split_indices = split.split_indices(self._patient_to_index)
+            train_mask = np.zeros(n_samples, dtype=bool)
+            train_mask[self._split_indices[0]] = True
+            fit_samples = _MaskedSamples(samples, train_mask)
+
+        # Fit a processor for every schema field that has none. Supplied
+        # (pre-fitted) processors are used as they are, never refitted.
+        learned_on_all = []
+        for schema, processors in (
+            (self.input_schema, self._input_processors),
+            (self.output_schema, self._output_processors),
+        ):
+            for key, spec in schema.items():
+                if key not in processors:
+                    processor = self._get_processor_instance(spec)
+                    processor.fit(fit_samples, key)
+                    processors[key] = processor
+                    if split is None and getattr(processor, "learns_statistics", False):
+                        learned_on_all.append(key)
+        if learned_on_all:
+            logger.warning(
+                "Processors for %s learn statistics from the data and were fitted "
+                "on all samples, so validation/test patients shaped them. Fit on "
+                "the training patients with set_task(task, split=PatientSplit(...)).",
+                ", ".join(learned_on_all),
+            )
 
         self._fitted = True
 
@@ -239,6 +331,7 @@ def save(self, path: str) -> None:
             "output_processors": self._output_processors,
             "patient_to_index": self._patient_to_index,
             "record_to_index": self._record_to_index,
+            "fit_split": self._split.to_dict() if self._split is not None else None,
         }
         with open(path, "wb") as f:
             pickle.dump(metadata, f)
@@ -342,6 +435,9 @@ def __init__(
 
         self.patient_to_index = metadata["patient_to_index"]
         self.record_to_index = metadata["record_to_index"]
+        # The split the processors were fitted on, or None if fitted on all
+        # samples (also None for caches written before this was recorded).
+        self.fit_split: dict | None = metadata.get("fit_split")
 
     def _remove_ignored_processors(self):
         """Remove any processors that are IgnoreProcessor instances."""
diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py
index 2dbc94186..3f4678618 100644
--- a/pyhealth/datasets/splitter.py
+++ b/pyhealth/datasets/splitter.py
@@ -1,3 +1,7 @@
+import math
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
 from itertools import chain
 from typing import List, Optional, Tuple, Union
 
@@ -10,6 +14,135 @@
 # TODO: add more splitting methods
 
 
+class Split(ABC):
+    """Interface for ``BaseDataset.set_task(task, split=...)``.
+
+    ``set_task`` fits every processor on the samples of part 0 and returns one
+    ``SampleDataset`` per part, in order, as views over one processed cache.
+    It uses the parts exactly as returned: they may overlap or leave samples
+    out. :class:`PatientSplit` is the reference implementation.
+
+    Examples:
+        >>> import numpy as np
+        >>> from pyhealth.datasets import Split
+        >>> class HoldoutSplit(Split):
+        ...     def __init__(self, test_patients):
+        ...         self.test_patients = sorted(test_patients)
+        ...     def split_indices(self, patient_to_index):
+        ...         test = set(self.test_patients)
+        ...         parts = ([], [])
+        ...         for pid, indices in patient_to_index.items():
+        ...             parts[pid in test].extend(indices)
+        ...         return [np.array(sorted(p), dtype=np.int64) for p in parts]
+        ...     def to_dict(self):
+        ...         return {"kind": "holdout", "test_patients": self.test_patients}
+        >>> HoldoutSplit(["p2"]).split_indices({"p1": [0, 1], "p2": [2]})
+        [array([0, 1]), array([2])]
+    """
+
+    @abstractmethod
+    def split_indices(
+        self, patient_to_index: Mapping[str, Sequence[int]]
+    ) -> list[np.ndarray]:
+        """Returns the sample indices of each part.
+
+        Args:
+            patient_to_index: Patient ID to the indices of its samples.
+
+        Returns:
+            One ``int64`` array of sample indices per part, training part
+            first. Sorted indices keep reads from the cache sequential.
+        """
+
+    @abstractmethod
+    def to_dict(self) -> dict:
+        """Returns a JSON-serialisable description of the split.
+
+        It goes into the processed-cache key, so it must change whenever the
+        parts would, and is saved as ``SampleDataset.fit_split``.
+        """
+
+
+@dataclass(frozen=True)
+class PatientSplit(Split):
+    """A patient-level split for ``BaseDataset.set_task(task, split=...)``.
+
+    With a split, ``set_task`` fits every processor on the training patients
+    only and returns one dataset per part, so validation and test patients
+    never shape preprocessing (vocabularies, learned statistics).
+
+    Patients are sorted by ID before shuffling with ``seed``, so the same
+    ``PatientSplit`` gives the same patients on every run and machine,
+    whatever order the samples are stored in.
+
+    Unlike :func:`split_by_patient`, which splits an already-processed
+    ``SampleDataset`` whose processors were fitted on all samples, the parts
+    are chosen before fitting. ``split_by_patient`` shuffles patients in the
+    order they appear in the dataset, so for the same ratios and seed it
+    generally picks different patients.
+
+    Args:
+        ratios: Two or more non-negative fractions summing to 1. The first
+            part is the training part, e.g. (train, test), (train, val, test)
+            or (train, val, cal, test).
+        seed: Seed for the shuffle.
+
+    Examples:
+        >>> from pyhealth.datasets import PatientSplit
+        >>> split = PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42)
+        >>> parts = split.split_indices({"p1": [0, 1], "p2": [2], "p3": [3, 4]})
+        >>> sorted(int(i) for part in parts for i in part)
+        [0, 1, 2, 3, 4]
+    """
+
+    ratios: Sequence[float]
+    seed: int = 0
+
+    def __post_init__(self):
+        ratios = tuple(float(r) for r in self.ratios)
+        if len(ratios) < 2:
+            raise ValueError(f"PatientSplit needs at least 2 ratios, training part first; got {ratios}.")
+        if any(r < 0 for r in ratios):
+            raise ValueError(f"PatientSplit ratios must be non-negative, got {ratios}.")
+        if not math.isclose(sum(ratios), 1.0, abs_tol=1e-6):
+            raise ValueError(f"PatientSplit ratios must sum to 1, got {ratios} (sum {sum(ratios)}).")
+        object.__setattr__(self, "ratios", ratios)
+
+    def to_dict(self) -> dict:
+        """JSON-serialisable description, used in cache keys and metadata."""
+        return {"kind": "patient", "ratios": list(self.ratios), "seed": int(self.seed)}
+
+    def split_indices(
+        self, patient_to_index: Mapping[str, Sequence[int]]
+    ) -> list[np.ndarray]:
+        """Splits sample indices by patient.
+
+        Patients are sorted, shuffled with ``seed`` and cut by ``ratios``.
+
+        Args:
+            patient_to_index: Patient ID to the indices of its samples.
+
+        Returns:
+            One sorted ``int64`` array of sample indices per ratio, in order
+            (train first). Each patient's samples land in exactly one part.
+        """
+        patients = sorted(patient_to_index, key=str)
+        order = np.random.default_rng(self.seed).permutation(len(patients))
+        cuts = [
+            int(len(patients) * sum(self.ratios[: k + 1]))
+            for k in range(len(self.ratios) - 1)
+        ]
+        parts = []
+        for part in np.split(order, cuts):
+            indices = np.fromiter(
+                chain.from_iterable(patient_to_index[patients[p]] for p in part),
+                dtype=np.int64,
+            )
+            indices.sort()
+            parts.append(indices)
+        return parts
+
+
 def _label_to_int(label) -> int:
     """Convert a stored label (int/np scalar/torch scalar) to Python int."""
     if torch.is_tensor(label):
@@ -117,12 +250,21 @@ def split_by_visit(
         seed: random seed for shuffling the dataset
 
     Returns:
-        train_dataset, val_dataset, test_dataset: three subsets of the dataset of
-            type `torch.utils.data.Subset`.
+        train_dataset, val_dataset, test_dataset: ``SampleDataset`` views
+            (``dataset.subset``) over the same processed samples.
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_visit
+        >>> train, val, test = split_by_visit(  # doctest: +SKIP
+        ...     samples, [0.8, 0.1, 0.1], seed=42
+        ... )
 
     Note:
-        The original dataset can be accessed by `train_dataset.dataset`,
-            `val_dataset.dataset`, and `test_dataset.dataset`.
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     rng = np.random.default_rng(seed)
     assert sum(ratios) == 1.0, "ratios must sum to 1.0"
@@ -152,12 +294,21 @@ def split_by_patient(
         seed: random seed for shuffling the dataset
 
     Returns:
-        train_dataset, val_dataset, test_dataset: three subsets of the dataset of
-            type `torch.utils.data.Subset`.
+        train_dataset, val_dataset, test_dataset: ``SampleDataset`` views
+            (``dataset.subset``) over the same processed samples.
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_patient
+        >>> train, val, test = split_by_patient(  # doctest: +SKIP
+        ...     samples, [0.8, 0.1, 0.1], seed=42
+        ... )
 
     Note:
-        The original dataset can be accessed by `train_dataset.dataset`,
-            `val_dataset.dataset`, and `test_dataset.dataset`.
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     rng = np.random.default_rng(seed)
     assert sum(ratios) == 1.0, "ratios must sum to 1.0"
@@ -192,14 +343,26 @@ def split_by_sample(
         dataset: a `SampleDataset` object
         ratios: a list/tuple of ratios for train / val / test
         seed: random seed for shuffling the dataset
+        get_index: if True, return three ``torch.Tensor`` index vectors
+            instead of datasets
 
     Returns:
-        train_dataset, val_dataset, test_dataset: three subsets of the dataset of
-            type `torch.utils.data.Subset`.
+        train_dataset, val_dataset, test_dataset: ``SampleDataset`` views
+            (``dataset.subset``) over the same processed samples, or index
+            tensors if ``get_index`` is True.
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_sample
+        >>> train, val, test = split_by_sample(  # doctest: +SKIP
+        ...     samples, [0.8, 0.1, 0.1], seed=42
+        ... )
 
     Note:
-        The original dataset can be accessed by `train_dataset.dataset`,
-            `val_dataset.dataset`, and `test_dataset.dataset`.
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     rng = np.random.default_rng(seed)
     assert sum(ratios) == 1.0, "ratios must sum to 1.0"
@@ -237,13 +400,22 @@ def split_by_visit_conformal(
         seed: random seed for shuffling the dataset
 
     Returns:
-        train_dataset, val_dataset, cal_dataset, test_dataset: four subsets
-            of the dataset of type `torch.utils.data.Subset`.
+        train_dataset, val_dataset, cal_dataset, test_dataset:
+            ``SampleDataset`` views (``dataset.subset``) over the same
+            processed samples.
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_visit_conformal
+        >>> train, val, cal, test = split_by_visit_conformal(  # doctest: +SKIP
+        ...     samples, [0.6, 0.1, 0.1, 0.2], seed=42
+        ... )
 
     Note:
-        The original dataset can be accessed by `train_dataset.dataset`,
-            `val_dataset.dataset`, `cal_dataset.dataset`, and
-            `test_dataset.dataset`.
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     rng = np.random.default_rng(seed)
     assert len(ratios) == 4, "ratios must have 4 elements for train/val/cal/test"
@@ -283,13 +455,22 @@ def split_by_patient_conformal(
         seed: random seed for shuffling the dataset
 
     Returns:
-        train_dataset, val_dataset, cal_dataset, test_dataset: four subsets
-            of the dataset of type `torch.utils.data.Subset`.
+        train_dataset, val_dataset, cal_dataset, test_dataset:
+            ``SampleDataset`` views (``dataset.subset``) over the same
+            processed samples.
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_patient_conformal
+        >>> train, val, cal, test = split_by_patient_conformal(  # doctest: +SKIP
+        ...     samples, [0.6, 0.1, 0.1, 0.2], seed=42
+        ... )
 
     Note:
-        The original dataset can be accessed by `train_dataset.dataset`,
-            `val_dataset.dataset`, `cal_dataset.dataset`, and
-            `test_dataset.dataset`.
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     rng = np.random.default_rng(seed)
     assert len(ratios) == 4, "ratios must have 4 elements for train/val/cal/test"
@@ -354,10 +535,23 @@ def split_by_patient_conformal_tuh(
             to 1.0.
         seed: random seed used to shuffle the patient list.
         get_index: if ``True``, return four :class:`torch.Tensor` index vectors
-            instead of :class:`~torch.utils.data.Subset` objects.
+            instead of datasets.
 
     Returns:
         ``(train_dataset, val_dataset, cal_dataset, test_dataset)``
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_patient_conformal_tuh
+        >>> train, val, cal, test = split_by_patient_conformal_tuh(  # doctest: +SKIP
+        ...     samples, [0.6, 0.2, 0.2], seed=42
+        ... )
+
+    Note:
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     assert len(ratios) == 3, (
         "ratios must have exactly 3 elements (train/val/cal). "
@@ -429,10 +623,23 @@ def split_by_sample_conformal_tuh(
         ratios: the fraction of the train pool assigned to train / val / cal respectively
         seed: random seed for shuffling the train pool
         get_index: if True, return four ``torch.Tensor`` index vectors instead
-            of ``Subset`` objects
+            of datasets
 
     Returns:
         train_dataset, val_dataset, cal_dataset, test_dataset
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_sample_conformal_tuh
+        >>> train, val, cal, test = split_by_sample_conformal_tuh(  # doctest: +SKIP
+        ...     samples, [0.6, 0.2, 0.2], seed=42
+        ... )
+
+    Note:
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     assert len(ratios) == 3, (
         "ratios must have exactly 3 elements (train/val/cal). "
@@ -511,10 +718,23 @@ def split_by_patient_tuh(
             to 1.0.
         seed: random seed used to shuffle the patient list.
         get_index: if ``True``, return three :class:`torch.Tensor` index
-            vectors instead of :class:`~torch.utils.data.Subset` objects.
+            vectors instead of datasets.
 
     Returns:
         ``(train_dataset, val_dataset, test_dataset)``
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_patient_tuh
+        >>> train, val, test = split_by_patient_tuh(  # doctest: +SKIP
+        ...     samples, [0.8, 0.2], seed=42
+        ... )
+
+    Note:
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     assert len(ratios) == 2, (
         "ratios must have exactly 2 elements (train/val). "
@@ -587,10 +807,23 @@ def split_by_sample_tuh(
             respectively.  Must be a length-2 sequence summing to 1.0.
         seed: random seed for shuffling the train pool
         get_index: if True, return three ``torch.Tensor`` index vectors instead
-            of ``Subset`` objects
+            of datasets
 
     Returns:
         train_dataset, val_dataset, test_dataset
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_sample_tuh
+        >>> train, val, test = split_by_sample_tuh(  # doctest: +SKIP
+        ...     samples, [0.8, 0.2], seed=42
+        ... )
+
+    Note:
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     assert len(ratios) == 2, (
         "ratios must have exactly 2 elements (train/val). "
@@ -649,17 +882,25 @@ def split_by_sample_conformal(
         dataset: a `SampleDataset` object
         ratios: a list/tuple of ratios for train / val / cal / test
         seed: random seed for shuffling the dataset
-        get_index: if True, return indices instead of Subset objects
+        get_index: if True, return indices instead of datasets
 
     Returns:
-        train_dataset, val_dataset, cal_dataset, test_dataset: four subsets
-            of the dataset of type `torch.utils.data.Subset`, or four tensors
-            of indices if get_index=True.
+        train_dataset, val_dataset, cal_dataset, test_dataset:
+            ``SampleDataset`` views (``dataset.subset``) over the same
+            processed samples, or four index tensors if ``get_index`` is True.
+
+    Examples:
+        >>> from pyhealth.datasets import split_by_sample_conformal
+        >>> train, val, cal, test = split_by_sample_conformal(  # doctest: +SKIP
+        ...     samples, [0.6, 0.1, 0.1, 0.2], seed=42
+        ... )
 
     Note:
-        The original dataset can be accessed by `train_dataset.dataset`,
-            `val_dataset.dataset`, `cal_dataset.dataset`, and
-            `test_dataset.dataset`.
+        The processors of ``dataset`` (vocabularies, learned statistics) were
+        fitted on all of its samples, so every part shaped them. To fit them on
+        the training part only, split inside ``set_task`` instead:
+        ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit`
+        or a custom :class:`Split`.
     """
     rng = np.random.default_rng(seed)
     assert len(ratios) == 4, "ratios must have 4 elements for train/val/cal/test"
diff --git a/pyhealth/processors/base_processor.py b/pyhealth/processors/base_processor.py
index 6f8f5ed63..650c143d1 100644
--- a/pyhealth/processors/base_processor.py
+++ b/pyhealth/processors/base_processor.py
@@ -63,8 +63,31 @@ class FeatureProcessor(Processor):
     Processor for individual fields (features).
 
     Example: Tokenization, image loading, normalization.
+
+    Set ``learns_statistics = True`` on processors whose ``fit()`` learns values
+    from the data (means, standard deviations, medians, bins). Fitting such a
+    processor on all samples leaks validation/test information, so PyHealth
+    warns unless it is fitted on a training split
+    (``set_task(task, split=PatientSplit(...))``).
+
+    Examples:
+        >>> import torch
+        >>> from pyhealth.processors.base_processor import FeatureProcessor
+        >>> class MeanCenter(FeatureProcessor):
+        ...     learns_statistics = True
+        ...     def fit(self, samples, field):
+        ...         values = [s[field] for s in samples]
+        ...         self.mean = sum(values) / len(values)
+        ...     def process(self, value):
+        ...         return torch.tensor(value - self.mean)
+        >>> p = MeanCenter()
+        >>> p.fit([{"x": 1.0}, {"x": 3.0}], "x")
+        >>> p.process(5.0)
+        tensor(3.)
     """
 
+    learns_statistics: bool = False
+
     def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None:
         """Fit the processor to the samples.
 
diff --git a/pyproject.toml b/pyproject.toml
index 49a2d73d5..62fe5750d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,9 @@ dependencies = [
   "pandas~=2.3.1",
   "pydantic~=2.11.7",
   "dask[complete]>=2026.1.0",
-  "litdata~=0.2.59",
+  # >=0.2.64: worker indexes merge in numeric rank order (patient_to_index relies
+  # on it). <0.2.68: 0.2.68+ reads SampleDataset.subset() parts ~10-35x slower.
+  "litdata>=0.2.64,<0.2.68",
   "pyarrow>=23.0.1",
   "orjson~=3.10",
   "narwhals~=2.13.0",
diff --git a/tests/core/test_caching.py b/tests/core/test_caching.py
index a346b98e1..65556ab0d 100644
--- a/tests/core/test_caching.py
+++ b/tests/core/test_caching.py
@@ -123,6 +123,7 @@ def test_set_task_signature(self):
             "num_workers",
             "input_processors",
             "output_processors",
+            "split",
         ]
         self.assertEqual(params, expected_params)
 
@@ -131,6 +132,7 @@ def test_set_task_signature(self):
         self.assertEqual(sig.parameters["num_workers"].default, None)
         self.assertEqual(sig.parameters["input_processors"].default, None)
         self.assertEqual(sig.parameters["output_processors"].default, None)
+        self.assertEqual(sig.parameters["split"].default, None)
 
     def test_set_task_writes_cache_and_metadata(self):
         """Ensure set_task materializes cache files and schema metadata."""
diff --git a/tests/core/test_litdata_merge.py b/tests/core/test_litdata_merge.py
new file mode 100644
index 000000000..5b190aad3
--- /dev/null
+++ b/tests/core/test_litdata_merge.py
@@ -0,0 +1,62 @@
+"""_litdata_merge must keep samples in worker order for any worker count.
+
+Workers write contiguous sample ranges, and patient_to_index (built before
+processing) assumes merged sample i is input sample i. Sorting the per-worker
+index files as strings puts rank 10 before rank 2, so this uses 12 ranks.
+"""
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import litdata
+from litdata.streaming.writer import BinaryWriter
+
+from pyhealth.datasets.base_dataset import _litdata_merge
+from pyhealth.utils import set_env
+
+N_WORKERS = 12
+PER_WORKER = 3
+
+
+def _write_worker_outputs(cache_dir: Path) -> None:
+    """Writes N_WORKERS contiguous ranges, as _proc_transform_fn does."""
+    with set_env(DATA_OPTIMIZER_NUM_WORKERS=str(N_WORKERS)):
+        for rank in range(N_WORKERS):
+            with set_env(DATA_OPTIMIZER_GLOBAL_RANK=str(rank)):
+                writer = BinaryWriter(cache_dir=str(cache_dir), chunk_bytes="64MB")
+                for k in range(PER_WORKER):
+                    i = rank * PER_WORKER + k
+                    writer.add_item(k, {"position": i, "patient_id": f"p{i:03d}"})
+                writer.done()
+
+
+class TestLitdataMerge(unittest.TestCase):
+    def test_merge_keeps_worker_order_beyond_ten_workers(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            cache = Path(tmp)
+            _write_worker_outputs(cache)
+            _litdata_merge(cache)
+            positions = [s["position"] for s in litdata.StreamingDataset(str(cache))]
+            self.assertEqual(positions, list(range(N_WORKERS * PER_WORKER)))
+            self.assertEqual(
+                sorted(p.name for p in cache.glob("*index.json")), ["index.json"]
+            )
+
+    def test_merge_is_noop_when_index_exists(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            cache = Path(tmp)
+            _write_worker_outputs(cache)
+            _litdata_merge(cache)
+            before = (cache / "index.json").read_text()
+            _litdata_merge(cache)
+            self.assertEqual((cache / "index.json").read_text(), before)
+
+    def test_merge_without_worker_indexes_raises(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            with self.assertRaises(ValueError):
+                _litdata_merge(Path(tmp))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_partial_processors.py b/tests/core/test_partial_processors.py
new file mode 100644
index 000000000..30b104915
--- /dev/null
+++ b/tests/core/test_partial_processors.py
@@ -0,0 +1,85 @@
+"""Supplying some processors must still fit the others.
+
+SampleBuilder used to fit input (or output) processors only when none were
+supplied. Supplying a processor for one field left every other field of that
+side without a processor, so its raw value (e.g. a Python list) reached the
+model without any error.
+"""
+
+import pickle
+import unittest
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset
+from pyhealth.datasets.sample_dataset import SampleBuilder
+from pyhealth.processors import (
+    BinaryLabelProcessor,
+    SequenceProcessor,
+    TensorProcessor,
+)
+
+
+def _samples():
+    return [
+        {"patient_id": "p1", "codes": ["a", "b"], "age": [40.0], "label": 1},
+        {"patient_id": "p2", "codes": ["c"], "age": [55.0], "label": 0},
+    ]
+
+
+class TestPartialProcessors(unittest.TestCase):
+    def setUp(self):
+        self.codes = SequenceProcessor()
+        self.codes.fit(_samples(), "codes")
+
+    def test_missing_input_processor_is_fitted(self):
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence", "age": "tensor"},
+            output_schema={"label": "binary"},
+            input_processors={"codes": self.codes},
+        )
+        builder.fit(_samples())
+        self.assertIs(builder.input_processors["codes"], self.codes)
+        self.assertIsInstance(builder.input_processors["age"], TensorProcessor)
+        out = builder.transform({"sample": pickle.dumps(_samples()[0])})
+        self.assertIsInstance(out["age"], torch.Tensor)
+
+    def test_supplied_processor_is_not_refitted(self):
+        vocab = dict(self.codes.code_vocab)
+        samples = _samples() + [
+            {"patient_id": "p3", "codes": ["z"], "age": [60.0], "label": 1}
+        ]
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence", "age": "tensor"},
+            output_schema={"label": "binary"},
+            input_processors={"codes": self.codes},
+        )
+        builder.fit(samples)
+        self.assertEqual(self.codes.code_vocab, vocab)  # "z" was not added
+
+    def test_missing_output_processor_is_fitted(self):
+        samples = [dict(s, label2=s["label"]) for s in _samples()]
+        label = BinaryLabelProcessor()
+        label.fit(samples, "label")
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence"},
+            output_schema={"label": "binary", "label2": "binary"},
+            output_processors={"label": label},
+        )
+        builder.fit(samples)
+        self.assertIs(builder.output_processors["label"], label)
+        self.assertIn("label2", builder.output_processors)
+
+    def test_create_sample_dataset_with_partial_processors(self):
+        dataset = create_sample_dataset(
+            samples=_samples(),
+            input_schema={"codes": "sequence", "age": "tensor"},
+            output_schema={"label": "binary"},
+            input_processors={"codes": self.codes},
+            dataset_name="test_partial_processors",
+        )
+        self.assertIsInstance(dataset[0]["age"], torch.Tensor)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_split_set_task.py b/tests/core/test_split_set_task.py
new file mode 100644
index 000000000..7c968a9e3
--- /dev/null
+++ b/tests/core/test_split_set_task.py
@@ -0,0 +1,357 @@
+"""set_task(task, split=PatientSplit(...)): fit processors on training patients only.
+
+The synthetic data makes leakage visible: test patients' values are shifted by
++100 and they carry a code that never occurs in training.
+"""
+
+import json
+import logging
+import tempfile
+import tracemalloc
+import unittest
+import uuid
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import BaseDataset, PatientSplit, Split
+from pyhealth.datasets.sample_dataset import SampleBuilder
+from pyhealth.processors import SequenceProcessor
+from pyhealth.processors.base_processor import FeatureProcessor
+from pyhealth.tasks import BaseTask
+
+N_PATIENTS = 40
+SPLIT = PatientSplit(ratios=(0.6, 0.2, 0.2), seed=3)
+
+CONFIG = """version: "1.0"
+tables:
+  events:
+    file_path: "events.csv"
+    patient_id: "patient_id"
+    timestamp: "time"
+    attributes:
+      - "code"
+      - "value"
+"""
+
+
+class MeanCenter(FeatureProcessor):
+    """Learns the mean of a numeric feature and subtracts it (a statistic)."""
+
+    learns_statistics = True
+
+    def __init__(self):
+        self.mean = None
+
+    def fit(self, samples, field):
+        values = [float(s[field][0]) for s in samples]
+        self.mean = sum(values) / len(values)
+
+    def process(self, value):
+        return torch.tensor([float(value[0]) - self.mean])
+
+
+class EventTask(BaseTask):
+    task_name = "split_test_task"
+    input_schema = {"codes": "sequence", "value": MeanCenter}
+    output_schema = {"label": "binary"}
+
+    def __call__(self, patient):
+        samples = []
+        for i, event in enumerate(patient.get_events("events")):
+            samples.append(
+                {
+                    "patient_id": patient.patient_id,
+                    "visit_id": f"{patient.patient_id}-{i}",
+                    "codes": event.code.split("|"),
+                    "value": [float(event.value)],
+                    "label": i % 2 if len(patient.get_events("events")) > 1 else int(patient.patient_id[-1]) % 2,
+                }
+            )
+        return samples
+
+
+class HoldoutSplit(Split):
+    """Custom split: the given patients form part 1, everyone else part 0."""
+
+    def __init__(self, test_patients):
+        self.test_patients = sorted(test_patients)
+
+    def split_indices(self, patient_to_index):
+        test = set(self.test_patients)
+        parts = ([], [])
+        for pid, indices in patient_to_index.items():
+            parts[pid in test].extend(indices)
+        return [np.array(sorted(p), dtype=np.int64) for p in parts]
+
+    def to_dict(self):
+        return {"kind": "holdout", "test_patients": self.test_patients}
+
+
+class OverlapSplit(Split):
+    """Custom split whose part 1 repeats part 0: set_task must accept it as given."""
+
+    def split_indices(self, patient_to_index):
+        everything = np.array(sorted(i for ix in patient_to_index.values() for i in ix), dtype=np.int64)
+        return [everything, everything[:5]]
+
+    def to_dict(self):
+        return {"kind": "overlap"}
+
+
+def _patient_ids():
+    return [f"p{i:03d}" for i in range(N_PATIENTS)]
+
+
+def _parts(split):
+    """Patients in each split part, from one-sample-per-patient indices."""
+    ids = _patient_ids()
+    parts = split.split_indices({pid: [k] for k, pid in enumerate(ids)})
+    return [{ids[int(k)] for k in part} for part in parts]
+
+
+def _write_dataset(root: Path, test_patients: set) -> None:
+    rows = ["patient_id,time,code,value"]
+    for k, pid in enumerate(_patient_ids()):
+        for v in range(1 + k % 3):  # 1-3 samples per patient
+            shift = 100.0 if pid in test_patients else 0.0
+            code = "A|TESTONLY" if pid in test_patients else ("A|B" if v % 2 else "A")
+            rows.append(f"{pid},2020-01-{1 + v:02d} 00:00:00,{code},{10.0 + v + shift}")
+    (root / "events.csv").write_text("\n".join(rows) + "\n")
+    (root / "config.yaml").write_text(CONFIG)
+
+
+class TestSplitSetTask(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.tmp = tempfile.TemporaryDirectory()
+        root = Path(cls.tmp.name) / "data"
+        root.mkdir()
+        cls.parts = _parts(SPLIT)
+        _write_dataset(root, test_patients=cls.parts[2])
+        cls.dataset = BaseDataset(
+            root=str(root),
+            tables=["events"],
+            dataset_name="SplitDataset",
+            config_path=str(root / "config.yaml"),
+            cache_dir=str(Path(cls.tmp.name) / "cache"),
+            num_workers=1,
+        )
+        cls.train, cls.val, cls.test = cls.dataset.set_task(EventTask(), split=SPLIT)
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.tmp.cleanup()
+
+    @staticmethod
+    def _patients(ds):
+        return {ds[i]["patient_id"] for i in range(len(ds))}
+
+    def test_parts_are_patient_disjoint_and_complete(self):
+        parts = [self._patients(d) for d in (self.train, self.val, self.test)]
+        self.assertEqual(parts, self.parts)
+        all_samples = self.dataset.set_task(EventTask())
+        self.assertEqual(sum(len(d) for d in (self.train, self.val, self.test)), len(all_samples))
+
+    def test_statistics_come_from_training_patients(self):
+        train_values = [
+            10.0 + v
+            for k, pid in enumerate(_patient_ids())
+            if pid in self.parts[0]
+            for v in range(1 + k % 3)
+        ]
+        mean = self.train.input_processors["value"].mean
+        self.assertAlmostEqual(mean, float(np.mean(train_values)))
+        self.assertLess(mean, 20.0)  # fitted on all samples it would include +100 values
+
+    def test_test_only_codes_map_to_unknown(self):
+        codes = self.train.input_processors["codes"]
+        self.assertNotIn("TESTONLY", codes.code_vocab)
+        sample = self.test[0]
+        self.assertIn(codes.code_vocab["<unk>"], sample["codes"].tolist())
+
+    def test_fit_split_is_recorded(self):
+        for part in (self.train, self.val, self.test):
+            self.assertEqual(part.fit_split, SPLIT.to_dict())
+        self.assertIsNone(self.dataset.set_task(EventTask()).fit_split)
+
+    def test_rerun_reuses_the_cache_and_the_same_split(self):
+        with self.assertLogs("pyhealth.datasets.base_dataset", "INFO") as logs:
+            again = self.dataset.set_task(EventTask(), split=SPLIT)
+        self.assertTrue(any("Found cached processed samples" in m for m in logs.output))
+        self.assertFalse(any("Fitting processors" in m for m in logs.output))
+        self.assertEqual([self._patients(d) for d in again], self.parts)
+
+    def test_other_seed_gives_other_split_and_cache(self):
+        other = PatientSplit(ratios=(0.6, 0.2, 0.2), seed=4)
+        train, _, _ = self.dataset.set_task(EventTask(), split=other)
+        self.assertEqual(self._patients(train), _parts(other)[0])
+        self.assertNotEqual(train.path, self.train.path)
+
+    def test_two_part_split(self):
+        parts = self.dataset.set_task(EventTask(), split=PatientSplit(ratios=(0.8, 0.2)))
+        self.assertEqual(len(parts), 2)
+
+    def test_four_part_split(self):
+        split = PatientSplit(ratios=(0.5, 0.1, 0.1, 0.3), seed=3)
+        parts = self.dataset.set_task(EventTask(), split=split)
+        self.assertEqual(len(parts), 4)
+        patients = [self._patients(d) for d in parts]
+        self.assertEqual(patients, _parts(split))
+        for i in range(4):
+            for j in range(i + 1, 4):
+                self.assertFalse(patients[i] & patients[j])
+        self.assertEqual(set().union(*patients), set(_patient_ids()))
+        self.assertEqual(sum(len(d) for d in parts), len(self.dataset.set_task(EventTask())))
+        self.assertEqual(parts[3].fit_split["ratios"], [0.5, 0.1, 0.1, 0.3])
+
+    def test_custom_split(self):
+        split = HoldoutSplit(self.parts[2])
+        train, held = self.dataset.set_task(EventTask(), split=split)
+        self.assertEqual(self._patients(held), self.parts[2])
+        self.assertEqual(self._patients(train), set(_patient_ids()) - self.parts[2])
+        self.assertEqual(held.fit_split, split.to_dict())
+        codes = train.input_processors["codes"]
+        self.assertNotIn("TESTONLY", codes.code_vocab)
+        self.assertIn(codes.code_vocab["<unk>"], held[0]["codes"].tolist())
+
+    def test_custom_split_parts_are_used_as_given(self):
+        everything, first = self.dataset.set_task(EventTask(), split=OverlapSplit())
+        self.assertEqual(len(everything), len(self.dataset.set_task(EventTask())))
+        self.assertEqual(len(first), 5)
+        self.assertEqual(
+            [first[i]["visit_id"] for i in range(5)],
+            [everything[i]["visit_id"] for i in range(5)],
+        )
+
+    def test_unsplit_path_and_cache_key_are_unchanged(self):
+        samples = self.dataset.set_task(EventTask())
+        proc_key = json.dumps(
+            {"input_processors": None, "output_processors": None}, sort_keys=True, default=str
+        )
+        self.assertTrue(samples.path.endswith(f"samples_{uuid.uuid5(uuid.NAMESPACE_DNS, proc_key)}.ld"))
+
+    def test_supplied_processor_is_kept_with_split(self):
+        codes = SequenceProcessor()
+        codes.fit([{"codes": ["A", "B", "Z"]}], "codes")
+        vocab = dict(codes.code_vocab)
+        train, _, test = self.dataset.set_task(
+            EventTask(), input_processors={"codes": codes}, split=SPLIT
+        )
+        self.assertEqual(train.input_processors["codes"].code_vocab, vocab)
+        self.assertLess(train.input_processors["value"].mean, 20.0)
+
+
+class TestSplitBuilder(unittest.TestCase):
+    def _samples(self, n=6):
+        return [
+            {"patient_id": f"p{i}", "codes": ["A"], "value": [float(i)], "label": i % 2}
+            for i in range(n)
+        ]
+
+    def test_warns_when_statistics_are_fitted_on_all_samples(self):
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence", "value": MeanCenter},
+            output_schema={"label": "binary"},
+        )
+        with self.assertLogs("pyhealth.datasets.sample_dataset", "WARNING") as logs:
+            builder.fit(self._samples())
+        self.assertIn("value", "\n".join(logs.output))
+
+    def test_no_warning_with_a_split(self):
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence", "value": MeanCenter},
+            output_schema={"label": "binary"},
+        )
+        with self.assertNoLogs("pyhealth.datasets.sample_dataset", "WARNING"):
+            builder.fit(self._samples(), split=PatientSplit(ratios=(0.5, 0.5)))
+
+    def test_no_warning_for_vocabulary_only(self):
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence"}, output_schema={"label": "binary"}
+        )
+        with self.assertNoLogs("pyhealth.datasets.sample_dataset", "WARNING"):
+            builder.fit(self._samples())
+
+    def test_split_fitting_streams_without_collecting_samples(self):
+        n, payload = 400, 4000  # ~50 MB if every sample were held at once
+
+        class LazyStream:
+            def __iter__(self):
+                for i in range(n):
+                    yield {
+                        "patient_id": f"p{i % 100}",
+                        "codes": ["A", "B"],
+                        "value": [float(i)],
+                        "label": i % 2,
+                        "payload": [float(i)] * payload,
+                    }
+
+            def __len__(self):
+                raise AssertionError("fit must not take len() of the stream")
+
+        builder = SampleBuilder(
+            input_schema={"codes": "sequence", "value": MeanCenter},
+            output_schema={"label": "binary"},
+        )
+        tracemalloc.start()
+        builder.fit(LazyStream(), split=PatientSplit(ratios=(0.7, 0.3), seed=0))
+        _, peak = tracemalloc.get_traced_memory()
+        tracemalloc.stop()
+        self.assertLess(peak, 5 * 2**20, f"fit peaked at {peak / 2**20:.1f} MiB")
+        train_idx = builder.split_indices[0]
+        self.assertEqual(sum(len(p) for p in builder.split_indices), n)
+        expected = float(np.mean([float(i) for i in train_idx]))
+        self.assertAlmostEqual(builder.input_processors["value"].mean, expected)
+
+    def test_invalid_ratios(self):
+        for ratios in ((0.5, 0.6), (1.0,), (0.5, -0.1, 0.6), (0.2, 0.2, 0.2, 0.3)):
+            with self.subTest(ratios=ratios), self.assertRaises(ValueError):
+                PatientSplit(ratios=ratios)
+
+
+class TestSplitSetTaskManyWorkers(unittest.TestCase):
+    """More than 10 workers: merged samples must stay in task order.
+
+    patient_to_index is built before processing, so a merge that reorders the
+    worker blocks (10.index.json before 2.index.json) hands every part the wrong samples.
+    """
+
+    N_WORKERS = 12
+
+    @classmethod
+    def setUpClass(cls):
+        cls.tmp = tempfile.TemporaryDirectory()
+        root = Path(cls.tmp.name) / "data"
+        root.mkdir()
+        cls.parts = _parts(SPLIT)
+        _write_dataset(root, test_patients=cls.parts[2])
+        cls.dataset = BaseDataset(
+            root=str(root),
+            tables=["events"],
+            dataset_name="SplitDatasetManyWorkers",
+            config_path=str(root / "config.yaml"),
+            cache_dir=str(Path(cls.tmp.name) / "cache"),
+            num_workers=1,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.tmp.cleanup()
+
+    def test_split_parts_hold_their_patients(self):
+        parts = self.dataset.set_task(EventTask(), num_workers=self.N_WORKERS, split=SPLIT)
+        got = [{d[i]["patient_id"] for i in range(len(d))} for d in parts]
+        self.assertEqual(got, self.parts)
+
+    def test_patient_to_index_points_at_own_samples(self):
+        samples = self.dataset.set_task(EventTask(), num_workers=self.N_WORKERS)
+        self.assertGreater(len(samples), self.N_WORKERS * 2)
+        for pid, indices in samples.patient_to_index.items():
+            for i in indices:
+                self.assertEqual(samples[i]["patient_id"], pid)
+
+
+if __name__ == "__main__":
+    unittest.main()