A basic default ChunkManager for arrays that report their own chunks

[hmaarrfk]

Is your feature request related to a problem?

I'm creating duckarrays for various file backed datastructures for mine that are naturally "chunked". i.e. different parts of the array may appear in completely different files.

Using these "chunks" and the "strides" algorithms can better decide on how to iterate in a convenient manner.

For example, an MP4 file's chunks may be defined as being delimited by I frames, while images stored in a TIFF may be delimited by a page.

So for me, chunks are not so useful for parallel computing, but more for computing locally and choosing the appropriate way to iterate through a large arrays (TB of uncompressed data).

Describe the solution you'd like

I think a default Chunk manager could simply implement compute as np.asarray as a default instance, and be a catchall to all other instances.

Advanced users could then go in an reimplement their own chunkmanager, but I was unable to use my duckarrays that incldued a chunk property because they weren't associated with any chunk manager.

Something as simple as:

diff --git a/xarray/core/parallelcompat.py b/xarray/core/parallelcompat.py
index c009ef48..bf500abb 100644
--- a/xarray/core/parallelcompat.py
+++ b/xarray/core/parallelcompat.py
@@ -681,3 +681,26 @@ class ChunkManagerEntrypoint(ABC, Generic[T_ChunkedArray]):
         cubed.store
         """
         raise NotImplementedError()
+
+
+class DefaultChunkManager(ChunkMangerEntrypoint):
+    def __init__(self) -> None:
+        self.array_cls = None
+
+    def is_chunked_array(self, data: Any) -> bool:
+        return is_duck_array(data) and hasattr(data, "chunks")
+
+    def chunks(self, data: T_ChunkedArray) -> T_NormalizedChunks:
+        return data.chunks
+
+    def compute(self, *data: T_ChunkedArray | Any, **kwargs) -> tuple[np.ndarray, ...]:
+        raise tuple(np.asarray(d) for d in data)
+
+    def normalize_chunks(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def from_array(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def apply_gufunc(self, *args, **kwargs):
+        raise NotImplementedError()

Describe alternatives you've considered

I created my own chunk manager, with my own chunk manager entry point.

Kinda tedious...

Additional context

It seems that this is related to: #7019

[hmaarrfk]

Even removing the "abstractmethod" decorator would be welcome ;)

[TomNicholas]

Hi @hmaarrfk ! I'm excited to see someone is actually using the ChunkManager system 😀

I'm curious to hear more about your use case - do you have the code for your custom ChunkManager somewhere?

I'm not totally sure I understand what the purpose of a "default" chunkmanager would be - IIUC you're trying to expose your underlying array's .chunks attribute but treat it like a normal duckarray in all other respects?

Even removing the "abstractmethod" decorator would be welcome ;)

Which one? Or you mean all of them except literally just __init__ and chunks?

[hmaarrfk]

do you have the code for your custom ChunkManager somewhere?

I've provided the patch that is essentially my code for the chunk manager ;).

I am a little stricter and I check for isinstance instead of hasattr because i don't want to overwrite the dask chunk manager you have all written. but if you were to have a "default" chunkmanager it could be more forgiving with the hasattr check since it would be the last chunkmanager to get checked.

IIUC you're trying to expose your underlying array's .chunks attribute but treat it like a normal duckarray in all other respects?

100% correct. Its my own implementation of "lazy access" to TIFF and MP4s as read-only arrays. It really only implements the very features I need so it isn't really "full feature complete". I've found this necessary to implement myself for performance reasons.

Or you mean all of them except literally just init and chunks?

I mean all of them except __init__.

• def chunks could be implemented as data.chunks and allow overwriting by subclasses.
• def compute could be implemented in the suggested as tuple(np.asarray(d) for d in data) and allow optional overwriting by subclasses.

edit: reordered some paragraphs

[TomNicholas]

Okay overall this seems like a good idea to me! Basically allowing for the idea that there are really three types of arrays: un-chunked duck arrays, chunked duck arrays where the chunks might be read but won't be changed, and re-chunkable duck arrays where the re-chunking happens via some (probably parallel) processing framework.

I am a little stricter and I check for isinstance instead of hasattr because i don't want to overwrite the dask chunk manager you have all written. but if you were to have a "default" chunkmanager it could be more forgiving with the hasattr check since it would be the last chunkmanager to get checked.

Sorry I'm not quite following this part. You're talking about inside is_chunked_array? I don't think anything should be relying on the order of which chunkmanager is checked against first.

[TomNicholas]

Also note there are two senses of the idea of a "default" chunkmanager to keep track of here: (1) what is used to request the .chunks attribute of a duckarray of unspecified type, and (2) which chunkmanager to use to .chunk an array that was not previously chunked (this has to be dask for backwards compatibility reasons).

[hmaarrfk]

Sorry I'm not quite following this part.

This is my code ;)

import numpy as np
from xarray.core.parallelcompat import ChunkManagerEntrypoint
from xarray.core.types import T_NormalizedChunks
from ._my_array import MyArray1, MyArray2

class MyChunkManager(ChunkManagerEntrypoint["MyArrayEntryPoint"]):
    array_cls: None
    available: bool = True

    def __init__(self) -> None:
        self.array_cls = (MyArray1, MyArray2)

    def chunks(self, data) -> T_NormalizedChunks:
        return data.chunks

    def compute(self, *data, **kwargs) -> tuple[np.ndarray, ...]:
        return tuple(np.asarray(d) for d in data)

    def apply_gufunc(self, *args, **kwargs):
        raise NotImplementedError()

    def from_array(self, *args, **kwargs):
        raise NotImplementedError()

    def normalize_chunks(self, *args, **kwargs):
        raise NotImplementedError()

Also note there are two senses of the idea of a "default" chunkmanager to keep track of here

You are correct. To be explicit, my problem is that implementing a whole chunkmanager, that allows for rechunking and many other fancy operations, seems a little overkill when I (or other users) try to expose a .chunks property. If I add a .chunks property without building out my ChunkManager Machinery, I get the following traceback

  File "/home/mark/git/xarray/xarray/core/dataset.py", line 842, in load
    return dataset.isel(isel).compute()
  File "/home/mark/git/xarray/xarray/core/dataset.py", line 1011, in compute
    chunkmanager = get_chunked_array_type(*lazy_data.values())
  File "/home/mark/git/xarray/xarray/core/parallelcompat.py", line 142, in get_chunked_array_type
    return new.load(**kwargs)
  File "/home/mark/git/xarray/xarray/core/dataset.py", line 842, in load
    raise TypeError(
TypeError: Could not find a Chunk Manager which recognises type <class 'mark._my_array.MyArray1'>

I'm sorry for jumping to a proposal for a solution. I'm trying to get better at framing my asks overall (even beyond open source).

[hmaarrfk]

The other reason why I think it is good for xarray to consider, is that I was pretty close to avoiding exposing .chunks because understanding the entrypoints is quite an "advanced python packaging" thing.

It is clever and avoids circular dependencies and allows for plugins to be built.

But a developer that writes python code and not packaging code it is quite difficult....

[TomNicholas]

I'm sorry for jumping to a proposal for a solution. I'm trying to get better at framing my asks overall (even beyond open source).

No this is great! Some initiative in thinking of possible solutions is always appreciated :)

---

So there's arguably a more general issue here: If you have a new duckarray type, xarray will happily wrap that and use it without you explicitly telling xarray what type it is (hence the term "duck-typed".) But currently if you have a duckarray type that also implements a .chunks attribute, xarray requires you to explicitly register that type using the chunkmanager system. This is unnecessarily restrictive for the case of arrays that are just duckarrays + .chunks.

The reason we have chunkmanagers at all is for when chunked computation requires calling a function like apply_gufunc, for which we need to know if the underlying type is dask or cubed or whatever, so that we can use the correct import of dask.apply_gufunc or cubed.apply_gufunc. But your use case doesn't need this complexity.

---

One solution would be to just remove .chunks from the ChunkManager entirely, and instead just call .chunks on the underlying array immediately. This would work, but it kind of breaks the nice way we now have all chunk-related stuff go through the ChunkManager class. Another disadvantage of this would be if you had an array with an attribute that was like .chunks but not actually called .chunks, we couldn't re-direct it.

Another solution is similar to what you're suggesting: define a default ChunkManager (called DuckArrayChunkManager or something) that forwards .chunks on but has no other methods defined. However this needs to work for any compliant array type, not just for MyArray. I guess that's what you were referring to here:

if you were to have a "default" chunkmanager it could be more forgiving with the hasattr check since it would be the last chunkmanager to get checked.

[TomNicholas]

FYI @andersy005

[TomNicholas]

@hmaarrfk would you be interested in submitting a PR to add the "default" chunk manager?

[dcherian]

Is your ask simply that we forward .chunks to .data.chunks, if available? That seems OK and a lot simpler than adding a "default chunk manager"

[TomNicholas]

Yeah we might have fallen down a rabbit hole here...