Zarr-v3 Consolidated Metadata

[TomAugspurger]

Implements the optional Consolidated Metadata feature of zarr-v3: zarr-developers/zarr-specs#309, along with reading zarr v2 metadata.

The implementation defines a new dataclass: ConsoliatedMetadata. It's an optional field on the existing GroupMetadata object. Opening as a draft until the PR to the zarr-specs repo finishes up.

closes #1161

TODO:

• Add unit tests and/or doctests in docstrings
• Add docstrings and API docs for any new/modified user-facing classes and functions
• New/modified features documented in docs/tutorial.rst
• Changes documented in docs/release.rst
• GitHub Actions have all passed
• Test coverage is 100% (Codecov passes)

[TomAugspurger]

This isn't handling nested groups correctly yet. My understanding (and I'll clarify this in the spec) is that given a structure like

/root
  /arr-1
  /arr-2
  /child-group
    /child-arr-1
    /child-arr-2

we'd expect that the metadata from child-group, child-arr-1, and child-arr-2 should end up in the consolidated metadata.

[jhamman]

Good start here Tom! I have a few questions about parts you haven't implemented yet that should help explore the design a bit more:

1. How are you thinking about wiring this into read operations (e.g. getitem, listdir, etc.)?
2. How should we think about consistency w.r.t. the consolidated metadata field? When should we invalidate the consolidated metadata and when is it safe to keep using it?
3. How can this approach be used to support reading/writing v2 consolidated metadata (.zmetadata)?

[TomAugspurger]

Quick update here:

• AsyncGroup now uses consolidated metadata in some operations (primarily getitem. I'm going through some more spots now). This means you can do AsyncGroup.getitem(key) to get a child node without any additional I/O.
• fc901eb added support for reading zarr V2 consolidated metadata. Needs some cleanup and probably more testing, but the basics seem to work.

[TomAugspurger]

As I start to use this a bit, I'm rethinking the in-memory representation of consolidated metadata for a nested hierarchy. Specifically the consolidated_metadata.metadata dictionary which maps keys to ArrayMetadata | GroupMetadata objects. Our options are:

1. A flat structure, where the keys include all path segments from the root group and the root ConsolidatedMetadata.metadata dict is equal to the length of all child nodes (not just immediate children). This matches the representation on disk.
2. A nested structure, where the keys include just the name (so not the leading segments of the path). The root ConsolidatedMetadata.metadata dict holds just immediate children nodes. Metadata for nested groups can still be accessed, but through the .consolidated_metadata on the children. This matches the logical representation.

The motivation for this rethink is from the current implementation having to be careful about just using consolidated_metdata.metadata directly. If you want to propagate consolidated metadata, you need to use the Group.getitem method. See Group.members where this becomes relevant.

Here's an example:

Given a hierarchy like

root/
  g0/
   c0/
    array0
    array1
   c1/
     array0
     array1
  g1/
    c0/
      array0
      array1
    c1/
      array0
      array1

We'll represent the consolidated metadata as a flat mapping of (store) keys to values.

"g0": {"attributes": ..., "node_type": "group"},
"g1": {"attributes": ..., "node_type": "group"},
"g0/c0": {"attributes": ..., "node_type": "group"},
...
"g0/c0/array0": {"shape": [...], "node_type": "array"},
...
"g1/c1/array1": {"shape": [...], "node_type": "array"}

But in memory, what is the Group.metadata.consolidated_metadata for each of theses groups?

Should it match the flat structure on disk, where the keys imply the structure?

# g0
GroupMetadata(
    attributes=...,
    consolidated_metadata=ConsolidatedMetadata(
	    metadata={
		    "c0": GroupMetadata(...),
		    "c0/array0": ArrayMetadata(...),
		    ...,
		    "c1/array1": ArrayMetadata(...),
		},
)

Or should it have a nested / tree-like structure, where just immediate children appear in group.metadata.consolidated_metadata.metadata, and nested members can be accessed through that dictionary?

# g0
GroupMetadata(
    attributes=...,
    consolidated_metadata=ConsolidatedMetadata(
	    metadata={
		    "c0": GroupMetadata(
			    attributes=...,
			    consoliated_metadata=ConsolidatedMetadata(
				    metadata={
					    "array0": ArrayMetadata(...),
					    "array1": ArrayMetadata(...),
				    }
			    )
			),
			"c1": GroupMetadata(
				attributes=...,
				consolidated_metadata=ConsolidatedMetadata(
					metadata={
					    "array0": ArrayMetadata(...),
					    "array1": ArrayMetadata(...),
					}
				)
			)
		},
)

Right now I've implemented option 1. I'll out option 2 today.

[d-v-b]

do we have to pick just 1 in-memory representation? Over in pydantic-zarr I do something similar to metadata consolidation, and I use a tree representation or a flat dict[str, Array | Group] representation situationally (see the to_flat and from_flat functions).

[TomAugspurger]

There's a behavior change here that's causing some xarray tests to fail. It comes down to whether we interpret zarr.open_consolidated(store, path="path/to/group") as either:

1. Open the consolidated metadata at the root of the store (.zmetadata) and then select the (nested) group path/to/group. Or,
2. Open the consolidated metadata at "path/to/group/"

Zarr 2.x implemented the first one. This branch currently implements the second version.

The rough flow is

1. create a Group / arrays at some level down in the store: e.g. a/b/group with arrays a/b/group/x, a/b/group/y, ...
2. Consolidate the store at the root level with zarr.consolidate_metadata(store).
3. open the (nested) group with zarr.open_consolidated(store=store, path="a/b/group")

I'll see how hard it is to implement what 2.x was doing, but I want to confirm that's the behavior we want.

A limitation of 2.x's way of doing things is that you can only have one consolidated metadata for the entire Store: at the root.

[jhamman]

@TomAugspurger -- this came together so nicely! Way to go 😄

I left a few comments of substance but I think this can go in very soon (tomorrow?)

[jhamman]

Suggested change

	zarr-python implements the `Consolidated Metadata_` extension to the Zarr Spec.
	Zarr-Python implements the `Consolidated Metadata_` extension to the Zarr Spec.

[jhamman]

Suggested change

	attribute of the Group.
	attribute of the Group metadata.

[jhamman]

These docs are great @TomAugspurger! 👏

[jhamman]

Suggested change

	async def open_consolidated(*args: Any, use_consolidated: bool = True, **kwargs: Any) -> AsyncGroup:
	"""
	Alias for :func:`open_group` with ``use_consolidated=True``.
	"""
	return await open_group(*args, use_consolidated=use_consolidated, **kwargs)
	async def open_consolidated(*args: Any, use_consolidated: bool = True, **kwargs: Any) -> AsyncGroup:
	"""
	Alias for :func:`open_group` with ``use_consolidated=True``.
	"""
	return await open_group(*args, use_consolidated=True, **kwargs)

Seems to me that open_consolidated should always use use_consolidated=True

[TomAugspurger]

My goal here is to avoid a user accidentally passing use_consolidated in **kwargs and us silently overwriting it. I'll update this to raise if use_consolidated isn't True.

[jhamman]

perhaps raise NodeTypeValidationError from #2310

[jhamman]

In the future, we may choose to inform the caller that the consolidated metadata does not match the root group metadata but for now, this seems fine.

[TomAugspurger]

Forgot to actually delete this block.

[jhamman]

Huge @TomAugspurger 🎉 🎉 🎉 !!!

[d-v-b]

awesome work!