def generate_roundrobin_read_list(num_files, world_size):
    """
    Generate the file IDs to read for each rank
    using round robin assignment.

    Parameters:
    -----------
    num_files : int
        Total number of files.
    world_size : int
        World size of group.

    Returns:
    --------
    read_list : np.array
        Array of target file IDs to read. Each worker is expected
        to read the list of file indexes in its rank's index in the list.
        e.g. rank 0 reads the indexed in read_list[0], rank 1 the
        ones in read_list[1] etc.

    Examples
    --------
    >>> tools.distpartitionning.utils.generate_roundrobin_read_list(10, 4)
    [[0, 4, 8], [1, 5, 9], [2, 6], [3, 7]]
    """
    assignment_lists = [[] for _ in range(world_size)]
    for rank, part_idx in zip(cycle(range(world_size)), range(num_files)):
        assignment_lists[rank].append(part_idx)

    return assignment_lists

    Generate the file IDs to read for each rank
    using round robin assignment.

    Parameters:
    -----------
    num_files : int
        Total number of files.
    world_size : int
        World size of group.

    Returns:
    --------
    read_list : np.array
        Array of target file IDs to read. Each worker is expected
        to read the list of file indexes in its rank's index in the list.
        e.g. rank 0 reads the indexed in read_list[0], rank 1 the
        ones in read_list[1] etc.

    Examples
    --------
    >>> tools.distpartitionning.utils.generate_roundrobin_read_list(10, 4)
    [[0, 4, 8], [1, 5, 9], [2, 6], [3, 7]]
    
tools/distpartitioning/utils.py