Fix performance bottleneck in train_test_split

[lukemerrick]

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Speeds up the train_test_split function in cases where the large number of chunks makes repeated membership tests against a list painfully slow. The fix is simply to convert dummy_subsampled_chunk_filename into a set for O(1) lookup complexity.

The core bottleneck is here:

litData/src/litdata/utilities/train_test_split.py

Lines 62 to 64 in 1a3c1c1

	subsampled_chunks = [
	_org_chunk for _org_chunk in original_chunks if _org_chunk["filename"] in dummy_subsampled_chunk_filename
	]

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Issue: #648

Did you have fun?

Unfortunately finding this issue was not a fun process, since it seemed like this function is the last place to expect a performance regression, but fixing it was thankfully not a lot of work!

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 83%. Comparing base (68bcbe6) to head (18b64af).
Report is 1 commits behind head on main.

Additional details and impacted files

@@         Coverage Diff         @@
##           main   #647   +/-   ##
===================================
  Coverage    83%    83%           
===================================
  Files        49     49           
  Lines      6756   6756           
===================================
  Hits       5636   5636           
  Misses     1120   1120           

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[deependujha]

nice catch @lukemerrick. Thanks for contributing!🎉

[bhimrazy]

🚀

[tchaton]

Nice !