from __future__ import annotations

from collections.abc import Iterable, Mapping, MutableMapping, Sequence

from dataclasses import dataclass, replace

from functools import lru_cache

from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple

import numpy as np

import numpy.typing as npt

from zarr.abc.codec import (

ArrayBytesCodec,

ArrayBytesCodecPartialDecodeMixin,

ArrayBytesCodecPartialEncodeMixin,

Codec,

CodecPipeline,

_codec_supports_sync,

)

from zarr.abc.store import (

ByteGetter,

ByteRequest,

ByteSetter,

RangeByteRequest,

Store,

SuffixByteRequest,

_store_supports_sync_io,

)

from zarr.codecs._deprecated_enum import _coerce_enum_input, _DeprecatedStrEnumMeta

from zarr.codecs.bytes import BytesCodec

from zarr.codecs.crc32c_ import Crc32cCodec

from zarr.core.array_spec import ArrayConfig, ArraySpec

from zarr.core.buffer import (

Buffer,

BufferPrototype,

NDBuffer,

default_buffer_prototype,

numpy_buffer_prototype,

)

from zarr.core.chunk_grids import ChunkGrid

from zarr.core.chunk_utils import (

ChunkTransform,

decode_and_scatter_chunk,

encode_or_elide_chunk,

evolve_codecs,

merge_and_encode_chunk,

)

from zarr.core.common import (

ShapeLike,

parse_chunk_shape,

parse_named_configuration,

parse_shapelike,

product,

)

from zarr.core.config import config as zarr_config

from zarr.core.dtype.common import HasEndianness

from zarr.core.dtype.npy.int import UInt64

from zarr.core.dtype.npy.structured import Struct

from zarr.core.indexing import (

BasicIndexer,

ChunkProjection,

CoordinateIndexer,

SelectorTuple,

SliceDimIndexer,

_lexicographic_order,

colexicographic_order_coords,

get_indexer,

lexicographic_order_coords,

morton_order_coords,

)

from zarr.core.metadata.v3 import (

ChunkGridMetadata,

RectilinearChunkGridMetadata,

RegularChunkGridMetadata,

parse_codecs,

)

from zarr.registry import get_ndbuffer_class, get_pipeline_class

from zarr.storage._common import StorePath

from zarr.storage._utils import _normalize_byte_range_index

if TYPE_CHECKING:

from collections.abc import Iterator

from typing import Self

from zarr.core.common import JSON

from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType

MAX_UINT_64 = 2**64 - 1

ShardMapping = Mapping[tuple[int, ...], Buffer | None]

ShardMutableMapping = MutableMapping[tuple[int, ...], Buffer | None]

IndexLocation = Literal["start", "end"]

"""Position of the shard index within the encoded shard."""

INDEX_LOCATION: Final = ("start", "end")

class ShardingCodecIndexLocation(metaclass=_DeprecatedStrEnumMeta):

"""

Deprecated. Pass a literal string (`"start"` or `"end"`) directly to

`ShardingCodec` instead.

"""

_members: ClassVar[dict[str, str]] = {"start": "start", "end": "end"}

SubchunkWriteOrder = Literal["morton", "unordered", "lexicographic", "colexicographic"]

SUBCHUNK_WRITE_ORDER: Final[tuple[str, str, str, str]] = (

"morton",

"unordered",

"lexicographic",

"colexicographic",

)

def _is_identity_full_read(indexer: Any, shard_shape: tuple[int, ...]) -> bool:

"""True when `indexer` selects every element of a `shard_shape` array in

natural order: one whole-dimension, step-1 `SliceDimIndexer` per dimension.

Structural on purpose, not `isinstance(indexer, BasicIndexer)`: a full

`arr[:]` read reaches the shard as an `OrthogonalIndexer`, so a type gate

would silently disable the bulk fast path for the most common case. Any

gather (integer-array / boolean / coordinate selection), subset, strided,

or integer-scalar selection fails the per-dimension check — output shape

alone is not enough, because a reordering or duplicating selection can have

the same shape as the shard while requiring `chunk_selection` /

`out_selection` to be honored.

"""

dim_indexers = getattr(indexer, "dim_indexers", None)

if dim_indexers is None or len(dim_indexers) != len(shard_shape):

return False

return all(

isinstance(dim_indexer, SliceDimIndexer)

and dim_indexer.dim_len == dim_len

and dim_indexer.start == 0

and dim_indexer.stop == dim_len

and dim_indexer.step == 1

for dim_indexer, dim_len in zip(dim_indexers, shard_shape, strict=True)

)

def _parse_index_location(data: object) -> IndexLocation:

if isinstance(data, str) and data in INDEX_LOCATION:

return data # type: ignore[return-value]

raise ValueError(f"index_location must be one of {list(INDEX_LOCATION)!r}. Got {data!r}.")

@dataclass(frozen=True)

