from __future__ import annotations

import asyncio

import os

import threading

from concurrent.futures import ThreadPoolExecutor

from dataclasses import dataclass, field

from itertools import batched, chain, pairwise

from typing import TYPE_CHECKING, Any

from warnings import warn

from zarr.abc.codec import (

ArrayArrayCodec,

ArrayBytesCodec,

ArrayBytesCodecPartialDecodeMixin,

ArrayBytesCodecPartialEncodeMixin,

BytesBytesCodec,

Codec,

CodecPipeline,

GetResult,

)

from zarr.core.chunk_utils import (

ChunkTransform,

_merge_chunk_array,

chunk_is_empty,

decode_and_scatter_chunk,

evolve_codecs,

fill_value_or_default,

merge_and_encode_chunk,

scatter_chunk,

)

from zarr.core.common import concurrent_iter, concurrent_map

from zarr.core.config import config

from zarr.errors import ZarrUserWarning

from zarr.registry import register_pipeline

if TYPE_CHECKING:

from collections.abc import Iterable, Iterator, Sequence

from typing import Self

from zarr.abc.store import ByteGetter, ByteSetter

from zarr.core.array_spec import ArraySpec

from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer

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

from zarr.core.indexing import SelectorTuple

from zarr.core.metadata.v3 import ChunkGridMetadata

_pool: ThreadPoolExecutor | None = None

_pool_size: int = 0

_pool_lock = threading.Lock()

def _resolve_max_workers() -> int:

"""Helper for getting the maximum number of workers available to the `FusedCodecPipeline`"""

import os as _os

default = _os.cpu_count() or 1

cfg = config.get("codec_pipeline.max_workers", default=None)

if cfg is None:

return default

try:

return max(1, int(cfg))

except (TypeError, ValueError):

# This value arrives via the config/env layer (e.g.

# `ZARR_CODEC_PIPELINE__MAX_WORKERS`), so tolerate bad input here

# instead of raising mid-read.

warn(

f"Ignoring invalid `codec_pipeline.max_workers` config value {cfg!r}; "

f"falling back to {default}.",

category=ZarrUserWarning,

stacklevel=2,

)

return default

def _get_pool(max_workers: int) -> ThreadPoolExecutor:

"""Get or create the module-level thread pool, sized to `max_workers`.

The pool grows on demand — if a request arrives for more workers than the

current pool has, it is replaced with a larger one. The previous pool is NOT

shut down here: another thread may be holding a reference to it and about to

submit (`shutdown` would make its `pool.map` raise "cannot schedule new

futures after shutdown"). The orphaned pool finishes its in-flight tasks and

is garbage-collected once no caller references it. The pool only grows, never

shrinks (a shrink request reuses the larger pool, leaving workers idle).

Callers that want sequential execution should not call this — they

should run the task list inline. `max_workers` must be >= 1.

"""

global _pool, _pool_size

if max_workers < 1:

raise ValueError(f"max_workers must be >= 1, got {max_workers}")

if _pool is None or _pool_size < max_workers:

with _pool_lock:

if _pool is None or _pool_size < max_workers:

# Replace without shutting down the old pool (see docstring):

# avoids a race with a concurrent in-flight pool.map on it.

_pool = ThreadPoolExecutor(max_workers=max_workers)

_pool_size = max_workers

return _pool

def _reset_pool_after_fork() -> None:

"""Discard the pool a forked child inherits from its parent.

Executor threads do not survive `fork()`, so the inherited pool accepts

submissions that no thread ever runs. The lock is recreated too: it may

have been held by a thread that no longer exists in the child.

"""

global _pool, _pool_size, _pool_lock

_pool = None # pragma: no cover

_pool_size = 0 # pragma: no cover

_pool_lock = threading.Lock() # pragma: no cover

# this is only available on certain operating systems

if hasattr(os, "register_at_fork"):

os.register_at_fork(after_in_child=_reset_pool_after_fork)

def _unzip2[T, U](iterable: Iterable[tuple[T, U]]) -> tuple[list[T], list[U]]:

out0: list[T] = []

out1: list[U] = []

for item0, item1 in iterable:

out0.append(item0)

out1.append(item1)

return (out0, out1)

def resolve_batched(codec: Codec, chunk_specs: Iterable[ArraySpec]) -> Iterable[ArraySpec]:

return [codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs]

def resolve_aa_specs(

aa_codecs: tuple[Codec, ...], chunk_spec: ArraySpec

) -> tuple[tuple[ArraySpec, ...], ArraySpec]:

