tf.data parallel filter dataset for each class and interleave not giving expected speedup 
@niemiaszek 
In order to expedite the trouble-shooting process, please provide a complete code snippet to reproduce the issue reported here. I have tried to replicate [this](https://colab.research.google.com/gist/sushreebarsa/29378816f9ccca87549cb6617278aa8d/61245.ipynb#scrollTo=EZfM7r8QhGC7) issue and faced different errors. Thank you! Sorry for no response so far. I will prepare standalone code with some sample TFRecords prepared on public dataset. Then it should be easier to fiddle with this issue. For now I'm quite busy with other stuff and I was hoping for some expertise just based on code snippet, as my main question (how to efficiently filter `tf.data.Dataset` by classes) seem to be rather general than case-dependent.

  @niemiaszek Thank you for your response!
Though, in order to analyze the issue we need the complete code, you may perform filters on batches. Either unbatch and use tf.data.Dataset.filter or just use tf.data.Dataset.map which is more preferable.  As an example you may refer to this following;
```
ds = tfds.load('cats_vs_dogs', split='train', as_supervised=True)
ds = ds.filter(lambda img, label: label == 1)
```

When using the [Dataset.map](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#map), and [Dataset.filter](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#filter) transformations, which apply a function to each element, the element structure determines the arguments of the function. Please refer to this [document](https://www.tensorflow.org/api_docs/python/tf/data/Dataset). 
Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61245">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61245">No</a>

### Issue type

Performance

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.12

### Custom code

Yes

### OS platform and distribution

ubuntu 22.04

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.8/8.6

### GPU model and memory

_No response_

### Current behavior?

I've tried to get some help on [Tensorflow Forum](https://discuss.tensorflow.org/t/dataset-with-tf-data-with-batches-of-randomly-n-chosen-speakers-and-their-m-utterances/17824) but I've got no response so I decided to post an issue here. Post on TFRecords Forum holds extensive description of my problem (tf.data dataloader for GE2E loss), but I'll try to trim it down only to technical matter.

I'm working on implementation of dataloader, where I want to get (n_speakers x n_utters) samples in every batch. I've already prepared dataset as 10 sharded TFRecord (as suggested in docs) and I wanted to build pairing mechanism using tf.data API. I've menaged to do so with `tf.data.Dataset.choose_from_datasets`, but this requires me to have a list of datasets, where each dataset contains only samples from one speaker. To prepare such dataset I've filtered my one big TFRecords Dataset with each speaker_id. My implementation feels criminal, because it iterates through whole dataset as many times as there are speakers. I've tried to help it with parallel calls (mapping the filter function with list of speakers), but there was no speedup. I've also tried using interleave (setting only `num_parallel_calls` to `AUTOTUNE`) on my dataset parsing, but all I got was slower execution time.

Is if there is any easy and reasonable solution to this filtering issue (to filter whole dataset by each class in one pass)? My only idea to skip filtering at all is to separate each speaker in different TFRecord at the moment of dataset creation. Also how should I use interleave with my code? 

Currently it takes a long time to make a first pass over dataset (~10min for 10GB of TFRecords, next iters when filtered datasets get cached its 2s) - it can be an issue with training with bigger dataset I suppose.

I'm attaching the relevant code but it's not standalone unfortunately, as I can not attach used data. Also some parsing code is missing, but I think it's not that relevant - my examples are stored as melspectrograms with speaker_id label. I could mock some data if someone was interested in helping me.




### Standalone code to reproduce the issue

```shell
dataset = tf.data.Dataset.from_tensor_slices(files)
tfrecords_reader = functools.partial(
    self._tfrecords_reader,
    speaker_id=self._speaker_id,
    speaker_gender=self._speaker_gender,
    speaker_id_key=self._speaker_id_key,
    speaker_gender_key=self._speaker_gender_key,
)

dataset = tf.data.TFRecordDataset(dataset)
dataset = dataset.map(tfrecords_reader, num_parallel_calls=AUTOTUNE)

# NOT PARALLEL WAY - surprisingly it takes as long as parallel method...
# datasets = []
# for speaker_id in self._speaker_list:
#     datasets.append(
#         dataset.filter(lambda x, y: y["speaker_id"] == speaker_id)
#         .cache()
#         .shuffle(BUFFER_SIZE, reshuffle_each_iteration=True)
#         .repeat()
#     )

# Filtering for speakers
speakers = tf.data.Dataset.from_tensor_slices(self._speaker_list)
dataset = speakers.map(
    lambda speaker_id: dataset.filter(lambda x, y: y["speaker_id"] == speaker_id)
    .cache()
    .shuffle(BUFFER_SIZE, reshuffle_each_iteration=True)
    .repeat(),
    num_parallel_calls=AUTOTUNE,
)

datasets_idx = tf.data.Dataset.range(len(self._speaker_list))
# Define those values statically for uniform_candidate_sampler which requires static int, even tho it looks silly
len_speakers = len(speakers)
len_pairs = len_speakers - 1

def pair_speakers(x):
    """Pair all speakers in dataset.
    
    Pairing is done in a way where each speaker is taken at least once.
    Also we can guarantee that those pairs are being re-generated each time.
    """
    to_pair = tf.reshape(tf.where(tf.range(len(speakers)) != x), [1, -1])  # Exclude speaker x from being drawed
    pairs = tf.random.uniform_candidate_sampler(to_pair, len_pairs, self._num_speakers - 1, True, len_speakers)[
        0
    ]  # sample n_speakers-1 pairings without replacement from len(speakers)-1 speakers
    paired = tf.concat([[x], pairs], 0)  # concat speaker x with its pairs
    return tf.repeat(
        paired, tf.repeat(self._num_utters, self._num_speakers)
    )  # repeat each speaker count by selected number of utterances

choice_dataset = datasets_idx.map(
pair_speakers, num_parallel_calls=AUTOTUNE
)  # [mapping of len(speakers) datasets, each num_speakers*num_utters]
choice_dataset = choice_dataset.flat_map(
    tf.data.Dataset.from_tensor_slices
)  # [single dataset of len(speakers)*num_speakers*num_utters,]

# Deterministicly takes samples from filtered datasets with order specified in choice dataset
dataset = tf.data.Dataset.choose_from_datasets(dataset, choice_dataset).batch(
    self._num_speakers * self._num_utters
)

# Apply transforms to batched dataset spectrograms. Faster transforms than on single examples
dataset.map(
    lambda x, y: (self._transform(x), y),
    num_parallel_calls=tf.data.experimental.AUTOTUNE,
)

dataset = dataset.prefetch(AUTOTUNE)
total_steps = len(self._speaker_list)

return dataset, total_steps
```


### Relevant log output

_No response_