Is there a way to control the chunksizes during an aligment or reindex operation?

[josephnowak]

After the update 2024.08.02 of Dask the vindex method is keeping the chunksizes consistent, which I think can have a negative impact in certain scenarios on operations like reindex, reindex_like, and align on the Xarray side, for example, if we try to align a DataArray that is a small subset of another one, the number of chunks can increase drastically (see the example for a more clear overview), so my question is if there is a way to control this behavior to prevent the generation of graphs extremely big that affect the performance negatively.

import dask.array as da

n = 10000
m = 10000
big = xr.DataArray(
    da.zeros(shape=(n, m), chunks=(30, 6000)),
    dims=["a", "b"],
    coords={"a": list(range(n)), "b": list(range(m))}
)

n = 10
m = 10
small = xr.DataArray(
    da.zeros(shape=(n, m), chunks=(10, 10)),
    dims=["a", "b"],
    coords={"a": list(range(n)), "b": list(range(m))}
)

# This is going to generate 1000000 chunks
print(small.reindex_like(big).chunksizes)
print(xr.align(big, small, join="outer")[1].chunksizes)

My proposal is that Xarray (not sure if this would be something better on the Dask side but they do not handle indexes as Xarray) could handle the alignment of the chunks in a more "sophisticated" way, it can be through a heuristic that decides the "ideal chunks" of the output, for example, use the biggest chunk of all the arrays as output, and add artificial data before reindexing (probably using the pad method would be ideal) to the small arrays to make them of at least the size of the "ideal chunk", this would guaranty that the number of chunks generated with the operation would be smaller, which would make the performance quite better.

[dcherian]

Not really. This is a bit hard given the modularity of the stack. Can you tell us what kind of behaviour you would like? If you could control the chunk size, what would you set?

cc @phofl

[josephnowak]

Thanks for your response (I thought the forum was not being used so I published this as an issue, I'm going to close it to avoid confusion), I would expect that the smaller array that is going to get the same shape of the bigger array share the same chunksizes to avoid creating many small chunks when the reindex happens.
In the example that I sent the smaller array is creating 1M chunks during the reindex operation, but if we could pad some artificial data before reindexing to make the smaller array of at least the size of the expected chunk and then chunk with -1 on all the padded dims, the reindex operation is going to generate much less tasks.
I think the following code illustrates what I think would be an "ideal" behavior (sorry if it is a little bit long)

def reindex_with_pad(
    data: xr.DataArray | xr.Dataset,
    coords: dict[str, list],
    preferred_chunks: dict[str, int],
    fill_value: Any,
    method: str = None,
) -> xr.Dataset | xr.DataArray:
    """
    This function is going to pad artificial data on the dimensions whose
    size is smaller than the preferred chunk before reindexing
    to reduce the amount of chunks, to achieve that it generates an
    autoincrement index that replace the coord

    """
    if method == "bfill":
        raise NotImplementedError("The bfill method is not implemented for the moment")

    reindex_data = data.copy()
    reindex_mapped_pad_coords = {}
    reindex_mapped_coords = {}
    inv_autoincrement_map = {}
    pad_widths = {}

    for dim, coord in coords.items():
        coord = np.array(coord)
        pad_width = preferred_chunks.get(dim, 0) - data.sizes[dim]
        reindex_mapped_coords[dim] = coord
        if pad_width <= 0:
            continue
        data_coord = data.coords[dim].to_numpy()
        total_coord = np.union1d(np.array(coord), data_coord)

        auto_map = {c: i for i, c in enumerate(total_coord)}
        mapped_data_coord = np.array([auto_map[c] for c in data_coord])
        reindex_mapped_coords[dim] = np.array([auto_map[c] for c in coord])
        reindex_mapped_pad_coords[dim] = np.concatenate(
            (
                mapped_data_coord,
                list(range(len(total_coord), len(total_coord) + pad_width)),
            ),
        )

        reindex_data.coords[dim] = mapped_data_coord
        inv_autoincrement_map[dim] = {v: k for k, v in auto_map.items()}
        pad_widths[dim] = (0, pad_width)

    if not reindex_mapped_pad_coords:
        # It is not necessary to pad data, reindex as usual
        return data.reindex(coords, fill_value=fill_value, method=method)

    # It can also be done using reindex directly but it generates more tasks
    reindex_data = reindex_data.pad(pad_widths, mode="edge")
    for dim, coord in reindex_mapped_pad_coords.items():
        reindex_data.coords[dim] = coord
    # Create a single chunk on the padded dimensions
    reindex_data = reindex_data.chunk(
        {dim: -1 for dim in reindex_mapped_pad_coords.keys()}
    )
    # Reindex with expected coords, this is going to automatically drop the padded data
    reindex_data = reindex_data.reindex(
        reindex_mapped_coords, method=method, fill_value=fill_value
    )

    # Revert the mapping
    for dim in reindex_mapped_pad_coords.keys():
        reindex_data.coords[dim] = np.array(
            [inv_autoincrement_map[dim][c] for c in reindex_data.coords[dim].to_numpy()]
        )
    return reindex_data

[phofl]

We do have an option that allows you to set the increase in chunk size you are willing to tolerate on the Dask side. array.shuffle.chunksize-tolerance. That might only get you part of the way there.

I've recently been thinking a lot if we should just aim for something like auto, i.e. always 128MiB or something similar. But Dask still has a few places where chunk sizes are blowing up accidentally (dask/dask#11426), so I took the path of preserving the input chunks.

This is indeed a pattern I hadn't really considered, I'll think about this a bit more

[josephnowak]

From my perspective, I think that the idea of preserving the input chunks is ideal for a default option (probably also for most of the cases where Dask is used), but probably it would be good if the user could set the "ideal" chunks to use:

# Use "with" because vindex does not accept parameters
with dask.config.set({"ideal_chunk": (N, M, ..., Z)}):
      arr_with_ideal_chunk_size = arr.vindex[...]