fix: pair each subsampled chunk with its own region of interest in train_test_split

src/litdata/CHANGELOG.md

@@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Pair each subsampled chunk with its corresponding region of interest in `train_test_split`. ([#949](https://github.com/Lightning-AI/litData/pull/949))
+
 - Fix `len(StreamingDataLoader)` undercounting per-worker partial batches with whole-chunk window shuffling. ([#939](https://github.com/Lightning-AI/litData/pull/939))
 
 ## [0.2.76] - 2026-09-30

src/litdata/utilities/train_test_split.py

@@ -50,8 +50,7 @@ def train_test_split(
     # we need subsampled chunk filenames, original chunk file, and subsampled_roi
 
     dummy_streaming_dataset = deepcopy_dataset(streaming_dataset)
-    # Note: We make this a set to speed up the computation of subsampled_chunks.
-    dummy_subsampled_chunk_filename = set(dummy_streaming_dataset.subsampled_files)
+    dummy_subsampled_chunk_filename = list(dummy_streaming_dataset.subsampled_files)
     dummy_subsampled_roi = dummy_streaming_dataset.region_of_interest
     subsampled_chunks: list[dict[str, Any]] = []
 
@@ -64,11 +63,16 @@ def train_test_split(
     # load chunks from `index.json` file
     data = load_index_file(input_dir.path)
 
-    # subsample the chunks defined in the index file
+    # subsample the chunks defined in the index file, keeping the dataset's own
+    # chunk order so that every chunk is paired with its own region of interest
     original_chunks = data["chunks"]
-    subsampled_chunks = [
-        _org_chunk for _org_chunk in original_chunks if _org_chunk["filename"] in dummy_subsampled_chunk_filename
-    ]
+    chunk_by_filename = {_chunk["filename"]: _chunk for _chunk in original_chunks}
+    missing_chunk_filename = next(
+        (_filename for _filename in dummy_subsampled_chunk_filename if _filename not in chunk_by_filename), None
+    )
+    if missing_chunk_filename is not None:
+        raise ValueError(f"Couldn't find chunk file {missing_chunk_filename} in the index file.")
+    subsampled_chunks = [chunk_by_filename[_filename] for _filename in dummy_subsampled_chunk_filename]
 
     new_datasets = [deepcopy_dataset(streaming_dataset) for _ in splits]
 

tests/utilities/test_train_test_split.py

@@ -165,3 +165,35 @@ def test_streaming_dataset_subset(tmpdir):
     assert len(sliced) == 10
     assert sliced[0] == 5
     assert sliced[9] == 14
+
+
+def test_split_a_shuffled_subsampled_dataset(tmpdir):
+    cache = Cache(str(tmpdir), chunk_size=10)
+    for i in range(100):
+        cache[i] = i
+    cache.done()
+    cache.merge()
+
+    ds = StreamingDataset(input_dir=str(tmpdir), subsample=0.45, shuffle=True, seed=1)
+
+    train, test = train_test_split(ds, [0.6, 0.4], shuffle=False)
+
+    # every split keeps the chunks and regions of interest of the dataset it splits
+    split_items = sorted([split[i] for split in (train, test) for i in range(len(split))])
+    dataset_items = sorted(ds[i] for i in range(len(ds)))
+    assert split_items == dataset_items
+
+
+def test_split_a_repeated_subsampled_dataset(tmpdir):
+    cache = Cache(str(tmpdir), chunk_size=10)
+    for i in range(100):
+        cache[i] = i
+    cache.done()
+    cache.merge()
+
+    # a subsample above 1 repeats chunks, so the same filename is sampled more than once
+    ds = StreamingDataset(input_dir=str(tmpdir), subsample=1.5)
+
+    train, test = train_test_split(ds, [0.5, 0.5])
+
+    assert len(train) + len(test) == len(ds) == 150