"""Resolve the per-stage chunk specs for a single chunk's codec chain.

Threads `chunk_spec` forward through the array->array codecs via

`resolve_metadata` (each codec sees the spec produced by the previous one),

returning `(aa_specs, ab_spec)`:

* `aa_specs[i]` is the spec the i-th AA codec operates on (its *input* on

encode / *output* on decode);

* `ab_spec` is the spec after all AA codecs — what the array->bytes codec

and the bytes->bytes codecs operate on.

This is the single source of truth for per-stage spec evolution, shared by

the synchronous `ChunkTransform` and the asynchronous

`AsyncChunkTransform`. It is pure metadata (only `resolve_metadata`), so

it places no synchronous-codec requirement on the codecs.

"""

aa_specs: list[ArraySpec] = []

spec = chunk_spec

for aa_codec in aa_codecs:

aa_specs.append(spec)

spec = aa_codec.resolve_metadata(spec)

return tuple(aa_specs), spec

def pipeline_supports_partial_decode(

array_bytes_codec: ArrayBytesCodec,

*,

array_array_codecs: tuple[ArrayArrayCodec, ...],

bytes_bytes_codecs: tuple[BytesBytesCodec, ...],

require_no_aa_bb: bool,

) -> bool:

"""Whether a codec pipeline can decode a partial selection without a full read.

Requires the array->bytes codec to implement

``ArrayBytesCodecPartialDecodeMixin``. When ``require_no_aa_bb`` is True it

additionally requires no array->array / bytes->bytes codecs, because those

can change the slice<->byte-range correspondence (an AA codec can make the

selection non-contiguous, a BB codec can rewrite the bytes), making partial

decode infeasible.

Both pipelines pass `require_no_aa_bb=True`: an outer AA/BB codec (e.g. a

compressor wrapping a sharding serializer) must see every byte of the

chunk, so a partial branch that only re-decodes/re-encodes the inner

sharding codec would silently bypass it.

"""

if require_no_aa_bb and (len(array_array_codecs) + len(bytes_bytes_codecs)) != 0:

return False

return isinstance(array_bytes_codec, ArrayBytesCodecPartialDecodeMixin)

def pipeline_supports_partial_encode(

array_bytes_codec: ArrayBytesCodec,

*,

array_array_codecs: tuple[ArrayArrayCodec, ...],

bytes_bytes_codecs: tuple[BytesBytesCodec, ...],

require_no_aa_bb: bool,

) -> bool:

"""Whether a codec pipeline can encode a partial selection without a full rewrite.

Mirror of `pipeline_supports_partial_decode` for encoding.

"""

if require_no_aa_bb and (len(array_array_codecs) + len(bytes_bytes_codecs)) != 0:

return False

return isinstance(array_bytes_codec, ArrayBytesCodecPartialEncodeMixin)

async def _cancel_and_drain(futures: Iterable[asyncio.Future[Any]]) -> None:

"""Cancel every not-yet-done future/task and await its outcome.

Used to clean up work spawned by a drain loop (`asyncio.as_completed` +

`await`) when the loop exits early via exception. Without this, tasks

already spawned keep running unattended after the caller has moved on,

and an eventual failure surfaces as an unraisable "exception was never

retrieved" warning instead of being observed here.

"""

pending = [f for f in futures if not f.done()]

if len(pending) == 0:

return

for f in pending:

f.cancel()

await asyncio.gather(*pending, return_exceptions=True)

async def _fetch_and_decode_as_completed(

batch: Sequence[tuple[ByteGetter | None, ArraySpec]],

transform: ChunkTransform,

) -> list[NDBuffer | None]:

"""Concurrently fetch each chunk's bytes and decode it as fetches complete.

Decoding overlaps with in-flight fetches: each chunk is decoded the moment

its bytes arrive (in a thread pool when one is available, otherwise inline)

rather than waiting for the whole batch to land. A `None` byte getter fetches

nothing and decodes to `None`. Results are returned in input order.

"""

max_workers = _resolve_max_workers()

pool = _get_pool(max_workers) if max_workers > 1 else None

loop = asyncio.get_running_loop()

decode_futures: list[asyncio.Future[NDBuffer | None]] = [loop.create_future() for _ in batch]

async def _fetch(

idx: int, byte_getter: ByteGetter | None, prototype: BufferPrototype

) -> tuple[int, Buffer | None]:

return idx, None if byte_getter is None else await byte_getter.get(prototype)

def _decode(buffer: Buffer | None, chunk_spec: ArraySpec) -> NDBuffer | None:

return None if buffer is None else transform.decode_chunk(buffer, chunk_spec)

fetch_tasks = concurrent_iter(

[

(idx, byte_getter, chunk_spec.prototype)

for idx, (byte_getter, chunk_spec) in enumerate(batch)

],

_fetch,

config.get("async.concurrency"),

)