class _ShardingByteGetter(ByteGetter):

"""In-memory byte getter for one inner chunk of a shard.

Implements `SyncByteGetter` (dict access needs no event loop), so the

synchronous codec pipeline takes its sync fast path on inner chunks

instead of scheduling one coroutine per chunk; the async `get` simply

delegates to `get_sync`.

"""

shard_dict: ShardMapping

chunk_coords: tuple[int, ...]

def get_sync(

self, prototype: BufferPrototype | None = None, byte_range: ByteRequest | None = None

) -> Buffer | None:

if prototype is not None and prototype != default_buffer_prototype():

raise ValueError(

"Non-default buffer prototypes are not supported within shards. "

f"Got {prototype}, expected {default_buffer_prototype()}."

)

value = self.shard_dict.get(self.chunk_coords)

if value is None:

return None

if byte_range is None:

return value

start, stop = _normalize_byte_range_index(value, byte_range)

return value[start:stop]

async def get(

self, prototype: BufferPrototype, byte_range: ByteRequest | None = None

) -> Buffer | None:

return self.get_sync(prototype, byte_range)

@dataclass(frozen=True)

class _ShardingByteSetter(_ShardingByteGetter, ByteSetter):

"""In-memory byte setter for one inner chunk of a shard.

Implements `SyncByteSetter`; the async methods delegate to the sync ones.

"""

shard_dict: ShardMutableMapping

def set_sync(self, value: Buffer) -> None:

self.shard_dict[self.chunk_coords] = value

def delete_sync(self) -> None:

del self.shard_dict[self.chunk_coords]

async def set(self, value: Buffer, byte_range: ByteRequest | None = None) -> None:

if byte_range is not None:

raise NotImplementedError("byte_range is not supported within shards.")

self.set_sync(value)

async def delete(self) -> None:

self.delete_sync()

async def set_if_not_exists(self, default: Buffer) -> None:

self.shard_dict.setdefault(self.chunk_coords, default)

class _ShardIndex(NamedTuple):

# the chunk grid shape of a single shard

chunks_per_shard: tuple[int, ...]

# dtype uint64, shape (chunks_per_shard_0, chunks_per_shard_1, ..., 2)

offsets_and_lengths: npt.NDArray[np.uint64]

def _localize_chunk(self, chunk_coords: tuple[int, ...]) -> tuple[int, ...]:

return tuple(

chunk_i % shard_i

for chunk_i, shard_i in zip(chunk_coords, self.chunks_per_shard, strict=False)

)

def is_all_empty(self) -> bool:

return bool(np.array_equiv(self.offsets_and_lengths, MAX_UINT_64))

def get_full_chunk_map(self) -> npt.NDArray[np.bool_]:

return np.not_equal(self.offsets_and_lengths[..., 0], MAX_UINT_64)

def is_dense(self, chunk_byte_length: int, *, data_section_start: int) -> bool:

"""True when the chunk payloads exactly tile the shard's data section.

Every chunk must be present with length `chunk_byte_length`, and the

sorted offsets must be exactly `data_section_start + i * chunk_byte_length`

for `i` in `0..n_chunks-1`: no gaps, no overlaps, and nothing outside the

data section (a corrupt index could otherwise point chunks into the

index region or out of the blob). Used to gate the vectorized

whole-shard decode: a dense fixed-size shard is a regular grid of

equal-length payloads, so it can be reshaped/scattered in bulk rather

than decoded chunk-by-chunk.

"""

offsets = self.offsets_and_lengths[..., 0].reshape(-1)

lengths = self.offsets_and_lengths[..., 1].reshape(-1)

# all present

if bool(np.any(offsets == MAX_UINT_64)):

return False

# all the same fixed length

if not bool(np.all(lengths == chunk_byte_length)):

return False

expected = np.uint64(data_section_start) + np.arange(

offsets.size, dtype=np.uint64

) * np.uint64(chunk_byte_length)

return bool(np.array_equal(np.sort(offsets), expected))

def get_chunk_slice(self, chunk_coords: tuple[int, ...]) -> tuple[int, int] | None:

localized_chunk = self._localize_chunk(chunk_coords)

chunk_start, chunk_len = self.offsets_and_lengths[localized_chunk]

if (chunk_start, chunk_len) == (MAX_UINT_64, MAX_UINT_64):

return None

else:

return (int(chunk_start), int(chunk_start + chunk_len))

def get_chunk_slices_vectorized(

self, chunk_coords_array: npt.NDArray[np.integer[Any]]

) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint64], npt.NDArray[np.bool_]]:

"""Get chunk slices for multiple coordinates at once.

Parameters

----------

chunk_coords_array : ndarray of shape (n_chunks, n_dims)

Array of chunk coordinates to look up.

Returns

-------

starts : ndarray of shape (n_chunks,)

Start byte positions for each chunk.

ends : ndarray of shape (n_chunks,)

End byte positions for each chunk.

valid : ndarray of shape (n_chunks,)

Boolean mask indicating which chunks are non-empty.

"""

