zarr_n5

 26class N5WrapperStore(WrapperStore[TStore], Generic[TStore]):
 27    """A read-only store for opening N5 hierarchies.
 28
 29    Requests for Zarr metadata documents are redirected to N5 attributes,
 30    and Zarr metadata calculated on the fly.
 31
 32    Note that N5 attributes can be omitted in groups.
 33    You may want to wrap this in an `ImplicitGroupWrapperStore` to replicate that behaviour.
 34
 35    Only compatible with DEFAULT-mode N5 arrays.
 36    """
 37
 38    _store: TStore
 39
 40    def intercept_metadata(self, key: str) -> None | str:
 41        """If the given key is for Zarr v3 metadata, return the key for N5 metadata in the equivalent node.
 42
 43        Otherwise, return None.
 44        """
 45        if "/" in key:
 46            pref, fname = key.rsplit("/", 1)
 47        else:
 48            pref = None
 49            fname = key
 50
 51        if fname != ZARR_V3_METADATA_KEY:
 52            return None
 53
 54        if pref is None:
 55            k2 = N5_METADATA_KEY
 56        else:
 57            k2 = f"{pref}/{N5_METADATA_KEY}"
 58
 59        return k2
 60
 61    async def get(
 62        self,
 63        key: str,
 64        prototype: BufferPrototype,
 65        byte_range: ByteRequest | None = None,
 66    ) -> Buffer | None:
 67        k2 = self.intercept_metadata(key)
 68        if k2 is None:
 69            return await self._store.get(key, prototype, byte_range)
 70
 71        res = await self._store.get(k2, prototype)
 72
 73        if res is None:
 74            return None
 75
 76        b = res.to_bytes()
 77        logger.debug("Got N5 metadata from %s: %s", k2, b)
 78        d = json.loads(b)
 79        n5_meta = N5GroupMetadata.from_jso(d)
 80        try:
 81            n5_meta = N5ArrayMetadata.from_group(n5_meta)
 82            out_d = n5_meta.to_zarr(N5Mode.DEFAULT)
 83        except KeyError:
 84            out_d = n5_meta.to_zarr()
 85
 86        b2 = json.dumps(out_d.to_dict()).encode()
 87        logger.debug("Inferred Zarr v3 metadata at %s: %s", key, b2)
 88
 89        b2 = slice_buf(b2, byte_range)
 90
 91        return prototype.buffer.from_bytes(b2)
 92
 93    async def get_partial_values(
 94        self,
 95        prototype: BufferPrototype,
 96        key_ranges: Iterable[tuple[str, ByteRequest | None]],
 97    ) -> list[Buffer | None]:
 98
 99        # Split the key ranges into metadata requests and other (chunk) requests.
100        # We always need to read the whole N5 metadata file
101        # to convert it into Zarr v3 metadata before slicing it,
102        # so this prevents reading it multiple times.
103        meta_reqs: defaultdict[str, list[tuple[int, ByteRequest | None]]] = defaultdict(
104            list
105        )
106        other_reqs: list[tuple[int, tuple[str, ByteRequest | None]]] = []
107        count = 0
108        for idx, (key, byte_range) in enumerate(key_ranges):
109            if is_zarr3_metadata(key):
110                meta_reqs[key].append((idx, byte_range))
111            else:
112                other_reqs.append((idx, (key, byte_range)))
113            count += 1
114
115        other_reqs_fut = self._store.get_partial_values(
116            prototype, (tup[1] for tup in other_reqs)
117        )
118        meta_req_list = list(meta_reqs.items())
119        meta_reqs_fut = asyncio.gather(
120            *(self.get(k, prototype) for k, _ in meta_req_list)
121        )
122        # Gather all requests to run concurrently
123        other_res, meta_res = await asyncio.gather(other_reqs_fut, meta_reqs_fut)
124        out: list[None | Buffer] = [None for _ in range(count)]
125
126        # Slice and insert the metadata responses into the pre-allocated output list
127        for res, (_, meta_req) in zip(meta_res, meta_req_list):
128            if res is None:
129                continue
130            blike = res.as_buffer_like()
131            for idx, byte_range in meta_req:
132                out[idx] = Buffer.from_bytes(slice_buf(blike, byte_range))
133
134        # Insert the non-metadata responses into the output;
135        # these are already sliced by the underlying store.
136        for res, (idx, _) in zip(other_res, other_reqs):
137            out[idx] = res
138
139        return out
140
141    async def exists(self, key: str) -> bool:
142        k2 = self.intercept_metadata(key)
143        return await self._store.exists(k2 or key)
144
145    @property
146    def supports_writes(self) -> bool:
147        return False
148
149    @property
150    def supports_deletes(self) -> bool:
151        return False
152
153    async def delete(self, key: str) -> None:
154        raise NotImplementedError
155
156    @property
157    def supports_listing(self) -> bool:
158        return self._store.supports_listing
159
160    def list(self) -> AsyncIterator[str]:
161        return self._store.list()
162
163    def list_prefix(self, prefix: str) -> AsyncIterator[str]:
164        return self._store.list_prefix(prefix)
165
166    def list_dir(self, prefix: str) -> AsyncIterator[str]:
167        return self._store.list_dir(prefix)
168
169    def __str__(self) -> str:
170        return f"{type(self).__name__}({self._store})"

