Explosive memory use with .sel(..) depending on number of chunks?

[jonaslb]

So I'm posting this as a discussion because I'm guessing the issue is "really" with dask, not xarray, but still I hope that someone has some insights into why this happens and what can be done to avoid it.

I have the following example code:

import dask.array as da
import numpy as np
import pandas as pd
import xarray as xr


def make_dataset() -> xr.DataArray:
    # Dataset coordinates: time, ens, step, latitude, longitude
    times = xr.DataArray(
        pd.date_range(
            start=pd.Timestamp("2022-01-01"),  # type: ignore[index]
            end=pd.Timestamp("2025-01-02"),  # type: ignore[index]
            freq="6h",
        ),
        dims=("time",),
    )
    steps = xr.DataArray(
        pd.timedelta_range(
            start="0h",
            end="3D",
            freq="1h",
        ),
        dims=("step",),
    )
    ens = xr.DataArray(
        np.arange(0, 51),
        dims=("ens",),
    )

    # with 0.1 resolution: data.sel(..) is fast and peak mem is 329 MB
    # with 0.01 resolution: data.sel(..) is slow and peak mem is >14 GB (OOM)
    lat = xr.DataArray(np.arange(80, 40, -0.01), dims=("latitude",))
    lon = xr.DataArray(np.arange(-30, 40, 0.01), dims=("longitude",))
    data = xr.DataArray(
        da.ones(
            # Just some example data, in reality this would be from stored data
            (len(times), len(steps), len(lat), len(lon), len(ens)),
            # changing chunks for lat/lon to eg 200, 200 decreases memory a lot
            # ie memory seems to scale with chunk count not data size
            chunks=(1, 8, 100, 100, -1),
        ),
        dims=("time", "step", "latitude", "longitude", "ens"),
        coords={
            "time": times,
            "step": steps,
            "latitude": lat,
            "longitude": lon,
            "ens": ens,
        },
    )
    return data


print("Preparing data")
data = make_dataset()
print("selecting some time range")
data = data.sel(time=slice(pd.Timestamp("2024-01-01"), pd.Timestamp("2024-07-01")))
print("selecting a subset of times")
data = data.sel(time=data.coords["time"].dt.hour.isin([0, 12]))
print("created a dataarray with filtered times")

It can be run like systemd-run --same-dir --user --wait --pty --collect --property=MemoryMax=23G --property=MemorySwapMax=0G -- uv run ./test_for_xarray.py (set those limits as you like to avoid nuking your pc).

I'm wondering why it has such a dramatic effect to do this filtering? I feel that it should be "simple" to do this filtering, certainly not in the tens of GB of memory to do it?

[jonaslb]

I created a dask issue about bad slicing performance: dask/dask#12020