try:

for fetch_coro in asyncio.as_completed(fetch_tasks):

idx, buffer = await fetch_coro

chunk_spec = batch[idx][1]

# Bridge both paths to asyncio.Future so the final collection loop

# can `await` uniformly without blocking the event loop. For the

# pool path that means `wrap_future` (not `pool.submit(...).result()`,

# which would block the loop thread for the duration of every decode

# — freezing any unrelated coroutines sharing this loop).

if pool is None:

decode_futures[idx].set_result(_decode(buffer, chunk_spec))

else:

decode_futures[idx] = asyncio.wrap_future(pool.submit(_decode, buffer, chunk_spec))

return await asyncio.gather(*decode_futures)

finally:

# On the happy path every future here is already done, so this is a

# no-op; on failure it stops abandoned fetches/decodes from

# continuing to run unattended after this function has raised. A

# single call over both iterables (not two sequential calls) so that

# outer-task cancellation during the first drain can't skip the

# second, leaving its futures/tasks unobserved.

await _cancel_and_drain(chain(fetch_tasks, decode_futures))

async def _encode_and_write_as_completed(

batch: Sequence[tuple[ByteSetter, NDBuffer | None, ArraySpec]],

transform: ChunkTransform,

) -> None:

"""Encode each chunk and write it out as encodes complete.

The reverse of `_fetch_and_decode_as_completed`: each chunk is encoded (in a

thread pool when one is available, otherwise inline) and its write is launched

the moment that chunk's encode finishes — overlapping IO writes with

still-running encodes rather than waiting for the whole batch to encode.

A `None` chunk array encodes to `None` and is `delete`d. Writes are bounded

by `async.concurrency`.

"""

max_workers = _resolve_max_workers()

pool = _get_pool(max_workers) if max_workers > 1 else None

loop = asyncio.get_running_loop()

semaphore = asyncio.Semaphore(config.get("async.concurrency"))

def _encode(

idx: int, chunk_array: NDBuffer | None, chunk_spec: ArraySpec

) -> tuple[int, Buffer | None]:

return idx, None if chunk_array is None else transform.encode_chunk(chunk_array, chunk_spec)

async def _write(idx: int, chunk_bytes: Buffer | None) -> None:

byte_setter = batch[idx][0]

async with semaphore:

if chunk_bytes is None:

await byte_setter.delete()

else:

await byte_setter.set(chunk_bytes)

# Submit every encode up front. The pool path bridges to asyncio.Future via

# `wrap_future` (not `pool.submit(...).result()`, which would block the loop

# thread); the inline path resolves immediately.

encode_futures: list[asyncio.Future[tuple[int, Buffer | None]]] = []

for idx, (_, chunk_array, chunk_spec) in enumerate(batch):

if pool is None:

fut: asyncio.Future[tuple[int, Buffer | None]] = loop.create_future()

fut.set_result(_encode(idx, chunk_array, chunk_spec))

else:

fut = asyncio.wrap_future(pool.submit(_encode, idx, chunk_array, chunk_spec))

encode_futures.append(fut)

# Kick off each chunk's write the instant its encode lands, so writes of

# already-compressed chunks proceed while the rest are still encoding.

write_tasks: list[asyncio.Task[None]] = []

try:

for encode_coro in asyncio.as_completed(encode_futures):

idx, chunk_bytes = await encode_coro

write_tasks.append(asyncio.ensure_future(_write(idx, chunk_bytes)))

await asyncio.gather(*write_tasks)

finally:

# On the happy path every future here is already done, so this is a

# no-op; on failure (an encode or a write raising) it stops

# already-spawned writes from continuing in the background after

# this function has raised. A single call over both iterables (not

# two sequential calls) so that outer-task cancellation during the

# first drain can't skip the second, leaving its futures/tasks

# unobserved.

await _cancel_and_drain(chain(write_tasks, encode_futures))

async def _async_read_fallback(

pipeline: CodecPipeline,

batch: list[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],

out: NDBuffer,

drop_axes: tuple[int, ...],

) -> tuple[GetResult, ...]:

"""Async fallback read used when no fast-path is available.

Fetches every chunk's bytes via `concurrent_map` (sized by

`async.concurrency`), decodes the batch through `pipeline.decode`,

then scatters each decoded chunk into `out` at its `out_selection`.

Used by both `BatchedCodecPipeline.read_batch` (non-partial-decode

branch) and `FusedCodecPipeline.read` (when the store does not advertise

sync IO / sync transform is unavailable).

"""

chunk_array_batch: list[NDBuffer | None]

if isinstance(pipeline, FusedCodecPipeline) and pipeline.sync_transform is not None:

chunk_array_batch = await _fetch_and_decode_as_completed(

[(byte_getter, chunk_spec) for byte_getter, chunk_spec, *_ in batch],

pipeline.sync_transform,

)

else:

chunk_bytes_batch = await concurrent_map(

[(byte_getter, array_spec.prototype) for byte_getter, array_spec, *_ in batch],

lambda byte_getter, prototype: byte_getter.get(prototype),

config.get("async.concurrency"),

)

chunk_array_batch = list(

await pipeline.decode(

[

(chunk_bytes, chunk_spec)

for chunk_bytes, (_, chunk_spec, *_) in zip(

chunk_bytes_batch, batch, strict=False

)

],

)

)

results: list[GetResult] = []

for chunk_array, (_, chunk_spec, chunk_selection, out_selection, _) in zip(

chunk_array_batch, batch, strict=True

):

selected = None if chunk_array is None else chunk_array[chunk_selection]

results.append(

scatter_chunk(

selected,

out,

chunk_spec=chunk_spec,

out_selection=out_selection,

drop_axes=drop_axes,

)

)

return tuple(results)

async def _async_write_fallback(

pipeline: CodecPipeline,

batch: list[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],

value: NDBuffer,

drop_axes: tuple[int, ...],

) -> None:

"""Async fallback write used when no fast-path is available.

For each chunk in `batch`: read its existing bytes from the store

(skipping the read for complete chunks), decode the batch via

`pipeline.decode`, merge `value` into each decoded chunk via

`_merge_chunk_array`, drop chunks that are all-fill when

`write_empty_chunks` is False, encode the surviving chunks via

`pipeline.encode`, then `set` the encoded bytes (or `delete`

if encoding produced `None` or the chunk dropped).

Used by both `BatchedCodecPipeline.write_batch` (non-partial-encode

branch) and `FusedCodecPipeline.write` (when the store does not advertise

sync IO / sync transform is unavailable).

"""

sync_transform = pipeline.sync_transform if isinstance(pipeline, FusedCodecPipeline) else None

if sync_transform is not None:

# Read each chunk's existing bytes (skipping complete chunks) and decode

# as fetches complete, overlapping the sync decode with in-flight reads.

chunk_array_decoded: Iterable[NDBuffer | None] = await _fetch_and_decode_as_completed(

[

(None if is_complete_chunk else byte_setter, chunk_spec)

for byte_setter, chunk_spec, _, _, is_complete_chunk in batch

],

sync_transform,

)

else:

async def _read_key(

byte_setter: ByteSetter | None, prototype: BufferPrototype

) -> Buffer | None:

if byte_setter is None:

return None

return await byte_setter.get(prototype=prototype)

chunk_bytes_batch: Iterable[Buffer | None] = await concurrent_map(

[

(

None if is_complete_chunk else byte_setter,

chunk_spec.prototype,

)

for byte_setter, chunk_spec, chunk_selection, _, is_complete_chunk in batch

],

_read_key,

config.get("async.concurrency"),

)

chunk_array_decoded = await pipeline.decode(

[

(chunk_bytes, chunk_spec)

for chunk_bytes, (_, chunk_spec, *_) in zip(chunk_bytes_batch, batch, strict=False)

],

)

chunk_array_merged = [

_merge_chunk_array(

chunk_array,

value,

out_selection,

chunk_spec,

chunk_selection,

is_complete_chunk,

drop_axes,

)

for chunk_array, (

_,

chunk_spec,

chunk_selection,

out_selection,

is_complete_chunk,

) in zip(chunk_array_decoded, batch, strict=False)

]

# _merge_chunk_array always returns a real NDBuffer (never None), so the only

# way a chunk drops to None here is the empty-chunk normalization.

chunk_array_batch: list[NDBuffer | None] = [

None if chunk_is_empty(chunk_array, chunk_spec) else chunk_array

for chunk_array, (_, chunk_spec, *_) in zip(chunk_array_merged, batch, strict=False)

]

if sync_transform is not None:

await _encode_and_write_as_completed(

[

(byte_setter, chunk_array, chunk_spec)

for chunk_array, (byte_setter, chunk_spec, *_) in zip(

chunk_array_batch, batch, strict=False

)

],

sync_transform,

)

else:

chunk_bytes_batch = await pipeline.encode(

[

(chunk_array, chunk_spec)

for chunk_array, (_, chunk_spec, *_) in zip(chunk_array_batch, batch, strict=False)

],

)

async def _write_key(byte_setter: ByteSetter, chunk_bytes: Buffer | None) -> None:

if chunk_bytes is None:

await byte_setter.delete()

else:

await byte_setter.set(chunk_bytes)

await concurrent_map(

[

(byte_setter, chunk_bytes)

for chunk_bytes, (byte_setter, *_) in zip(chunk_bytes_batch, batch, strict=False)

],

_write_key,

config.get("async.concurrency"),

)

@dataclass(slots=True, kw_only=True)

class AsyncChunkTransform:

"""A per-chunk asynchronous codec chain — the async mirror of ChunkTransform.

Decodes/encodes a SINGLE chunk through the full codec chain, awaiting each

codec's per-chunk async method (`_decode_single`/`_encode_single`) with the

correctly-evolved per-stage spec (via the shared `resolve_aa_specs`).

Unlike ChunkTransform it places no `SupportsSyncCodec` requirement on the

codecs, so it works for async-only codecs. Unlike the batched codec API it

operates on one chunk at a time — the mini-batch fan-out is a

BatchedCodecPipeline concern and deliberately not reintroduced here.

"""

codecs: tuple[Codec, ...]

_aa_codecs: tuple[ArrayArrayCodec, ...] = field(init=False, repr=False, compare=False)

_ab_codec: ArrayBytesCodec = field(init=False, repr=False, compare=False)

_bb_codecs: tuple[BytesBytesCodec, ...] = field(init=False, repr=False, compare=False)

def __post_init__(self) -> None:

# `AsyncChunkTransform` is (re)constructed per decode/encode call from a

# codec chain that already went through `codecs_from_list` when the

# pipeline itself was built, so re-splitting it here must not re-emit

# that chain's advisory warnings on every call.

aa, ab, bb = codecs_from_list_unchecked(list(self.codecs))

self._aa_codecs = aa

self._ab_codec = ab

self._bb_codecs = bb