33class ImplicitGroupWrapperStore(WrapperStore[TStore], Generic[TStore]):
34    """A store which supplies empty group metadata documents if they do not exist.
35
36    Used to replicate N5's behaviour where any directory (or prefix) is a valid group,
37    even when no metadata document exists.
38    Wrap over an `N5WrapperStore`.
39
40    Inferred group metadata's attributes will contain the key/value `"_implicit": true`.
41    """
42
43    _store: TStore
44
45    async def get(
46        self,
47        key: str,
48        prototype: BufferPrototype,
49        byte_range: ByteRequest | None = None,
50    ) -> Buffer | None:
51        res = await self._store.get(key, prototype, byte_range)
52        if res is not None or not is_zarr3_metadata(key):
53            return res
54
55        b = slice_buf(IMPLICIT_GROUP_BYTES, byte_range)
56        return prototype.buffer.from_bytes(b)
57
58    async def get_partial_values(
59        self,
60        prototype: BufferPrototype,
61        key_ranges: Iterable[tuple[str, ByteRequest | None]],
62    ) -> list[Buffer | None]:
63        key_ranges = list(key_ranges)
64        reses = await super().get_partial_values(prototype, key_ranges)
65        out = []
66        for (key, byte_range), res in zip(key_ranges, reses):
67            if res is None and is_zarr3_metadata(key):
68                res = prototype.buffer.from_bytes(
69                    slice_buf(IMPLICIT_GROUP_BYTES, byte_range)
70                )
71            out.append(res)
72
73        return out
74
75    async def exists(self, key: str) -> bool:
76        if is_zarr3_metadata(key):
77            return True
78        return await super().exists(key)
79
80    def __str__(self) -> str:
81        return f"{type(self).__name__}({self._store})"

 70@dataclass(frozen=True)
 71class N5DefaultCodec(ArrayBytesCodec):
 72    """Zarr codec for default-mode N5 data.
 73
 74    Only full-chunk reads are supported.
 75    Use the `N5DefaultCodec.from_compressor` constructor if initialising manually.
 76
 77    - Reads and validates the N5 block header
 78    - Applies the wrapped codecs to the N5 block body
 79    - Truncates or pads the resulting array to match the requested chunk
 80
 81    Should be the only codec present.
 82    """
 83
 84    codecs: CodecTuple
 85    """Codecs to be applied to the N5 block body."""
 86
 87    def __init__(self, *, codecs: Iterable[Codec | dict[str, JSON]]) -> None:
 88        cs = parse_codecs(fix_bytes(codecs))
 89        if not 2 <= len(cs) <= 3:
 90            raise ValueError(f"expected 2-3 codecs, got {len(cs)}")
 91        check_valid_transpose(cs[0])
 92        check_valid_bytes(cs[1])
 93        if len(cs) > 2:
 94            check_valid_compressor(cs[2])
 95
 96        object.__setattr__(self, "codecs", cs)
 97
 98    @property
 99    def codec_pipeline(self) -> CodecPipeline:
100        """Get the `CodecPipeline` comprising the wrapped codecs."""
101        return get_pipeline_class().from_codecs(self.codecs)
102
103    @classmethod
104    def from_compressor(cls, ndim: int, compressor: BytesBytesCodec | None = None):
105        """Construct the codec from minimal information."""
106        transpose = cls.make_transpose(ndim)
107        endian = cls.make_bytes()
108        codecs: CodecTuple
109        if compressor is None:
110            codecs = (transpose, endian)
111        else:
112            codecs = (transpose, endian, compressor)
113        return cls(codecs=codecs)
114
115    @classmethod
116    def make_transpose(cls, ndim: int) -> TransposeCodec:
117        """Generate the `TransposeCodec` needed for this data.
118
119        N5 data is always fully transposed.
120        """
121        order = list(range(ndim))
122        return TransposeCodec(order=tuple(reversed(order)))
123
124    @classmethod
125    def make_bytes(cls) -> BytesCodec:
126        """
127        Generate the `BytesCodec` needed for this data.
128
129        N5 data is always big-endian.
130        """
131        return BytesCodec(endian=N5_ENDIAN)
132
133    def compute_encoded_size(
134        self, input_byte_length: int, chunk_spec: ArraySpec
135    ) -> int:
136        header_length = N5BlockHeader.calc_size(chunk_spec.ndim, False)
137
138        for c in self.codecs:
139            input_byte_length = c.compute_encoded_size(input_byte_length, chunk_spec)
140            chunk_spec = c.resolve_metadata(chunk_spec)
141
142        return input_byte_length + header_length
143
144    def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
145        for c in self.codecs:
146            chunk_spec = c.resolve_metadata(chunk_spec)
147        return chunk_spec
148
149    def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
150        transpose = self.make_transpose(array_spec.ndim).evolve_from_array_spec(
151            array_spec
152        )
153        endian = self.make_bytes().evolve_from_array_spec(array_spec)
154
155        codecs: CodecTuple
156        match len(self.codecs):
157            case 2:
158                codecs = (transpose, endian)
159            case 3:
160                compressor: BytesBytesCodec = self.codecs[2].evolve_from_array_spec(  # type:ignore
161                    array_spec
162                )
163                codecs = (transpose, endian, compressor)
164            case _:
165                raise ValueError("unsupported number of codecs")
166
167        return type(self)(codecs=codecs)
168
169    def validate(
170        self,
171        *,
172        shape: tuple[int, ...],
173        dtype: ZDType[TBaseDType, TBaseScalar],
174        chunk_grid: ChunkGridMetadata,
175    ) -> None:
176        expected_ndim = len(self.codecs[0].order)
177        if len(shape) != expected_ndim:
178            raise ValueError(f"array is {len(shape)}D, codec is {expected_ndim}D")
179        if dtype._zarr_v3_name not in COMPATIBLE_DATA_TYPES:
180            raise ValueError(f"N5 does not support data type {dtype._zarr_v3_name}")
181
182        return super().validate(shape=shape, dtype=dtype, chunk_grid=chunk_grid)
183
184    async def _decode_single(
185        self, chunk_data: Buffer, chunk_spec: ArraySpec
186    ) -> NDBuffer:
187        b = chunk_data.as_buffer_like()
188        header = N5BlockHeader.from_bytes(b)
189        offset = header.size()
190
191        body_buf = chunk_data[offset:]
192        body_nd = chunk_spec.prototype.nd_buffer.empty(
193            header.shape, chunk_spec.dtype.to_native_dtype(), chunk_spec.order
194        )
195        if header.shape == chunk_spec.shape:
196            body_spec = chunk_spec
197            all_eq = True
198        else:
199            body_spec = ArraySpec(
200                header.shape,
201                chunk_spec.dtype,
202                chunk_spec.fill_value,
203                chunk_spec.config,
204                chunk_spec.prototype,
205            )
206            all_eq = False
207        maybe_body_nd, *_ = await self.codec_pipeline.decode([(body_buf, body_spec)])
208        # TODO: use codec_pipeline.read() instead; this should avoid the copy for truncated-block cases
209        if maybe_body_nd is None:
210            raise RuntimeError("unexpected nullish buffer")
211        else:
212            body_nd = maybe_body_nd
213
214        if all_eq:
215            # don't need to truncate or pad
216            return body_nd
217
218        # whether we can get the chunk we want by trimming down the N5 block body
219        can_trim = True
220
221        min_shape = []
222        slice_lst = []
223        for hs, cs in zip(header.shape, chunk_spec.shape):
224            if cs > hs:
225                # requested chunk is larger than the N5 block in some dimension
226                can_trim = False
227            min_len = min(hs, cs)
228            min_shape.append(min_len)
229            slice_lst.append(slice(0, min_len))
230
231        slicing = tuple(slice_lst)
232
233        if can_trim:
234            return body_nd[slicing]
235
236        out = chunk_spec.prototype.nd_buffer.create(
237            shape=chunk_spec.shape,
238            dtype=chunk_spec.dtype.to_native_dtype(),
239            order=chunk_spec.order,
240            fill_value=chunk_spec.fill_value,
241        )
242        out[slicing] = body_nd[slicing]
243        return out
244
245    # async def _encode_single(
246    #     self, chunk_data: NDBuffer, chunk_spec: ArraySpec
247    # ) -> Buffer | None:
248    #     header = N5BlockHeader(N5Mode.DEFAULT, chunk_spec.shape)
249    #     for c in self.codecs:
250    #         chunk_data = c._encode_single(chunk_data, chunk_spec)  # type:ignore
251    #         chunk_spec = c.resolve_metadata(chunk_spec)
252
253    #     buf: Buffer = chunk_data  # type: ignore
254
255    #     bio = BytesIO()
256    #     bio.write(header.to_bytes())
257    #     # TODO: avoid this copy?
258    #     bio.write(buf.as_buffer_like())
259    #     return Buffer.from_bytes(bio.getbuffer())
260
261    @classmethod
262    def from_dict(
263        cls,
264        data: dict[str, JSON],
265    ) -> Self:
266        _, configuration_parsed = parse_named_configuration(
267            data, N5_DEFAULT_NAME, require_configuration=True
268        )
269
270        return cls(**configuration_parsed)  # type: ignore[arg-type]
271
272    def to_dict(
273        self,
274    ) -> dict[str, JSON]:
275        return {
276            "name": N5_DEFAULT_NAME,
277            "configuration": {"codecs": [c.to_dict() for c in self.codecs]},
278        }