pydata / pydata/xarray

map_blocks should dispatch to ChunkManager

Open
#8,545 5 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

topic-chunked-arrays
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

Is your feature request related to a problem?

#7019 generalized most of xarrays internals to be able to use any chunked array type that we can create a ChunkManagerEntrypoint for. Most functions now go through this (e.g. apply_ufunc), but I did not redirect xarray.map_blocks to go through ChunkManagerEntrypoint.

This redirection works by dispatching to high-level dask.array primitives such as dask.array.apply_gufunc, dask.array.blockwise, and dask.array.map_blocks. However the current implementation of xarray.map_blocks is much lower-level, building a custom HLG, so it was not obvious how to swap it out.

Describe the solution you'd like

I would like to either:

  1. Replace the current internals of xarray.map_blocks with a simple call to ChunkManagerEntrypoint.map_blocks. This would be the cleanest separation of concerns we could do here. Presumably there is some obvious reason why this cannot or should not be done, but I have yet to understand what that reason is. (either @dcherian or @tomwhite can you enlighten me perhaps? 🙏)

  2. (More likely) refactor so that the existing guts of xarray.map_blocks are only called from the ChunkManagerEntrypoint, and a non-dask chunked array (i.e. cubed, but in theory other types too) would be able to specify how it wants to perform the map_blocks.

Describe alternatives you've considered

Leaving it as the status quo breaks the nice abstraction and separation of concerns that #7019 introduced.

Additional context

Split off from https://github.com/pydata/xarray/issues/8414

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 with the current xarray.map_blocks implementation and the ChunkManagerEntrypoint abstraction introduced in #7019. Compare its custom high-level-graph construction with the dask.array primitives mentioned in the issue, including apply_gufunc, blockwise, and map_blocks. Done means map_blocks dispatches through ChunkManagerEntrypoint while preserving support for existing dask behavior and allowing other chunked array types to provide their own implementation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.