async def decode_chunk(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> NDBuffer:

"""Decode one chunk through the chain (bb -> ab -> aa), async."""

aa_specs, ab_spec = resolve_aa_specs(self._aa_codecs, chunk_spec)

data: Buffer = chunk_bytes

for bb_codec in reversed(self._bb_codecs):

data = await bb_codec._decode_single(data, ab_spec)

chunk_array: NDBuffer = await self._ab_codec._decode_single(data, ab_spec)

for aa_codec, aa_spec in zip(reversed(self._aa_codecs), reversed(aa_specs), strict=True):

chunk_array = await aa_codec._decode_single(chunk_array, aa_spec)

return chunk_array

async def encode_chunk(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> Buffer | None:

"""Encode one chunk through the chain (aa -> ab -> bb), async.

Returns None if any stage drops the chunk (e.g. an all-fill chunk under

write_empty_chunks=False), matching ChunkTransform.encode_chunk.

"""

aa_specs, ab_spec = resolve_aa_specs(self._aa_codecs, chunk_spec)

aa_data: NDBuffer = chunk_array

for aa_codec, aa_spec in zip(self._aa_codecs, aa_specs, strict=True):

aa_result = await aa_codec._encode_single(aa_data, aa_spec)

if aa_result is None:

return None

aa_data = aa_result

ab_result = await self._ab_codec._encode_single(aa_data, ab_spec)

if ab_result is None:

return None

bb_data: Buffer = ab_result

for bb_codec in self._bb_codecs:

bb_result = await bb_codec._encode_single(bb_data, ab_spec)

if bb_result is None:

return None

bb_data = bb_result

return bb_data

@dataclass(frozen=True)

class BatchedCodecPipeline(CodecPipeline):

"""Default codec pipeline.

This batched codec pipeline divides the chunk batches into batches of a configurable

batch size ("mini-batch"). Fetching, decoding, encoding and storing are performed in

lock step for each mini-batch. Multiple mini-batches are processing concurrently.

"""

array_array_codecs: tuple[ArrayArrayCodec, ...]

array_bytes_codec: ArrayBytesCodec

bytes_bytes_codecs: tuple[BytesBytesCodec, ...]

batch_size: int

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

# Re-splits an already-`codecs_from_list`-validated (and warned-about)

# chain against the evolved spec, so this uses the quiet variant rather

# than routing through `from_codecs` (which would re-warn).

evolved_codecs = evolve_codecs(self, array_spec)

array_array_codecs, array_bytes_codec, bytes_bytes_codecs = codecs_from_list_unchecked(

evolved_codecs

)

return type(self)(

array_array_codecs=array_array_codecs,

array_bytes_codec=array_bytes_codec,

bytes_bytes_codecs=bytes_bytes_codecs,

batch_size=self.batch_size,

)

@classmethod

def from_codecs(cls, codecs: Iterable[Codec], *, batch_size: int | None = None) -> Self:

array_array_codecs, array_bytes_codec, bytes_bytes_codecs = codecs_from_list(codecs)

return cls(

array_array_codecs=array_array_codecs,

array_bytes_codec=array_bytes_codec,

bytes_bytes_codecs=bytes_bytes_codecs,

batch_size=batch_size or config.get("codec_pipeline.batch_size"),

)

@property

def supports_partial_decode(self) -> bool:

# Only a single ArrayBytesCodec that supports partial decoding, and no

# AA/BB codecs (they break the slice<->byte-range correspondence).

return pipeline_supports_partial_decode(

self.array_bytes_codec,

array_array_codecs=self.array_array_codecs,

bytes_bytes_codecs=self.bytes_bytes_codecs,

require_no_aa_bb=True,

)

@property

def supports_partial_encode(self) -> bool:

return pipeline_supports_partial_encode(

self.array_bytes_codec,

array_array_codecs=self.array_array_codecs,

bytes_bytes_codecs=self.bytes_bytes_codecs,

require_no_aa_bb=True,

)

def __iter__(self) -> Iterator[Codec]:

yield from self.array_array_codecs

yield self.array_bytes_codec

yield from self.bytes_bytes_codecs

def validate(

self,

*,

shape: tuple[int, ...],

dtype: ZDType[TBaseDType, TBaseScalar],

chunk_grid: ChunkGridMetadata,

) -> None:

for codec in self:

codec.validate(shape=shape, dtype=dtype, chunk_grid=chunk_grid)

def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int:

for codec in self:

byte_length = codec.compute_encoded_size(byte_length, array_spec)

array_spec = codec.resolve_metadata(array_spec)

return byte_length

def _codecs_with_resolved_metadata_batched(

self, chunk_specs: Iterable[ArraySpec]

) -> tuple[

list[tuple[ArrayArrayCodec, list[ArraySpec]]],

tuple[ArrayBytesCodec, list[ArraySpec]],

list[tuple[BytesBytesCodec, list[ArraySpec]]],

]:

aa_codecs_with_spec: list[tuple[ArrayArrayCodec, list[ArraySpec]]] = []

chunk_specs = list(chunk_specs)

for aa_codec in self.array_array_codecs:

aa_codecs_with_spec.append((aa_codec, chunk_specs))

chunk_specs = [aa_codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs]

ab_codec_with_spec = (self.array_bytes_codec, chunk_specs)

chunk_specs = [

self.array_bytes_codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs

]

bb_codecs_with_spec: list[tuple[BytesBytesCodec, list[ArraySpec]]] = []

for bb_codec in self.bytes_bytes_codecs:

bb_codecs_with_spec.append((bb_codec, chunk_specs))

chunk_specs = [bb_codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs]

return (aa_codecs_with_spec, ab_codec_with_spec, bb_codecs_with_spec)

async def decode_batch(

self,

chunk_bytes_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],

) -> Iterable[NDBuffer | None]:

chunk_bytes_batch: Iterable[Buffer | None]

chunk_bytes_batch, chunk_specs = _unzip2(chunk_bytes_and_specs)

(

aa_codecs_with_spec,

ab_codec_with_spec,

bb_codecs_with_spec,

) = self._codecs_with_resolved_metadata_batched(chunk_specs)

for bb_codec, chunk_spec_batch in bb_codecs_with_spec[::-1]:

chunk_bytes_batch = await bb_codec.decode(

zip(chunk_bytes_batch, chunk_spec_batch, strict=False)

)

ab_codec, chunk_spec_batch = ab_codec_with_spec

chunk_array_batch = await ab_codec.decode(

zip(chunk_bytes_batch, chunk_spec_batch, strict=False)

)

for aa_codec, chunk_spec_batch in aa_codecs_with_spec[::-1]:

chunk_array_batch = await aa_codec.decode(

zip(chunk_array_batch, chunk_spec_batch, strict=False)

)

return chunk_array_batch

async def decode_partial_batch(

self,

batch_info: Iterable[tuple[ByteGetter, SelectorTuple, ArraySpec]],

) -> Iterable[NDBuffer | None]:

codec = self.array_bytes_codec

if not self.supports_partial_decode or not isinstance(

codec, ArrayBytesCodecPartialDecodeMixin

):

raise ValueError("This codec pipeline does not support partial decoding.")

return await codec.decode_partial(batch_info)

async def encode_batch(

self,

chunk_arrays_and_specs: Iterable[tuple[NDBuffer | None, ArraySpec]],

) -> Iterable[Buffer | None]:

chunk_array_batch: Iterable[NDBuffer | None]

chunk_specs: Iterable[ArraySpec]

chunk_array_batch, chunk_specs = _unzip2(chunk_arrays_and_specs)

for aa_codec in self.array_array_codecs:

chunk_array_batch = await aa_codec.encode(

zip(chunk_array_batch, chunk_specs, strict=False)

)

chunk_specs = resolve_batched(aa_codec, chunk_specs)

chunk_bytes_batch = await self.array_bytes_codec.encode(

zip(chunk_array_batch, chunk_specs, strict=False)

)

chunk_specs = resolve_batched(self.array_bytes_codec, chunk_specs)

for bb_codec in self.bytes_bytes_codecs:

chunk_bytes_batch = await bb_codec.encode(

zip(chunk_bytes_batch, chunk_specs, strict=False)

)

chunk_specs = resolve_batched(bb_codec, chunk_specs)

return chunk_bytes_batch

async def encode_partial_batch(

self,

batch_info: Iterable[tuple[ByteSetter, NDBuffer, SelectorTuple, ArraySpec]],

) -> None:

codec = self.array_bytes_codec

if not self.supports_partial_encode or not isinstance(

codec, ArrayBytesCodecPartialEncodeMixin

):

raise ValueError("This codec pipeline does not support partial encoding.")

await codec.encode_partial(batch_info)

async def read_batch(

self,

batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],

out: NDBuffer,

drop_axes: tuple[int, ...] = (),

) -> tuple[GetResult, ...]:

if self.supports_partial_decode:

results: list[GetResult] = []

batch_info_list = list(batch_info)

chunk_array_batch = await self.decode_partial_batch(

[

(byte_getter, chunk_selection, chunk_spec)

for byte_getter, chunk_spec, chunk_selection, *_ in batch_info_list

]

)

for chunk_array, (_, chunk_spec, _, out_selection, _) in zip(

chunk_array_batch, batch_info_list, strict=False

):

if chunk_array is not None:

if drop_axes:

chunk_array = chunk_array.squeeze(axis=drop_axes)

out[out_selection] = chunk_array

results.append(GetResult(status="present"))

else:

out[out_selection] = fill_value_or_default(chunk_spec)

results.append(GetResult(status="missing"))

return tuple(results)

return await _async_read_fallback(self, list(batch_info), out, drop_axes)

async def write_batch(

self,

batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],

value: NDBuffer,

drop_axes: tuple[int, ...] = (),

) -> None:

if self.supports_partial_encode:

# Pass scalar values as is

if len(value.shape) == 0:

await self.encode_partial_batch(

[

(byte_setter, value, chunk_selection, chunk_spec)

for byte_setter, chunk_spec, chunk_selection, out_selection, _ in batch_info

],

)

else:

await self.encode_partial_batch(

[

(byte_setter, value[out_selection], chunk_selection, chunk_spec)

for byte_setter, chunk_spec, chunk_selection, out_selection, _ in batch_info

],

)

return

await _async_write_fallback(self, list(batch_info), value, drop_axes)

async def decode(

self,

chunk_bytes_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],

) -> Iterable[NDBuffer | None]:

output: list[NDBuffer | None] = []

for batch_info in batched(chunk_bytes_and_specs, self.batch_size):

output.extend(await self.decode_batch(batch_info))

return output

async def encode(

self,

chunk_arrays_and_specs: Iterable[tuple[NDBuffer | None, ArraySpec]],

) -> Iterable[Buffer | None]:

output: list[Buffer | None] = []

for single_batch_info in batched(chunk_arrays_and_specs, self.batch_size):

output.extend(await self.encode_batch(single_batch_info))

return output

async def read(

self,

batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],

out: NDBuffer,

drop_axes: tuple[int, ...] = (),

) -> tuple[GetResult, ...]:

batch_results = await concurrent_map(

[

(single_batch_info, out, drop_axes)

for single_batch_info in batched(batch_info, self.batch_size)

],

self.read_batch,

config.get("async.concurrency"),

)

results: list[GetResult] = []

for batch in batch_results:

results.extend(batch)

return tuple(results)

async def write(

self,

batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],

value: NDBuffer,

drop_axes: tuple[int, ...] = (),

) -> None:

await concurrent_map(

[

(single_batch_info, value, drop_axes)

for single_batch_info in batched(batch_info, self.batch_size)

],

self.write_batch,

config.get("async.concurrency"),

)

def codecs_from_list(

codecs: Iterable[Codec],

) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]:

