Function _shuffle

dask/array/_shuffle.py:181–310

(chunks, indexer, axis, in_name, out_name, token)

def _shuffle(chunks, indexer, axis, in_name, out_name, token):
    _validate_indexer(chunks, indexer, axis)

    if len(indexer) == len(chunks[axis]):
        # check if the array is already shuffled the way we want
        ctr = 0
        for idx, c in zip(indexer, chunks[axis]):
            if idx != list(range(ctr, ctr + c)):
                break
            ctr += c
        else:
            return chunks, {}

    chunksize_tolerance = config.get("array.chunk-size-tolerance")
    chunk_size_limit = int(sum(chunks[axis]) / len(chunks[axis]) * chunksize_tolerance)

    # Figure out how many groups we can put into one chunk
    current_chunk, new_chunks = [], []
    for idx in indexer:
        if len(current_chunk) + len(idx) > chunk_size_limit and len(current_chunk) > 0:
            new_chunks.append(current_chunk)
            current_chunk = idx.copy()
        else:
            current_chunk.extend(idx)
            if len(current_chunk) > chunk_size_limit / chunksize_tolerance:
                new_chunks.append(current_chunk)
                current_chunk = []
    if len(current_chunk) > 0:
        new_chunks.append(current_chunk)

    # force 64 bit to avoid potential integer overflows on win32 and numpy<2
    chunk_boundaries = np.cumsum(np.array(chunks[axis], dtype="uint64"))

    # Get existing chunk tuple locations
    chunk_tuples = list(
        product(*(range(len(c)) for i, c in enumerate(chunks) if i != axis))
    )

    intermediates = dict()
    merges = dict()
    dtype = np.min_scalar_type(max(*chunks[axis], chunk_size_limit))
    split_name = f"shuffle-split-{token}"
    slices = [slice(None)] * len(chunks)
    split_name_suffixes = count()
    sorter_name = "shuffle-sorter-"
    taker_name = "shuffle-taker-"

    old_blocks = {
        old_index: (in_name,) + old_index
        for old_index in np.ndindex(tuple([len(c) for c in chunks]))
    }
    for new_chunk_idx, new_chunk_taker in enumerate(new_chunks):
        new_chunk_taker = np.array(new_chunk_taker)
        sorter = np.argsort(new_chunk_taker).astype(dtype)
        sorter_key = None

        sorted_array = new_chunk_taker[sorter]
        source_chunk_nr, taker_boundary = np.unique(