train_test_split pairs chunks with the wrong region_of_interest for shuffled subsampled datasets (and fails for subsample > 1) #946

🐛 Bug

litdata.train_test_split rebuilds the chunk list by filtering index.json with set(dataset.subsampled_files). That list follows the index.json order and has duplicates removed, while dataset.region_of_interest keeps the order of subsampled_files. The two lists are then paired by position.

• With StreamingDataset(..., subsample=<1, shuffle=True), subsampled_files is in shuffled order, so chunks get another chunk's ROI. The resulting splits silently include items that are not in the input dataset and miss items that are.
• With subsample > 1, chunks appear more than once in subsampled_files, the set collapses them, and the two lists have different lengths: ValueError: Lists must be of the same size. from shuffle_lists_together (or, with shuffle=False, the assert len(chunks) == len(roi_list) in subsample_filenames_and_roi).

To Reproduce

import tempfile
from litdata import StreamingDataset, optimize, train_test_split

def fn(i):
    return i

def items(ds):
    return sorted(ds[i] for i in range(len(ds)))

if __name__ == "__main__":
    out = tempfile.mkdtemp()
    optimize(fn=fn, inputs=list(range(100)), output_dir=out, chunk_size=10, num_workers=1)

    ds = StreamingDataset(out, subsample=0.45, shuffle=True, seed=1)
    a, b = train_test_split(ds, [0.6, 0.4], shuffle=False)
    union = sorted(items(a) + items(b))
    print("not in input:", sorted(set(union) - set(items(ds))))
    print("missing     :", sorted(set(items(ds)) - set(union)))

    ds2 = StreamingDataset(out, subsample=1.5)
    train_test_split(ds2, [0.5, 0.5])

Output on main (c6f075c):

not in input: [65, 66, 67, 68, 69]
missing     : [95, 96, 97, 98, 99]
...
  File ".../litdata/utilities/subsample.py", line 23, in shuffle_lists_together
    raise ValueError("Lists must be of the same size.")
ValueError: Lists must be of the same size.

Expected behavior

The splits together contain exactly the items of the input dataset, for any subsample / shuffle setting of the input.

Additional context

A possible fix is to build the chunk list from subsampled_files directly, through a filename -> chunk dict (so the lookup stays O(n) as in #647), keeping its order and repetitions. Behaviour for non-shuffled inputs is unchanged. I have a small patch with a regression test and can open a PR if that approach is OK.

Environment: litdata main @ c6f075c, Python 3.12, macOS, CPU only.

Disclosure: found, reproduced and the draft fix tested locally with the help of an AI coding agent (Claude Code); reviewed before filing.