# Handle 0-dimensional arrays (n_dims == 0): the shard holds a single

# chunk, so every coordinate maps to the same flat entry.

if chunk_coords_array.shape[1] == 0:

offsets_and_lengths = self.offsets_and_lengths.reshape(-1, 2)

offsets_and_lengths = np.broadcast_to(

offsets_and_lengths, (chunk_coords_array.shape[0], 2)

)

else:

# Localize coordinates via modulo (vectorized)

shard_shape = np.array(self.chunks_per_shard, dtype=np.uint64)

localized = chunk_coords_array.astype(np.uint64) % shard_shape

# Build index tuple for advanced indexing

index_tuple = tuple(localized[:, i] for i in range(localized.shape[1]))

# Fetch all offsets and lengths at once

offsets_and_lengths = self.offsets_and_lengths[index_tuple]

starts = offsets_and_lengths[:, 0]

lengths = offsets_and_lengths[:, 1]

# Check for valid (non-empty) chunks

valid = starts != MAX_UINT_64

# Compute end positions

ends = starts + lengths

return starts, ends, valid

def set_chunk_slice(self, chunk_coords: tuple[int, ...], chunk_slice: slice | None) -> None:

localized_chunk = self._localize_chunk(chunk_coords)

if chunk_slice is None:

self.offsets_and_lengths[localized_chunk] = (MAX_UINT_64, MAX_UINT_64)

else:

self.offsets_and_lengths[localized_chunk] = (

chunk_slice.start,

chunk_slice.stop - chunk_slice.start,

)

@classmethod

def create_empty(cls, chunks_per_shard: tuple[int, ...]) -> _ShardIndex:

offsets_and_lengths = np.zeros(chunks_per_shard + (2,), dtype="<u8", order="C")

offsets_and_lengths.fill(MAX_UINT_64)

return cls(chunks_per_shard, offsets_and_lengths)

class _ShardReader(ShardMapping):

buf: Buffer

index: _ShardIndex

@classmethod

async def from_bytes(

cls, buf: Buffer, codec: ShardingCodec, chunks_per_shard: tuple[int, ...]

) -> _ShardReader:

shard_index_size = codec._shard_index_size(chunks_per_shard)

obj = cls()

obj.buf = buf

if codec.index_location == "start":

shard_index_bytes = obj.buf[:shard_index_size]

else:

shard_index_bytes = obj.buf[-shard_index_size:]

obj.index = await codec._decode_shard_index(shard_index_bytes, chunks_per_shard)

return obj

@classmethod

def create_empty(

cls, chunks_per_shard: tuple[int, ...], buffer_prototype: BufferPrototype | None = None

) -> _ShardReader:

if buffer_prototype is None:

buffer_prototype = default_buffer_prototype()

index = _ShardIndex.create_empty(chunks_per_shard)

obj = cls()

obj.buf = buffer_prototype.buffer.create_zero_length()

obj.index = index

return obj

def __getitem__(self, chunk_coords: tuple[int, ...]) -> Buffer:

chunk_byte_slice = self.index.get_chunk_slice(chunk_coords)

if chunk_byte_slice is not None:

return self.buf[chunk_byte_slice[0] : chunk_byte_slice[1]]

raise KeyError

def __len__(self) -> int:

return int(self.index.offsets_and_lengths.size / 2)

def __iter__(self) -> Iterator[tuple[int, ...]]:

return iter(lexicographic_order_coords(self.index.chunks_per_shard))

def to_dict_vectorized(self) -> dict[tuple[int, ...], Buffer | None]:

"""Build a dict of chunk coordinates to buffers using vectorized lookup.

The full per-shard chunk coordinate grid (both the array used for the

vectorized index lookup and the plain tuples used as dict keys) is

cached on `chunks_per_shard`, so neither is rebuilt on every call. For a

shard with tens of thousands of chunks this avoids reconstructing that

many tuples on every partial write.

Returns

-------

dict mapping chunk coordinate tuples to Buffer or None

"""

chunks_per_shard = self.index.chunks_per_shard

# The same chunk-grid coordinates are needed in two forms, and neither can

# stand in for the other:

# - `chunk_coords_array`: an (n_chunks, n_dims) numpy array, fed to the

# vectorized index lookup, which does modulo + advanced indexing on it.

# A list of tuples can't be used for that without first being arrayified.

# - `chunk_coords_keys`: the same coordinates as hashable Python tuples,

# used as the result dict's keys. numpy array rows are unhashable

# (mutable), so they can't key a dict.

# Both are cached per shape (see indexing.py), so neither is rebuilt here;

# row i of the array and key i refer to the same chunk.

