pydata / pydata/xarray

xarray - dask `map_blocks` duplicates coordinate tasks across calls on the same object

Open
#11,588 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug needs triage
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

What happened?

Using xarray with dask:
Calling map_blocks repeatedly on the same object puts a fresh private copy of
that object's coordinate arrays into the graph for every call, even though
the coordinate data is identical each time. The graph therefore grows linearly
in the number of calls.

This is distinct from #8409 ("Task graphs on .map_blocks with many chunks can
be huge"
), which was about full indexes being duplicated per chunk within one
call
. That is fixed — one call over 400 blocks now correctly produces one task
per distinct coordinate slice. The remaining problem is duplication across
calls
.

Cause

In subset_dataset_to_block, the key for a non-dask (coordinate) variable is
built from gname:

chunk_variable_task = (
    f"{name}-{gname}-{dask.base.tokenize(subsetter)}",
) + this_var_chunk_tuple

and gname is

gname = f"{dask.utils.funcname(func)}-{dask.base.tokenize(npargs[0], args, kwargs)}"

so it varies with func, args and kwargs. None of those affect the value of
a coordinate slice — that is fully determined by the variable and the subsetter.
Two calls differing only in a non-dask argument therefore produce disjoint
coordinate keys holding byte-identical data.

The keys make this visible. For a 9000x9000 grid chunked at 2250, two calls
differing only in one argument:

call 1:  x-lambda-c0d85a4de6a237ed6c335850717be8dd-223df680899f846c2ce66989f7775642
call 2:  x-lambda-b1806737c4170931cfb492be030d770f-223df680899f846c2ce66989f7775642
                  ^^^^^^^^ gname differs ^^^^^^^^  ^^^^ subsetter token identical ^^^^
What did you expect to happen?

Coordinate tasks should be shared between calls, since their content does not
depend on func, args or kwargs. x-tasks should stay at 20 regardless of
how many times map_blocks is called.

Minimal Complete Verifiable Example
import numpy as np, xarray as xr, dask.array as da, cloudpickle

PX, CH = 45000, 2250
arr = da.zeros((PX, PX), chunks=(CH, CH), dtype="float32")
xa = xr.DataArray(
    arr,
    coords={"y": np.arange(PX) * -10.0, "x": np.arange(PX) * 10.0},
    dims=["y", "x"],
    name="m",
)

def build(n_calls):
    graph = {}
    for i in range(n_calls):
        out = xa.map_blocks(lambda b, s: b, [[f"arg_{i}"]], template=xa)
        graph.update(dict(out.data.dask))
    return graph

for n in (1, 4, 16):
    g = build(n)
    xs = sum(1 for k in g if str(k[0]).startswith("x-"))
    print(f"{n:>2} calls: x-tasks={xs:>4}  graph={len(cloudpickle.dumps(g))/1e6:>6.2f} MB")
 1 calls: x-tasks=  20  graph=  0.97 MB
 4 calls: x-tasks=  80  graph=  3.80 MB
16 calls: x-tasks= 320  graph= 15.12 MB

The grid has 20 chunks per dimension, so 20 x-tasks is correct for one call.
Each holds a 2250-element float64 slice (~18 KB). Across calls the same 20
slices are re-materialised every time.

MVCE confirmation
  • [x ] Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
  • [ x] Complete example — the example is self-contained, including all data and the text of any traceback.
  • [ x] Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
  • [ x] New issue — a search of GitHub Issues suggests this is not a duplicate.
  • [x ] Recent environment — the issue occurs with the latest version of xarray and its dependencies.
Proposed fix

Key the coordinate task on the content that determines it:

chunk_variable_task = (
    f"{name}-{dask.base.tokenize(variable, subsetter)}",
) + this_var_chunk_tuple

Uniqueness is preserved — different coordinates hash apart, and tokenizing
variable (not just subsetter) is what keeps two same-named but different
coordinates from colliding. Identical slices now collide deliberately, which is
what a content-addressed cache key should do.

With that change:

 1 calls: x-tasks=  20  graph=  0.94 MB
 4 calls: x-tasks=  20  graph=  1.49 MB
16 calls: x-tasks=  20  graph=  3.73 MB

4.1x smaller at 16 calls, and flat in coordinate tasks. test_dask.py
(222 passed, 2 skipped, 3 xfailed) and test_groupby.py (441 passed, 64
skipped) are unaffected.

Why this matters in practice

Any pipeline that calls map_blocks once per group over a shared grid hits
this. In our case a Sentinel-1 mosaicking step rasterises masks per orbit and
per acquisition date, which is 59 map_blocks calls over one 45000x45000 grid.
Measured on the real graph:

tasks bytes
coordinate tasks needed ~40 ~0.7 MB
coordinate tasks shipped 2,543 44.6 MB

That was 38% of a 117.6 MB task graph, and the graph took the scheduler
long enough to process (repeated 15s event-loop stalls) that workers timed out
their heartbeats and were lost.

Measured on that real graph, same inputs both sides, xarray main with and
without the change:

before after
x-coord tasks 1,180 67
y-coord tasks 1,363 58
total tasks 665,832 663,414
serialized graph 117.6 MB 73.8 MB

A 37% smaller graph. The drop in total tasks (2,418) is exactly the drop in
coordinate tasks (2,543 - 125), i.e. nothing else in the graph changed.

Environment

xarray main (8de862c), also reproduced on 2026.7.0; dask 2026.7.1;
Python 3.12 and 3.14.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at the subset_dataset_to_block entry point, where coordinate task keys are assembled from gname and the variable subsetter. Review test_dask.py first, then verify that repeated map_blocks calls share coordinate tasks while distinct coordinates remain unique. Done means the regression is covered and the reported x-task counts stay flat across calls without breaking the existing dask and groupby tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
78/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.