Use a new low-memory approach for tf dataset index shuffling
This PR tries out a new approach to generating the index tensor in `to_tf_dataset`, which should reduce memory usage for very large datasets. I'll need to do some testing before merging it!

Fixes #5855