How to make sure batches are from 1 dataset

Hi
I have multiple dataset, then I want to concat them and have a distributed sampler on top of it, I need to make sure in each batch, data comes from 1 type of dataset, how could I do it? could you provide me with sample codes, thnaks