chunk_coords_array = _lexicographic_order(chunks_per_shard)

chunk_coords_keys = lexicographic_order_coords(chunks_per_shard)

starts, ends, valid = self.index.get_chunk_slices_vectorized(chunk_coords_array)

result: dict[tuple[int, ...], Buffer | None] = {}

for i, coords in enumerate(chunk_coords_keys):

if valid[i]:

result[coords] = self.buf[int(starts[i]) : int(ends[i])]

else:

result[coords] = None

return result

def _drops_only_unit_axes(shape: tuple[int, ...], full: tuple[int, ...]) -> bool:

"""Return whether ``shape`` is ``full`` with zero or more length-1 axes removed.

``(2, 2)`` is ``(2, 1, 2)`` minus its unit axis, and ``(1, 2)`` is

``(1, 2, 1)`` minus its last; ``(2, 2)`` is not ``(4,)``, and

``(3, 2, 1)`` is not ``(3, 2)`` because it adds an axis.

"""

remaining = iter(full)

for size in shape:

for full_size in remaining:

if full_size == size:

break

if full_size != 1:

return False

else:

return False

return all(full_size == 1 for full_size in remaining)

def _check_index_codecs_fixed_size(index_codecs: tuple[Codec, ...]) -> None:

"""Reject index codecs whose encoded size is not fixed.

The spec forbids variable-size codecs (such as compressors) in `index_codecs`:

the shard index is located by its byte size, which must be known before the

index is read.

"""

variable_size = [codec for codec in index_codecs if not codec.is_fixed_size]

if variable_size:

raise ValueError(

"Sharding `index_codecs` must produce a fixed-size encoding, but these codecs "

f"do not: {variable_size}. Compression codecs cannot be used for the shard index."

)

def _parse_inner_chunk_shape(data: ShapeLike) -> tuple[int, ...]:

"""Parse the inner chunk shape of a sharding codec: a shape (see `parse_shapelike`)

whose every size is a chunk edge length, so at least 1 (see `parse_chunk_shape`)."""

return parse_chunk_shape(parse_shapelike(data))

@dataclass(frozen=True)

class ShardingCodec(

ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin

):

"""Sharding codec.

`subchunk_write_order` controls the physical order of subchunks within a shard. It is a

write-time setting only: it is not stored in array metadata, so reopening a sharded array

does not recover it (the setting reverts to the `morton` default per codec instance).

"""

is_fixed_size = False

chunk_shape: tuple[int, ...]

codecs: tuple[Codec, ...]

index_codecs: tuple[Codec, ...]

index_location: IndexLocation = "end"

subchunk_write_order: SubchunkWriteOrder = "morton"

def __init__(

self,

*,

chunk_shape: ShapeLike,

codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(endian="little"),),

index_codecs: Iterable[Codec | dict[str, JSON]] = (

BytesCodec(endian="little"),

Crc32cCodec(),

),

index_location: ShardingCodecIndexLocation | IndexLocation = "end",

subchunk_write_order: SubchunkWriteOrder = "morton",

) -> None:

chunk_shape_parsed = _parse_inner_chunk_shape(chunk_shape)

codecs_parsed = parse_codecs(codecs)

index_codecs_parsed = parse_codecs(index_codecs)

_check_index_codecs_fixed_size(index_codecs_parsed)

index_location_coerced = _coerce_enum_input(

index_location, "index_location", "ShardingCodec"

)

index_location_parsed = _parse_index_location(index_location_coerced)

if subchunk_write_order not in SUBCHUNK_WRITE_ORDER:

raise ValueError(

f"Unrecognized subchunk write order: {subchunk_write_order}. Only {SUBCHUNK_WRITE_ORDER} are allowed."

)

object.__setattr__(self, "chunk_shape", chunk_shape_parsed)

object.__setattr__(self, "codecs", codecs_parsed)

object.__setattr__(self, "index_codecs", index_codecs_parsed)

object.__setattr__(self, "index_location", index_location_parsed)

object.__setattr__(self, "subchunk_write_order", subchunk_write_order)

# Use instance-local lru_cache to avoid memory leaks

object.__setattr__(self, "_get_chunk_spec", lru_cache()(self._get_chunk_spec))

object.__setattr__(self, "_get_index_chunk_spec", lru_cache()(self._get_index_chunk_spec))

object.__setattr__(self, "_get_chunks_per_shard", lru_cache()(self._get_chunks_per_shard))

object.__setattr__(self, "_shard_index_size", lru_cache()(self._shard_index_size))

object.__setattr__(

self, "_get_inner_chunk_transform", lru_cache()(self._get_inner_chunk_transform)

)

object.__setattr__(

self, "_get_index_chunk_transform", lru_cache()(self._get_index_chunk_transform)

)

# todo: typedict return type

def __getstate__(self) -> dict[str, Any]:

