[Yi-111-a]

Closes #946

Problem

train_test_split() builds subsampled_chunks by filtering the index.json chunk list by filename:

dummy_subsampled_chunk_filename = set(dummy_streaming_dataset.subsampled_files)
...
subsampled_chunks = [
    _org_chunk for _org_chunk in original_chunks if _org_chunk["filename"] in dummy_subsampled_chunk_filename
]

That keeps the chunks in index.json order, but subsampled_files and region_of_interest are built in lockstep by subsample_streaming_dataset() and are in the dataset's own subsampled order. So subsampled_chunks[i] and region_of_interest[i] describe different chunks whenever the dataset is shuffled (shuffle=True), and train_test_split silently returns items from the wrong chunk.

The set(...) has a second effect: when subsample > 1 the dataset samples chunks repeatedly, so subsampled_files contains duplicates. Collapsing them to a set makes subsampled_chunks shorter than region_of_interest, and shuffle_lists_together() then raises ValueError: Lists must be of the same size.

Fix

Look the chunks up by filename, in subsampled_files order, keeping duplicates:

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]

This is the same O(n) lookup the set was there for, but keeps the correspondence that shuffle_lists_together() / subsample_filenames_and_roi() assert on. A missing chunk is now reported with its filename instead of failing later.

Tests

tests/utilities/test_train_test_split.py:

• test_split_a_shuffled_subsampled_dataset — a shuffled, subsampled dataset splits into exactly the items it contains (fails before with mismatched items).
• test_split_a_repeated_subsampled_dataset — subsample=1.5 repeats chunks and splits cleanly (fails before with Lists must be of the same size.).

Both new tests fail on main and pass with this change. tests/utilities/ is otherwise unchanged (the two test_encryption.py failures in my local run are only a missing cryptography wheel).

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 83.33333% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 82%. Comparing base (c6f075c) to head (104c2cb).
⚠️ Report is 1 commits behind head on main.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@          Coverage Diff          @@
##            main    #949   +/-   ##
=====================================
  Coverage     82%     82%           
=====================================
  Files         72      72           
  Lines      16445   16449    +4     
=====================================
+ Hits       13505   13510    +5     
+ Misses      2940    2939    -1     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Copilot]

🟢 Approval recommended

No active review issues were identified for the change preserving correspondence between subsampled chunks and their regions of interest.

No open issues 🎉

What changed in this PR

This PR addresses #946 by preserving the correspondence between subsampled chunks and their regions of interest in train_test_split.

Changes:

• Looks up chunks in dataset order, retaining repeated filenames for subsample > 1.
• Reports missing chunk filenames with a direct error.
• Adds regression tests for shuffled subsamples and repeated chunks.

File	Description
tests/​utilities/​test_train_test_split.py	Adds checks for item preservation with shuffled subsamples and split lengths with repeated chunks.
src/​litdata/​utilities/​train_test_split.py	Preserves chunk order and duplicates when matching regions of interest, and validates missing filenames.

_{Review effort: Balanced}