Use entrypoint.name as the key for registering chunk key encodings.

- Change register_chunk_key_encoding function to take key as first arg
similar to codec.

Author: RFLeijenaar

run.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+
+import numpy as np
+
+import zarr
+import zarr.codecs
+from zarr.core.chunk_key_encodings import ChunkKeyEncoding
+from zarr.registry import register_chunk_key_encoding
+from zarr.storage import LocalStore, WrapperStore
+# from numcodecs.zarr3 import
+
+
+class FanoutStore(WrapperStore):
+    pass
+
+
+class FanoutChunkKeyEncoding(ChunkKeyEncoding):
+    name = "fanout"
+
+    def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]:
+        if chunk_key == "c":
+            return ()
+        return tuple(map(int, chunk_key[1:].split(self.separator)))
+
+    def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
+        return self.separator.join(map(str, ("c",) + chunk_coords))
+
+register_chunk_key_encoding("fanout", FanoutChunkKeyEncoding)
+
+
+def main():
+    store = Path("~/data/test.zarr").expanduser()
+    store = FanoutStore(LocalStore(store))
+
+    # z = zarr.open_array(store)
+
+    z = zarr.create_array(
+        store,
+        shape=(100_000, 3),
+        # shards=(1_000_000, 3),
+        chunks=(1_000, 3),
+        dtype=np.uint16,
+        chunk_key_encoding={"name": "v2"},
+        # chunk_key_encoding=FanoutChunkKeyEncoding(separator="/"),
+        # chunk_key_encoding="fanout",
+        overwrite=True,
+        # compressors=zarr.codecs.BloscCodec(),
+        # compressors=GZip(),
+    )
+    rng = np.random.default_rng()
+    data = rng.integers(0, 65_535, size=(10_000_000, 3), dtype=np.uint16)
+    z[:] = data
+    print(z.info)
+
+
+if __name__ == "__main__":
+    main()

src/zarr/core/chunk_key_encodings.py

@@ -108,5 +108,5 @@ def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding:
     return chunk_key_encoding
 
 
-register_chunk_key_encoding(DefaultChunkKeyEncoding, qualname="default")
-register_chunk_key_encoding(V2ChunkKeyEncoding, qualname="v2")
+register_chunk_key_encoding("default", DefaultChunkKeyEncoding)
+register_chunk_key_encoding("v2", V2ChunkKeyEncoding)

src/zarr/fanout.py

@@ -0,0 +1,13 @@
+from zarr.core.chunk_key_encodings import ChunkKeyEncoding
+
+
+class FanoutChunkKeyEncoding(ChunkKeyEncoding):
+    name = "fanout"
+
+    def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]:
+        if chunk_key == "c":
+            return ()
+        return tuple(map(int, chunk_key[1:].split(self.separator)))
+
+    def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
+        return self.separator.join(map(str, ("c",) + chunk_coords))

src/zarr/registry.py

@@ -25,7 +25,6 @@
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
     from zarr.core.common import JSON
 
-# CHANGE: Consider adding here
 __all__ = [
     "Registry",
     "get_buffer_class",
@@ -48,9 +47,9 @@ def __init__(self) -> None:
         super().__init__()
         self.lazy_load_list: list[EntryPoint] = []
 
-    def lazy_load(self) -> None:
+    def lazy_load(self, use_entrypoint_name: bool = False) -> None:
         for e in self.lazy_load_list:
-            self.register(e.load())
+            self.register(e.load(), qualname=e.name if use_entrypoint_name else None)
 
         self.lazy_load_list.clear()
 
@@ -158,8 +157,8 @@ def register_buffer(cls: type[Buffer], qualname: str | None = None) -> None:
     __buffer_registry.register(cls, qualname)
 
 
-def register_chunk_key_encoding(cls: type, qualname: str | None = None) -> None:
-    __chunk_key_encoding_registry.register(cls, qualname)
+def register_chunk_key_encoding(key: str, cls: type) -> None:
+    __chunk_key_encoding_registry.register(cls, key)
 
 
 def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]:
@@ -300,12 +299,11 @@ def get_ndbuffer_class(reload_config: bool = False) -> type[NDBuffer]:
 
 
 def get_chunk_key_encoding_class(key: str) -> type[ChunkKeyEncoding]:
-    __chunk_key_encoding_registry.lazy_load()
+    __chunk_key_encoding_registry.lazy_load(use_entrypoint_name=True)
     if key not in __chunk_key_encoding_registry:
         raise KeyError(
             f"Chunk key encoding '{key}' not found in registered chunk key encodings: {list(__chunk_key_encoding_registry)}."
         )
-
     return __chunk_key_encoding_registry[key]