# `subchunk_write_order` is not part of codec metadata (`to_dict`), so carry it

# explicitly to survive a pickle round-trip (otherwise it reverts to `morton`).

return {"subchunk_write_order": self.subchunk_write_order, **self.to_dict()}

def __setstate__(self, state: dict[str, Any]) -> None:

config = state["configuration"]

object.__setattr__(self, "chunk_shape", _parse_inner_chunk_shape(config["chunk_shape"]))

object.__setattr__(self, "codecs", parse_codecs(config["codecs"]))

object.__setattr__(self, "index_codecs", parse_codecs(config["index_codecs"]))

object.__setattr__(self, "index_location", _parse_index_location(config["index_location"]))

object.__setattr__(self, "subchunk_write_order", state["subchunk_write_order"])

# Use instance-local lru_cache to avoid memory leaks

object.__setattr__(self, "_get_chunk_spec", lru_cache()(self._get_chunk_spec))

object.__setattr__(self, "_get_index_chunk_spec", lru_cache()(self._get_index_chunk_spec))

object.__setattr__(self, "_get_chunks_per_shard", lru_cache()(self._get_chunks_per_shard))

object.__setattr__(self, "_shard_index_size", lru_cache()(self._shard_index_size))

object.__setattr__(

self, "_get_inner_chunk_transform", lru_cache()(self._get_inner_chunk_transform)

)

object.__setattr__(

self, "_get_index_chunk_transform", lru_cache()(self._get_index_chunk_transform)

)

@classmethod

def from_dict(cls, data: dict[str, JSON]) -> Self:

_, configuration_parsed = parse_named_configuration(data, "sharding_indexed")

return cls(**configuration_parsed) # type: ignore[arg-type]

@property

def codec_pipeline(self) -> CodecPipeline:

return get_pipeline_class().from_codecs(self.codecs)

def _get_inner_pipeline(self, shard_spec: ArraySpec) -> CodecPipeline:

"""The nested pipeline for inner-chunk IO, evolved against the inner

chunk spec.

Evolving matters for two reasons: it threads the spec through the inner

codec chain (spec-changing codecs see the spec they will actually

operate on), and — for a synchronous pipeline — it builds the sync

transform, so inner-chunk IO over the (sync-capable) sharding byte

getters takes the pipeline's sync fast path instead of scheduling one

coroutine per inner chunk. The bare `codec_pipeline` property returns an

unevolved pipeline, which a synchronous pipeline can only run through

its async fallback.

Memoized per (pipeline class, batch size, shard_spec): evolving builds

a ChunkTransform, which is wasteful to redo on every shard operation.

The pipeline class and batch size participate in the key so the

`codec_pipeline.path` and `codec_pipeline.batch_size` configs are still

honored after the first use (`from_codecs` captures batch_size at

construction). A benign construction race between threads is possible

(last writer wins) — same as the other caches here.

"""

cache: dict[tuple[type[CodecPipeline], int, ArraySpec], CodecPipeline] | None = getattr(

self, "_inner_pipeline_cache", None

)

if cache is None:

cache = {}

object.__setattr__(self, "_inner_pipeline_cache", cache)

key = (get_pipeline_class(), zarr_config.get("codec_pipeline.batch_size"), shard_spec)

pipeline = cache.get(key)

if pipeline is None:

chunk_spec = self._get_chunk_spec(shard_spec)

pipeline = self.codec_pipeline.evolve_from_array_spec(chunk_spec)

cache[key] = pipeline

return pipeline

def to_dict(self) -> dict[str, JSON]:

return {

"name": "sharding_indexed",

"configuration": {

"chunk_shape": self.chunk_shape,

"codecs": tuple(s.to_dict() for s in self.codecs),

"index_codecs": tuple(s.to_dict() for s in self.index_codecs),

"index_location": self.index_location,

},

}

def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:

"""Thread the spec through the inner chain.

Each codec is evolved against the spec produced by the previous one.

Evolving every codec against the same unthreaded spec is the bug shape that

strips `BytesCodec.endian` behind a dtype-changing codec — and this

method runs on the real array-creation path, baking the damaged chain

into the evolved instance before the transform builders ever run.

Parameters

----------

array_spec

The base spec to be evolved as we thread it.

Returns

-------

This codec with the evolved code chain.

"""

from zarr.core.chunk_utils import evolve_codecs

shard_spec = self._get_chunk_spec(array_spec)

evolved_codecs = evolve_codecs(self.codecs, shard_spec)

if evolved_codecs != self.codecs:

return replace(self, codecs=evolved_codecs)

return self

def validate(

self,

*,

shape: tuple[int, ...],

dtype: ZDType[TBaseDType, TBaseScalar],

chunk_grid: ChunkGridMetadata,

) -> None:

if len(self.chunk_shape) != len(shape):

raise ValueError(

"The shard's `chunk_shape` and array's `shape` need to have the same number of dimensions."

)

