Skip to content

feat(context): scope ArrayIO's extension resolution with a Context #420

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
d-v-b wants to merge 1 commit into
base: claude/arrayio-lazy-metadata-48d8e5
Choose a base branch
from
Draft

feat(context): scope ArrayIO's extension resolution with a Context #420

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/420.feature.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Added a `context` to `zarr.ArrayIO` and `zarr.AsyncArrayIO`: a value defining the data types, codecs and codec pipeline in scope for one array. Extension names in a stored document or in creation keyword arguments resolve through the context instead of the default registries; a kind the context leaves unset falls back to the defaults, while a kind it sets is a closed world, and a name outside it fails to resolve with an error naming the scope.
90 changes: 68 additions & 22 deletions src/zarr/core/array.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,7 @@
from zarr.abc.store import Store
from zarr.codecs.sharding import IndexLocation, ShardingCodec
from zarr.core.buffer import Buffer
from zarr.core.context import Context
from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
from zarr.storage import StoreLike
from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3
Expand Down Expand Up @@ -220,9 +221,13 @@ def _chunk_sizes_from_shape(
return tuple(result)


def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata:
def parse_array_metadata(
data: Any, path: str | None = None, *, context: Context | None = None
) -> ArrayMetadata:
"""Array metadata from a metadata object or a metadata document, naming the array at
`path` in warnings about how an invalid document was read.
`path` in warnings about how an invalid document was read. With a `context`, a
document's extension names resolve through it instead of the default registries;
a metadata object is already resolved and passes through any context.

`ArrayV2Metadata` accepts a chunk size of 0, as it always has, though only an
invalid document holds one: such metadata is read as the documents it would store
Expand All @@ -235,7 +240,7 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata:
if isinstance(data, dict):
zarr_format = data.get("zarr_format")
if zarr_format == 3:
meta_out = ArrayV3Metadata.from_dict(data, path=path)
meta_out = ArrayV3Metadata.from_dict(data, path=path, context=context)
if len(meta_out.storage_transformers) > 0:
msg = (
f"Array metadata contains storage transformers: {meta_out.storage_transformers}."
Expand All @@ -244,16 +249,19 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata:
raise ValueError(msg)
return meta_out
elif zarr_format == 2:
return ArrayV2Metadata.from_dict(data, path=path)
return ArrayV2Metadata.from_dict(data, path=path, context=context)
else:
raise ValueError(f"Invalid zarr_format: {zarr_format}. Expected 2 or 3")
raise TypeError # pragma: no cover


def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None) -> CodecPipeline:
def create_codec_pipeline(
metadata: ArrayMetadata, *, store: Store | None = None, context: Context | None = None
) -> CodecPipeline:
pipeline_class = get_pipeline_class() if context is None else context.pipeline_class()
if store is not None:
try:
return get_pipeline_class().from_array_metadata_and_store(
return pipeline_class.from_array_metadata_and_store(
array_metadata=metadata, store=store
)
except NotImplementedError:
Expand All @@ -266,7 +274,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
# "disables partial reads" warning) fire here; `evolve_from_array_spec`
# re-splits the same already-warned-about chain via
# `codecs_from_list_unchecked`, so it does not re-emit them.
pipeline = get_pipeline_class().from_codecs(metadata.codecs)
pipeline = pipeline_class.from_codecs(metadata.codecs)

# Use the regular chunk shape if available, otherwise use a
# placeholder. The ChunkTransform is shape-agnostic — the actual
Expand All @@ -285,7 +293,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
return pipeline.evolve_from_array_spec(chunk_spec)
elif isinstance(metadata, ArrayV2Metadata):
v2_codec = V2Codec(filters=metadata.filters, compressor=metadata.compressor)
pipeline = get_pipeline_class().from_codecs([v2_codec])
pipeline = pipeline_class.from_codecs([v2_codec])
chunk_spec = ArraySpec(
shape=metadata.chunks,
dtype=metadata.dtype,
Expand Down Expand Up @@ -408,6 +416,7 @@ def __init__(
metadata: ArrayV2Metadata | ArrayV2MetadataDict,
store_path: StorePath,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> None: ...

@overload
Expand All @@ -416,15 +425,17 @@ def __init__(
metadata: ArrayV3Metadata | ArrayMetadataJSON_V3,
store_path: StorePath,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> None: ...

def __init__(
self,
metadata: ArrayMetadata | ArrayMetadataDict,
store_path: StorePath,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> None:
object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config))
object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config, context))
# An AsyncArray is the resolving layer: build its codec pipeline now, so
# pipeline-construction warnings fire at creation, as they always have.
_ = self._io._resolved_codec_pipeline
Expand Down Expand Up @@ -4607,21 +4618,25 @@ def build_array_metadata(
attributes: dict[str, JSON] | None = None,
chunk_key_encoding: ChunkKeyEncodingLike | None = None,
dimension_names: DimensionNamesLike = None,
context: Context | None = None,
) -> ArrayV2Metadata | ArrayV3Metadata:
"""Build an array metadata document from creation-style keyword arguments.

This is the pure half of `init_array`: it normalizes the same keyword
arguments into an in-memory metadata object without touching any store.
The parameters have the same meaning as in
[`zarr.api.asynchronous.create_array`][].
[`zarr.api.asynchronous.create_array`][]. With a `context`, the data type
and Zarr format 3 codec names resolve through it instead of the default
registries; Zarr format 2 filters and compressors are numcodecs objects,
which no context governs.
"""

if zarr_format is None:
zarr_format = _default_zarr_format()

from zarr.codecs.sharding import ShardingCodec

zdtype = parse_dtype(dtype, zarr_format=zarr_format)
zdtype = parse_dtype(dtype, zarr_format=zarr_format, context=context)
shape_parsed = parse_shapelike(shape)
item_size = zdtype.item_size if isinstance(zdtype, HasItemSize) else 1
chunk_key_encoding_parsed = _parse_chunk_key_encoding(
Expand Down Expand Up @@ -4697,6 +4712,7 @@ def build_array_metadata(
filters=filters,
serializer=serializer,
dtype=zdtype,
context=context,
)
sub_codecs = cast("tuple[Codec, ...]", (*array_array, array_bytes, *bytes_bytes))
# A stored chunk grid passed as chunks= / shards= (e.g. by the "keep"
Expand Down Expand Up @@ -5320,6 +5336,7 @@ def _parse_chunk_encoding_v3(
filters: FiltersLike,
serializer: SerializerLike,
dtype: ZDType[TBaseDType, TBaseScalar],
context: Context | None = None,
) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]:
"""
Generate chunk encoding classes for v3 arrays with optional defaults.
Expand All @@ -5336,15 +5353,17 @@ def _parse_chunk_encoding_v3(
maybe_array_array = (filters,)
else:
maybe_array_array = cast("Iterable[Codec | dict[str, JSON]]", filters)
out_array_array = tuple(_parse_array_array_codec(c) for c in maybe_array_array)
out_array_array = tuple(
_parse_array_array_codec(c, context=context) for c in maybe_array_array
)

if serializer == "auto":
out_array_bytes = default_serializer_v3(dtype)
else:
# TODO: ensure that the serializer is compatible with the ndarray produced by the
# array-array codecs. For example, if a sequence of array-array codecs produces an
# array with a single-byte data type, then the serializer should not specify endiannesss.
out_array_bytes = _parse_array_bytes_codec(serializer)
out_array_bytes = _parse_array_bytes_codec(serializer, context=context)

if compressors is None:
out_bytes_bytes: tuple[BytesBytesCodec, ...] = ()
Expand All @@ -5357,7 +5376,9 @@ def _parse_chunk_encoding_v3(
else:
maybe_bytes_bytes = cast("Iterable[Codec | dict[str, JSON]]", compressors)

out_bytes_bytes = tuple(_parse_bytes_bytes_codec(c) for c in maybe_bytes_bytes)
out_bytes_bytes = tuple(
_parse_bytes_bytes_codec(c, context=context) for c in maybe_bytes_bytes
)

# TODO: ensure that the serializer is compatible with the ndarray produced by the
# array-array codecs. For example, if a sequence of array-array codecs produces an
Expand Down Expand Up @@ -6254,10 +6275,18 @@ class AsyncArrayIO:
opened from a stored document that needed a repair stores that repair
before its first chunk write. Call `store_metadata` first to make that
write explicit.

A `context` (see [`zarr.core.context.Context`][]) defines the data types,
codecs and codec pipeline in scope for this array. It is consumed where
the document resolves into classes — building the data path — so the
document operations never touch it: a name outside the context fails only
when chunks are read or written, even when the default registries know it.
Without one, the default registries are used.
"""

_store_path: StorePath
_config: ArrayConfig
_context: Context | None
_source: ArrayDocumentModel | ArrayMetadata
_document_cache: ArrayDocumentModel | None
_resolved_cache: ArrayMetadata | None
Expand All @@ -6270,9 +6299,11 @@ def __init__(
document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict,
store_path: StorePath,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> None:
self._store_path = store_path
self._config = parse_array_config(config)
self._context = context
self._document_cache = None
self._resolved_cache = None
self._chunk_grid_cache = None
Expand All @@ -6284,7 +6315,7 @@ def __init__(
# A resolved metadata object passes through; a raw document dict is
# read eagerly, as repaired when it needs that, warning as `AsyncArray`
# always has.
self._source = parse_array_metadata(document, str(store_path))
self._source = parse_array_metadata(document, str(store_path), context=context)

@classmethod
async def build(
Expand All @@ -6306,12 +6337,15 @@ async def build(
chunk_key_encoding: ChunkKeyEncodingLike | None = None,
dimension_names: DimensionNamesLike = None,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> AsyncArrayIO:
"""Declare an array bound to a location in a store, without storing anything.

The keyword arguments have the same meaning as in
[`zarr.api.asynchronous.create_array`][]. The returned instance holds
the normalized metadata in memory; `store_metadata` stores it.
the normalized metadata in memory; `store_metadata` stores it. With a
`context`, the data type and codec names resolve through it instead of
the default registries.
"""
metadata = build_array_metadata(
shape=shape,
Expand All @@ -6327,9 +6361,10 @@ async def build(
attributes=attributes,
chunk_key_encoding=chunk_key_encoding,
dimension_names=dimension_names,
context=context,
)
store_path = await make_store_path(store, path=name)
return cls(metadata, store_path, config)
return cls(metadata, store_path, config, context=context)

@classmethod
async def open(
Expand All @@ -6339,6 +6374,7 @@ async def open(
name: str | None = None,
zarr_format: ZarrFormat | None = None,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> AsyncArrayIO:
"""Read the metadata document of an existing array and bind to it.

Expand All @@ -6354,7 +6390,7 @@ async def open(
"""
store_path = await make_store_path(store, path=name)
data = await get_array_metadata(store_path, zarr_format=zarr_format)
return cls(_array_document_model(data), store_path, config)
return cls(_array_document_model(data), store_path, config, context=context)

@property
def document(self) -> ArrayDocumentModel:
Expand Down Expand Up @@ -6382,7 +6418,9 @@ def metadata(self) -> ArrayMetadata:
data = _document_json(source)
if isinstance(source, ZarrV2ArrayMetadata):
data["attributes"] = dict(source.attributes)
self._resolved_cache = parse_array_metadata(data, str(self._store_path))
self._resolved_cache = parse_array_metadata(
data, str(self._store_path), context=self._context
)
return self._resolved_cache

@property
Expand Down Expand Up @@ -6495,7 +6533,7 @@ def _shard_grid_shape(self) -> tuple[int, ...]:
def _resolved_codec_pipeline(self) -> CodecPipeline:
if self._codec_pipeline_cache is None:
self._codec_pipeline_cache = create_codec_pipeline(
metadata=self.metadata, store=self._store_path.store
metadata=self.metadata, store=self._store_path.store, context=self._context
)
return self._codec_pipeline_cache

Expand Down Expand Up @@ -6714,8 +6752,9 @@ def __init__(
document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict,
store_path: StorePath,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> None:
self._async_io = AsyncArrayIO(document, store_path, config)
self._async_io = AsyncArrayIO(document, store_path, config, context=context)

@classmethod
def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO:
Expand Down Expand Up @@ -6743,6 +6782,7 @@ def build(
chunk_key_encoding: ChunkKeyEncodingLike | None = None,
dimension_names: DimensionNamesLike = None,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> ArrayIO:
"""Declare an array bound to a location in a store, without storing anything.

Expand All @@ -6767,6 +6807,7 @@ def build(
chunk_key_encoding=chunk_key_encoding,
dimension_names=dimension_names,
config=config,
context=context,
)
)
)
Expand All @@ -6779,13 +6820,18 @@ def open(
name: str | None = None,
zarr_format: ZarrFormat | None = None,
config: ArrayConfigLike | None = None,
context: Context | None = None,
) -> ArrayIO:
"""Read the metadata document of an existing array and bind to it.

See `AsyncArrayIO.open`.
"""
return cls._wrap(
sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config))
sync(
AsyncArrayIO.open(
store, name=name, zarr_format=zarr_format, config=config, context=context
)
)
)

@property
Expand Down
Loading
Loading