[Data] Dataset.train_test_split reads dataset twice

[bveeramani]

What happened + What you expected to happen

I'm following the Fine-tuning a Torch object detection model example, and noticed when I call train_test_split the dataset is read twice.

I think it's because we call Dataset.count() in the method's implementation, and for non-trivial pipelines this causes the pipeline to execute.

ray/python/ray/data/dataset.py

Line 2131 in d8835df

ds_length = ds.count()

Versions / Dependencies

902b55a

Reproduction script

See descripition.

Issue Severity

Low: It annoys or frustrates me.