if isinstance(chunk_grid, RegularChunkGridMetadata):

edges_per_dim: tuple[tuple[int, ...], ...] = tuple((s,) for s in chunk_grid.chunk_shape)

elif isinstance(chunk_grid, RectilinearChunkGridMetadata):

edges_per_dim = tuple(

(s,) if isinstance(s, int) else s for s in chunk_grid.chunk_shapes

)

else:

raise TypeError(

f"Sharding is only compatible with regular and rectilinear chunk grids, "

f"got {type(chunk_grid)}"

)

for i, (edges, inner) in enumerate(zip(edges_per_dim, self.chunk_shape, strict=False)):

for edge in set(edges):

if edge % inner != 0:

raise ValueError(

f"Chunk edge length {edge} in dimension {i} is not "

f"divisible by the shard's inner chunk size {inner}."

)

self._validate_inner_codecs(dtype)

def _validate_inner_codecs(self, dtype: ZDType[TBaseDType, TBaseScalar]) -> None:

"""Validate the codecs that encode each inner chunk, as array metadata validates

its codecs: each against the chunk the codecs before it leave it, a regular

grid of one inner chunk. A sharding codec among them is thereby checked against

the chunk it splits, which `validate` of the outer codec alone does not see.

"""

spec = ArraySpec(

shape=self.chunk_shape,

dtype=dtype,

fill_value=dtype.default_scalar(),

config=ArrayConfig.from_dict({}),

prototype=default_buffer_prototype(),

)

for codec in self.codecs:

codec.validate(

shape=spec.shape,

dtype=spec.dtype,

chunk_grid=RegularChunkGridMetadata(chunk_shape=spec.shape),

)

spec = codec.resolve_metadata(spec)

def _get_inner_chunk_transform(self, shard_spec: ArraySpec) -> Any:

"""The synchronous transform for the inner codec chain.

Memoized by the instance-local `lru_cache` wrapping installed in

`__init__`/`__setstate__` (the single cache mechanism for these

builders — do not add another layer inside the body).

Codecs are evolved with the spec THREADED forward (`evolve_codecs`):

each inner codec is evolved against the spec produced by the previous

one, not the original chunk spec. Evolving every codec against the same

unthreaded spec is the bug shape that stripped `BytesCodec.endian` at

the pipeline level (see `evolve_codecs`) — the inner chain must use the

same single source of truth.

"""

chunk_spec = self._get_chunk_spec(shard_spec)

return ChunkTransform(codecs=evolve_codecs(self.codecs, chunk_spec))

def _get_index_chunk_transform(self, chunks_per_shard: tuple[int, ...]) -> Any:

"""The synchronous transform for the index codec chain.

Memoized via instance-local `lru_cache`.

"""

index_spec = self._get_index_chunk_spec(chunks_per_shard)

return ChunkTransform(codecs=evolve_codecs(self.index_codecs, index_spec))

def _decode_shard_index_sync(

self, index_bytes: Buffer, chunks_per_shard: tuple[int, ...]

) -> _ShardIndex:

"""Decode shard index synchronously using ChunkTransform."""

self._check_shard_index_size(index_bytes, chunks_per_shard)

index_transform = self._get_index_chunk_transform(chunks_per_shard)

index_spec = self._get_index_chunk_spec(chunks_per_shard)

index_array = index_transform.decode_chunk(index_bytes, index_spec)

return _ShardIndex(chunks_per_shard, index_array.as_numpy_array())

def _encode_shard_index_sync(self, index: _ShardIndex) -> Buffer:

"""Encode shard index synchronously using ChunkTransform."""

index_transform = self._get_index_chunk_transform(index.chunks_per_shard)

index_spec = self._get_index_chunk_spec(index.chunks_per_shard)

index_nd = get_ndbuffer_class().from_numpy_array(index.offsets_and_lengths)

result: Buffer | None = index_transform.encode_chunk(index_nd, index_spec)

if result is None:

raise RuntimeError("Encoding the shard index produced no bytes.")

return result

def _shard_reader_from_bytes_sync(

self, buf: Buffer, chunks_per_shard: tuple[int, ...]

) -> _ShardReader:

"""Sync version of _ShardReader.from_bytes."""

shard_index_size = self._shard_index_size(chunks_per_shard)

if self.index_location == "start":

shard_index_bytes = buf[:shard_index_size]

else:

shard_index_bytes = buf[-shard_index_size:]

index = self._decode_shard_index_sync(shard_index_bytes, chunks_per_shard)

reader = _ShardReader()

reader.buf = buf

reader.index = index

return reader

def _decode_sync(

self,

shard_bytes: Buffer,

shard_spec: ArraySpec,

) -> NDBuffer:

"""Decode a full shard synchronously.

Sync counterpart to `_decode_single`. Same semantics (decode every

inner chunk and assemble the full shard array) but routes through

`ChunkTransform` instead of the async codec pipeline, so it can

run on the sync codec-pipeline fast path without an event loop.

For a partial read where the caller only needs a slice of the shard,

use `_decode_partial_sync` instead — it fetches only the byte

ranges that overlap the selection.

This method does not parallelize decompression, but should.

See TODO: make issue for handling subchunk parallelism

"""

shard_shape = shard_spec.shape

chunk_shape = self.chunk_shape

chunks_per_shard = self._get_chunks_per_shard(shard_spec)

chunk_spec = self._get_chunk_spec(shard_spec)

inner_transform = self._get_inner_chunk_transform(shard_spec)

indexer = BasicIndexer(

tuple(slice(0, s) for s in shard_shape),

shape=shard_shape,

chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape),

)

out = chunk_spec.prototype.nd_buffer.empty(

shape=shard_shape,

dtype=shard_spec.dtype.to_native_dtype(),

order=shard_spec.order,

)

shard_dict = self._shard_reader_from_bytes_sync(shard_bytes, chunks_per_shard)

if shard_dict.index.is_all_empty():

out.fill(shard_spec.fill_value)

return out

for chunk_coords, chunk_selection, out_selection, _ in indexer:

# the GetResult status is discarded: missing INNER chunks of a

# present shard always fill (read_missing_chunks is a store-key

# level promise, applied to top-level statuses at the array layer)

decode_and_scatter_chunk(

shard_dict.get(chunk_coords),

out,

chunk_spec=chunk_spec,

chunk_selection=chunk_selection,

out_selection=out_selection,

drop_axes=(),

decode=inner_transform.decode_chunk,

)

return out

def _encode_sync(

self,

shard_array: NDBuffer,

shard_spec: ArraySpec,

) -> Buffer | None:

"""Encode a full shard synchronously.

Sync counterpart to `_encode_single`. This is reached when a

`ShardingCodec` is an *inner* codec of another sharding codec (nested

sharding): the outer codec encodes each inner chunk through its

`ChunkTransform`, which calls this method on the inner `ShardingCodec`.

Each inner chunk is encoded through the inner `ChunkTransform` and

collected into an intermediate `dict`. The dict's key order is

immaterial — the physical on-disk layout is decided downstream by the

`subchunk_write_order` loop in `_encode_shard_dict_sync` (this method

does NOT impose a layout). Empty inner chunks become `None` entries when

`write_empty_chunks` is False, signalling `_encode_shard_dict_sync` to

elide them from the data section and mark them empty in the shard index.

Returns `None` if every inner chunk was elided (an all-empty shard) —

callers treat that as "delete the shard key".

This method does not parallelize compression, but should.

See TODO: make issue for handling subchunk parallelism

For a partial write that only touches some inner chunks, use

`_encode_partial_sync` instead.

"""

shard_shape = shard_spec.shape

chunks_per_shard = self._get_chunks_per_shard(shard_spec)

chunk_spec = self._get_chunk_spec(shard_spec)

inner_transform = self._get_inner_chunk_transform(shard_spec)

indexer = BasicIndexer(

tuple(slice(0, s) for s in shard_shape),

shape=shard_shape,

chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape),

)

# Key order here is immaterial; _encode_shard_dict_sync lays the present

# chunks out in subchunk_write_order.

shard_builder: dict[tuple[int, ...], Buffer | None] = dict.fromkeys(

lexicographic_order_coords(chunks_per_shard)

)

for chunk_coords, _chunk_selection, out_selection, _ in indexer:

# None = chunk normalized to missing (see encode_or_elide_chunk)

shard_builder[chunk_coords] = encode_or_elide_chunk(

shard_array[out_selection], chunk_spec, inner_transform.encode_chunk

)

return self._encode_shard_dict_sync(

shard_builder,

chunks_per_shard=chunks_per_shard,

buffer_prototype=default_buffer_prototype(),

)

def _encode_partial_sync(

self,

byte_setter: Any,

value: NDBuffer,

selection: SelectorTuple,

shard_spec: ArraySpec,

) -> None:

"""Sync equivalent of `_encode_partial_single`.

Receives the source data for the written region (not a pre-merged

shard array) and the selection within the shard, matching the

calling convention of the async partial-encode path used by

`BatchedCodecPipeline`.

This method does not parallelize compression, but should.

See TODO: make issue for handling subchunk parallelism

Loads the existing shard, merges the written region into the affected

inner chunks, and rewrites the whole shard.

"""

chunks_per_shard = self._get_chunks_per_shard(shard_spec)

chunk_spec = self._get_chunk_spec(shard_spec)

inner_transform = self._get_inner_chunk_transform(shard_spec)

indexer, value = self._get_shard_indexer_and_value(selection, shard_spec, value)

