consolidate_metadata does not scale well

[bluppfisk]

Zarr version

2.18.7

Numcodecs version

0.15.1

Python Version

3.12

Operating System

Linux

Installation

pip

Description

consolidate_metadata() iterates over all keys in the store, only to access the ones that end in .zattrs, .zarray and .zgroup. For zarr stores that contain multiple millions of chunk files, this will start to become slow as the zarr grows, no matter how fast os.scandir is.

Perhaps it is my bad for storing additional metadata in the metadata files (although I am just following xarray's example of storing _ARRAY_DIMENSIONS etc), but it feels this could be made a lot faster by targeting the files directly (after all they are known to exist in each variable).

I'd be happy to submit a PR for this.

[TomAugspurger]

Thanks for the report. Can you maybe do some profiling to confirm exactly where the slowdown is for your usecase? That's very plausible, but it'd be good to confirm.

it feels this could be made a lot faster by targeting the files directly (after all they are known to exist in each variable).

Assuming it is listing nodes that's taking time, could you suggest a potential fix? I wonder if we'd need an additional method in the Store interface that avoids listing entire directories (and now that I think about it: are we ever listing chunk files? I don't see why we would need to do that... We should only need to recurse deeper when we see a Group node. We shouldn't ever need to list files under an Array node).

One more note: this isn't unique to consolidated metadata. IIUC, this should affect group.members(), which is used by the metadata consolidation methods.

[bluppfisk]

hi Tom,

I should be more precise and say that I am using zarr 2.18 and zarr 2 format.

in consolidate_metadata I see the following lines:

    out = {
        "zarr_consolidated_format": 1,
        "metadata": {key: json_loads(store[key]) for key in store if is_zarr_key(key)},
    }

the DirectoryStore's __iter__ method invokes - after jumping through some hoops - the following code from storage.py::DirectoryStore

    @staticmethod
    def _keys_fast(path, walker=os.walk):
        for dirpath, _, filenames in walker(path):
            dirpath = os.path.relpath(dirpath, path)
            if dirpath == os.curdir:
                for f in filenames:
                    yield f
            else:
                dirpath = dirpath.replace("\\", "/")
                for f in filenames:
                    yield "/".join((dirpath, f))

This means it will iterate over every file in the variable path's folder. If there are millions or billions (as is the case for me) of files there, this normally will take O(n) as it must iterate over every file, recursively, only to discard those that do not match is_zarr_key, which will only ever be true for 3 files (per variable).

It would seem faster to list those files:

for var_name in variables:
   path = base_path / var_name
   for metadata_key in [".zattrs", ".zgroup", ".zarray"]:
      metadata[f"{path}/{metadata_key}"] = json.loads(f"{path}/{metadata_key}")

Perhaps I'm missing something.

[TomAugspurger]

Ah, OK. I don't think anyone is actively working on zarr-python 2.x at the moment.

[bluppfisk]

understood - unfortunately we have an existing project on zarr 2 and since it modifies directorystore to allow rolling (removing the oldest chunks so it doesn't grow infinitely), we're not ready to upgrade to 3 yet. Would PRs be accepted for zarr 2 behaviour?

[d-v-b]

Would PRs be accepted for zarr 2 behaviour?

I don't think we have any plans for performance-related releases of v2, only bug / security fixes.

[bluppfisk]

Fair enough. Just in case someone comes across this issue, I'm happy to share our implementation:

import os

from zarr.convenience import open_consolidated
from zarr.storage import BaseStore, normalize_store_arg
from zarr.util import json_dumps, json_loads


def consolidate_metadata(store: BaseStore, metadata_key=".zmetadata", *, path=""):
    """
    Consolidate all metadata for groups and arrays within the given store
    into a single resource and put it under the given key.

    Parameters
    ----------
    store : MutableMapping or string
        Store or path to directory in file system or name of zip file.
    metadata_key : str
        Key to put the consolidated metadata under.
    path : str or None
        Path corresponding to the group that is being consolidated. Not required
        for zarr v2 stores.

    Returns
    -------
    g : :class:`zarr.hierarchy.Group`
        Group instance, opened with the new consolidated metadata.

    See Also
    --------
    open_consolidated

    """
    store = normalize_store_arg(store, mode="w")

    # iterate over every true directory in the store.path root
    subdirs = [
        f"{entry.name}/"
        for entry in os.scandir(store.path)
        if entry.is_dir(follow_symlinks=False)
    ]
    subdirs += [""]
    loose_keys = [
        f"{subdir}{key}"
        for key in [".zattrs", ".zarray", ".zgroup"]
        for subdir in subdirs
    ]

    out = {
        "zarr_consolidated_format": 1,
        "metadata": {key: json_loads(store[key]) for key in loose_keys if key in store},
    }
    store[metadata_key] = json_dumps(out)

    return open_consolidated(store, metadata_key=metadata_key, path=path)

[d-v-b]

thanks for the contribution here @bluppfisk, I'm going to close this as unplanned, since we are not touching the 2.x codebase. Feel free to re-open a new issue if you find performance problems / opportunities with the 3.x implementation of consolidated metadata.