"""Split `codecs` into `(array_array, array_bytes, bytes_bytes)`, validating order.

Emits user-facing advisory warnings about the codec chain (e.g. sharding's

"disables partial reads" warning). Use this for the FIRST construction of a

codec chain from user-supplied codecs. Use `codecs_from_list_unchecked` when

re-splitting a chain that was already validated and warned about by a prior

`codecs_from_list` call (e.g. `evolve_from_array_spec` re-splitting the same

codecs against an evolved spec) — re-warning there would fire the same

advisory once per reconstruction instead of once per user-facing chain.

"""

from zarr.codecs.sharding import ShardingCodec

codecs = tuple(codecs) # materialize to avoid generator consumption issues

if any(isinstance(codec, ShardingCodec) for codec in codecs) and len(codecs) > 1:

warn(

"Combining a `sharding_indexed` codec disables partial reads and "

"writes, which may lead to inefficient performance.",

category=ZarrUserWarning,

stacklevel=3,

)

return codecs_from_list_unchecked(codecs)

def codecs_from_list_unchecked(

codecs: Iterable[Codec],

) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]:

"""Split `codecs` into `(array_array, array_bytes, bytes_bytes)`, validating order.

Same structural validation as `codecs_from_list` (raises on bad codec

ordering or a missing/duplicate array-bytes codec) but does NOT emit

user-facing advisory warnings. See `codecs_from_list` for when to use each.

"""

codecs = tuple(codecs) # materialize to avoid generator consumption issues

array_array: tuple[ArrayArrayCodec, ...] = ()

array_bytes_maybe: ArrayBytesCodec | None = None

bytes_bytes: tuple[BytesBytesCodec, ...] = ()

for prev_codec, cur_codec in pairwise((None, *codecs)):

if isinstance(cur_codec, ArrayArrayCodec):

if isinstance(prev_codec, ArrayBytesCodec | BytesBytesCodec):

msg = (

f"Invalid codec order. ArrayArrayCodec {cur_codec}"

"must be preceded by another ArrayArrayCodec. "

f"Got {type(prev_codec)} instead."

)

raise TypeError(msg)

array_array += (cur_codec,)

elif isinstance(cur_codec, ArrayBytesCodec):

if isinstance(prev_codec, BytesBytesCodec):

msg = (

f"Invalid codec order. ArrayBytes codec {cur_codec}"

f" must be preceded by an ArrayArrayCodec. Got {type(prev_codec)} instead."

)

raise TypeError(msg)

if array_bytes_maybe is not None:

msg = (

f"Got two instances of ArrayBytesCodec: {array_bytes_maybe} and {cur_codec}. "

"Only one array-to-bytes codec is allowed."

)

raise ValueError(msg)

array_bytes_maybe = cur_codec

elif isinstance(cur_codec, BytesBytesCodec):

if isinstance(prev_codec, ArrayArrayCodec):

msg = (

f"Invalid codec order. BytesBytesCodec {cur_codec}"

"must be preceded by either another BytesBytesCodec, or an ArrayBytesCodec. "

f"Got {type(prev_codec)} instead."

)

raise TypeError(msg)

bytes_bytes += (cur_codec,)

else:

raise TypeError

if array_bytes_maybe is None:

raise ValueError("Required ArrayBytesCodec was not found.")

else:

return array_array, array_bytes_maybe, bytes_bytes

register_pipeline(BatchedCodecPipeline)

@dataclass(frozen=True)

class FusedCodecPipeline(CodecPipeline):

"""Codec pipeline that runs codec compute synchronously, in bulk.

This is an opt-in alternative to `BatchedCodecPipeline`. The win is NOT

"separating IO from compute" — the codecs (notably `ShardingCodec`) still

perform their own storage IO. The win is replacing the batched pipeline's

per-chunk *async scheduling* (≈one coroutine per chunk, which dominates real

codec work) with synchronous, batched/coalesced execution:

1. When every codec implements `SupportsSyncCodec`, a `ChunkTransform`

runs the codec chain synchronously (no event loop, no per-chunk coroutine)

— optionally across a thread pool for CPU-heavy decode/encode.

2. Sharded reads use the codec's synchronous IO methods: byte-range reads

coalesced via `Store.get_ranges_sync`, and a vectorized whole-shard bulk

decode for dense, fixed-size, uncompressed shards. Sharded writes go

through the codec's synchronous full-shard-rewrite path.

3. When the store lacks synchronous IO (e.g. ZipStore) the pipeline falls

back to the async path, equivalent to `BatchedCodecPipeline`.

IO ownership: the sharding codec holds the byte getter/setter and reads/

writes storage directly (the same model as zarrs; unlike tensorstore, which

keeps codecs storage-free). A storage-free codec is a possible future

direction (see the pure-codec design notes) but is explicitly NOT what this

pipeline does.

"""

codecs: tuple[Codec, ...]

array_array_codecs: tuple[ArrayArrayCodec, ...]