is_complete = self._is_complete_shard_write(indexer, chunks_per_shard)

is_scalar = len(value.shape) == 0

# Load existing inner-chunk bytes into a dict (same structure as

# the async path's shard_dict).

if is_complete:

shard_dict: dict[tuple[int, ...], Buffer | None] = dict.fromkeys(

lexicographic_order_coords(chunks_per_shard)

)

else:

existing_bytes = byte_setter.get_sync(prototype=shard_spec.prototype)

if existing_bytes is not None:

shard_reader_fb = self._shard_reader_from_bytes_sync(

existing_bytes, chunks_per_shard

)

# Build the dict with one vectorized index lookup over all chunks,

# matching the async _encode_partial_single path. A per-coordinate

# __getitem__ loop here is O(n_chunks) Python overhead that dominates

# partial writes into shards with many inner chunks. The coordinate

# array and keys are cached on the reader, so neither is rebuilt here.

shard_dict = shard_reader_fb.to_dict_vectorized()

else:

shard_dict = dict.fromkeys(lexicographic_order_coords(chunks_per_shard))

# Merge, encode, and store each affected inner chunk into shard_dict via

# the canonical merge_and_encode_chunk (None = normalized to missing).

#

# Scalar fast path: when the written value is a scalar broadcast, every

# *complete* inner chunk is byte-for-byte identical — same fill, same

# empty-check, same encoded bytes. Compute that outcome once and reuse it

# for all complete chunks instead of re-merging, re-checking, and

# re-encoding tens of thousands of identical chunks. Incomplete (edge)

# chunks still merge against their own existing data individually.

# `_sentinel` distinguishes "not computed yet" from a memoized `None`

# (an empty chunk).

_sentinel = object()

scalar_complete_result: Buffer | object | None = _sentinel

for chunk_coords, chunk_sel, out_sel, is_complete_chunk in indexer:

if is_scalar and is_complete_chunk:

if scalar_complete_result is _sentinel:

scalar_complete_result = merge_and_encode_chunk(

None,

value,

chunk_spec=chunk_spec,

chunk_selection=chunk_sel,

out_selection=out_sel,

is_complete=is_complete_chunk,

drop_axes=(),

decode=inner_transform.decode_chunk,

encode=inner_transform.encode_chunk,

)

shard_dict[chunk_coords] = scalar_complete_result # type: ignore[assignment]

continue

# A complete chunk fully overwrites: skip decoding what it replaces.

existing_raw = None if is_complete_chunk else shard_dict.get(chunk_coords)

shard_dict[chunk_coords] = merge_and_encode_chunk(

existing_raw,

value,

chunk_spec=chunk_spec,

chunk_selection=chunk_sel,

out_selection=out_sel,

is_complete=is_complete_chunk,

drop_axes=(),

decode=inner_transform.decode_chunk,

encode=inner_transform.encode_chunk,

)

blob = self._encode_shard_dict_sync(

shard_dict,

chunks_per_shard=chunks_per_shard,

buffer_prototype=default_buffer_prototype(),

)

if blob is None:

byte_setter.delete_sync()

else:

byte_setter.set_sync(blob)

def _build_shard_layout(

self,

shard_dict: ShardMapping,

chunks_per_shard: tuple[int, ...],

) -> tuple[_ShardIndex, list[Buffer]] | None:

"""Lay out the present inner chunks of a shard. Pure compute, no IO.

Packs the encoded inner chunks (in the codec's `subchunk_write_order`)

into a contiguous data section and builds a shard index pointing each

present chunk at its ABSOLUTE byte offset within the final blob: when

the index is stored at the start, offsets are pre-shifted by the index

size (known without encoding — the index codecs are fixed-size, which

every index read path already relies on), so the index can be encoded

exactly once by the caller.

Returns `(index, data_buffers)`, or `None` for an all-empty shard (no

chunks present). Shared by the sync and async `_encode_shard_dict*` so

the layout/offset logic cannot drift between them.

"""

index = _ShardIndex.create_empty(chunks_per_shard)

buffers: list[Buffer] = []

chunk_start = (

self._shard_index_size(chunks_per_shard) if self.index_location == "start" else 0

)

for chunk_coords in self._subchunk_order_iter(chunks_per_shard, self.subchunk_write_order):

value = shard_dict.get(chunk_coords)

if value is None or len(value) == 0:

continue

chunk_length = len(value)

buffers.append(value)

index.set_chunk_slice(chunk_coords, slice(chunk_start, chunk_start + chunk_length))

chunk_start += chunk_length

if len(buffers) == 0:

return None

return index, buffers

def _assemble_shard(

self,

index_bytes: Buffer,

buffers: list[Buffer],

buffer_prototype: BufferPrototype,

*,

chunks_per_shard: tuple[int, ...],

) -> Buffer:

"""Concatenate the encoded index and data buffers into the shard blob.