N5FSStore

zarr/__init__.py

@@ -9,7 +9,7 @@
                            zeros_like)
 from zarr.errors import CopyError, MetadataError
 from zarr.hierarchy import Group, group, open_group
-from zarr.n5 import N5Store
+from zarr.n5 import N5Store, N5FSStore
 from zarr.storage import (ABSStore, DBMStore, DictStore, DirectoryStore,
                           LMDBStore, LRUStoreCache, MemoryStore, MongoDBStore,
                           NestedDirectoryStore, RedisStore, SQLiteStore,

zarr/n5.py

@@ -11,7 +11,8 @@
 from numcodecs.registry import get_codec, register_codec
 
 from .meta import ZARR_FORMAT, json_dumps, json_loads
-from .storage import NestedDirectoryStore, _prog_ckey, _prog_number
+from .storage import FSStore
+from .storage import NestedDirectoryStore, _prog_ckey, _prog_number, normalize_storage_path
 from .storage import array_meta_key as zarr_array_meta_key
 from .storage import attrs_key as zarr_attrs_key
 from .storage import group_meta_key as zarr_group_meta_key
@@ -281,12 +282,298 @@ def _contains_attrs(self, path):
         return len(attrs) > 0
 
 
+class N5FSStore(FSStore):
+    """Implentation of the N5 format (https://github.com/saalfeldlab/n5) using `fsspec`,
+    which allows storage on a variety of filesystems. Based on `zarr.N5Store`.
+    Parameters
+    ----------
+    path : string
+        Location of directory to use as the root of the storage hierarchy.
+    normalize_keys : bool, optional
+        If True, all store keys will be normalized to use lower case characters
+        (e.g. 'foo' and 'FOO' will be treated as equivalent). This can be
+        useful to avoid potential discrepancies between case-senstive and
+        case-insensitive file system. Default value is False.
+
+    Examples
+    --------
+    Store a single array::
+
+        >>> import zarr
+        >>> store = zarr.N5FSStore('data/array.n5', auto_mkdir=True)
+        >>> z = zarr.zeros((10, 10), chunks=(5, 5), store=store, overwrite=True)
+        >>> z[...] = 42
+
+    Store a group::
+
+        >>> store = zarr.N5FSStore('data/group.n5', auto_mkdir=True)
+        >>> root = zarr.group(store=store, overwrite=True)
+        >>> foo = root.create_group('foo')
+        >>> bar = foo.zeros('bar', shape=(10, 10), chunks=(5, 5))
+        >>> bar[...] = 42
+
+    Notes
+    -----
+    This is an experimental feature.
+    Safe to write in multiple threads or processes.
+
+    Be advised that the `_dimension_separator` property of this store
+    (and arrays it creates) is ".", but chunks saved by this store will
+    in fact be "/" separated, as proscribed by the N5 format.
+
+    This is counter-intuitive (to say the least), but not arbitrary.
+    Chunks in N5 format are stored with reversed dimension order
+    relative to Zarr chunks: a chunk of a 3D Zarr array would be stored
+    on a file system as `/0/1/2`, but in N5 the same chunk would be
+    stored as `/2/1/0`. Therefore, stores targeting N5 must intercept
+    chunk keys and flip the order of the dimensions before writing to
+    storage, and this procedure requires chunk keys with "." separated
+    dimensions, hence the Zarr arrays targeting N5 have the deceptive
+    "." dimension separator.
+    """
+    _array_meta_key = 'attributes.json'
+    _group_meta_key = 'attributes.json'
+    _attrs_key = 'attributes.json'
+
+    def __init__(self, *args, **kwargs):
+        if 'dimension_separator' in kwargs:
+            kwargs.pop('dimension_separator')
+            warnings.warn('Keyword argument `dimension_separator` will be ignored')
+        dimension_separator = "."
+        super().__init__(*args, dimension_separator=dimension_separator, **kwargs)
+
+    def _swap_separator(self, key):
+        segments = list(key.split('/'))
+        if segments:
+            last_segment = segments[-1]
+            if _prog_ckey.match(last_segment):
+                coords = list(last_segment.split('.'))
+                last_segment = '/'.join(coords[::-1])
+                segments = segments[:-1] + [last_segment]
+                key = '/'.join(segments)
+        return key
+
+    def _normalize_key(self, key):
+        if is_chunk_key(key):
+            key = invert_chunk_coords(key)
+
+        key = normalize_storage_path(key).lstrip("/")
+        if key:
+            *bits, end = key.split("/")
+
+            if end not in (self._array_meta_key, self._group_meta_key, self._attrs_key):
+                end = end.replace(".", "/")
+                key = "/".join(bits + [end])
+        return key.lower() if self.normalize_keys else key
+
+    def __getitem__(self, key):
+        if key.endswith(zarr_group_meta_key):
+
+            key = key.replace(zarr_group_meta_key, self._group_meta_key)
+            value = group_metadata_to_zarr(self._load_n5_attrs(key))
+
+            return json_dumps(value)
+
+        elif key.endswith(zarr_array_meta_key):
+
+            key = key.replace(zarr_array_meta_key, self._array_meta_key)
+            value = array_metadata_to_zarr(self._load_n5_attrs(key))
+
+            return json_dumps(value)
+
+        elif key.endswith(zarr_attrs_key):
+
+            key = key.replace(zarr_attrs_key, self._attrs_key)
+            value = attrs_to_zarr(self._load_n5_attrs(key))
+
+            if len(value) == 0:
+                raise KeyError(key)
+            else:
+                return json_dumps(value)
+
+        elif is_chunk_key(key):
+            key = self._swap_separator(key)
+
+        return super().__getitem__(key)
+
+    def __setitem__(self, key, value):
+        if key.endswith(zarr_group_meta_key):
+
+            key = key.replace(zarr_group_meta_key, self._group_meta_key)
+
+            n5_attrs = self._load_n5_attrs(key)
+            n5_attrs.update(**group_metadata_to_n5(json_loads(value)))
+
+            value = json_dumps(n5_attrs)
+
+        elif key.endswith(zarr_array_meta_key):
+
+            key = key.replace(zarr_array_meta_key, self._array_meta_key)
+
+            n5_attrs = self._load_n5_attrs(key)
+            n5_attrs.update(**array_metadata_to_n5(json_loads(value)))
+
+            value = json_dumps(n5_attrs)
+
+        elif key.endswith(zarr_attrs_key):
+
+            key = key.replace(zarr_attrs_key, self._attrs_key)
+
+            n5_attrs = self._load_n5_attrs(key)
+            zarr_attrs = json_loads(value)
+
+            for k in n5_keywords:
+                if k in zarr_attrs.keys():
+                    raise ValueError(
+                        "Can not set attribute %s, this is a reserved N5 keyword" % k
+                    )
+
+            # replace previous user attributes
+            for k in list(n5_attrs.keys()):
+                if k not in n5_keywords:
+                    del n5_attrs[k]
+
+            # add new user attributes
+            n5_attrs.update(**zarr_attrs)
+
+            value = json_dumps(n5_attrs)
+
+        elif is_chunk_key(key):
+            key = self._swap_separator(key)
+
+        super().__setitem__(key, value)
+
+    def __delitem__(self, key):
+
+        if key.endswith(zarr_group_meta_key):  # pragma: no cover
+            key = key.replace(zarr_group_meta_key, self._group_meta_key)
+        elif key.endswith(zarr_array_meta_key):  # pragma: no cover
+            key = key.replace(zarr_array_meta_key, self._array_meta_key)
+        elif key.endswith(zarr_attrs_key):  # pragma: no cover
+            key = key.replace(zarr_attrs_key, self._attrs_key)
+        elif is_chunk_key(key):
+            key = self._swap_separator(key)
+
+        super().__delitem__(key)
+
+    def __contains__(self, key):
+        if key.endswith(zarr_group_meta_key):
+
+            key = key.replace(zarr_group_meta_key, self._group_meta_key)
+            if key not in self:
+                return False
+            # group if not a dataset (attributes do not contain 'dimensions')
+            return "dimensions" not in self._load_n5_attrs(key)
+
+        elif key.endswith(zarr_array_meta_key):
+
+            key = key.replace(zarr_array_meta_key, self._array_meta_key)
+            # array if attributes contain 'dimensions'
+            return "dimensions" in self._load_n5_attrs(key)
+
+        elif key.endswith(zarr_attrs_key):
+
+            key = key.replace(zarr_attrs_key, self._attrs_key)
+            return self._contains_attrs(key)
+
+        elif is_chunk_key(key):
+            key = self._swap_separator(key)
+
+        return super().__contains__(key)
+
+    def __eq__(self, other):
+        return isinstance(other, N5FSStore) and self.path == other.path
+
+    def listdir(self, path=None):
+        if path is not None:
+            path = invert_chunk_coords(path)
+
+        # We can't use NestedDirectoryStore's listdir, as it requires
+        # array_meta_key to be present in array directories, which this store
+        # doesn't provide.
+        children = super().listdir(path=path)
+        if self._is_array(path):
+
+            # replace n5 attribute file with respective zarr attribute files
+            children.remove(self._array_meta_key)
+            children.append(zarr_array_meta_key)
+            if self._contains_attrs(path):
+                children.append(zarr_attrs_key)
+
+            # special handling of directories containing an array to map
+            # inverted nested chunk keys back to standard chunk keys
+            new_children = []
+            root_path = self.dir_path(path)
+            for entry in children:
+                entry_path = os.path.join(root_path, entry)
+                if _prog_number.match(entry) and self.fs.isdir(entry_path):
+                    for file_name in self.fs.find(entry_path):
+                        file_path = os.path.join(root_path, file_name)
+                        rel_path = file_path.split(root_path)[1]
+                        new_child = rel_path.lstrip('/').replace('/', ".")
+                        new_children.append(invert_chunk_coords(new_child))
+                else:
+                    new_children.append(entry)
+            return sorted(new_children)
+
+        elif self._is_group(path):
+
+            # replace n5 attribute file with respective zarr attribute files
+            children.remove(self._group_meta_key)
+            children.append(zarr_group_meta_key)
+            if self._contains_attrs(path):  # pragma: no cover
+                children.append(zarr_attrs_key)
+            return sorted(children)
+        else:
+            return children
+
+    def _load_n5_attrs(self, path):
+        try:
+            s = super().__getitem__(path)
+            return json_loads(s)
+        except KeyError:
+            return {}
+
+    def _is_group(self, path):
+
+        if path is None:
+            attrs_key = self._attrs_key
+        else:
+            attrs_key = os.path.join(path, self._attrs_key)
+
+        n5_attrs = self._load_n5_attrs(attrs_key)
+        return len(n5_attrs) > 0 and "dimensions" not in n5_attrs
+
+    def _is_array(self, path):
+
+        if path is None:
+            attrs_key = self._attrs_key
+        else:
+            attrs_key = os.path.join(path, self._attrs_key)
+
+        return "dimensions" in self._load_n5_attrs(attrs_key)
+
+    def _contains_attrs(self, path):
+
+        if path is None:
+            attrs_key = self._attrs_key
+        else:
+            if not path.endswith(self._attrs_key):
+                attrs_key = os.path.join(path, self._attrs_key)
+            else:  # pragma: no cover
+                attrs_key = path
+
+        attrs = attrs_to_zarr(self._load_n5_attrs(attrs_key))
+        return len(attrs) > 0
+
+
 def is_chunk_key(key):
+    rv = False
     segments = list(key.split('/'))
     if segments:
         last_segment = segments[-1]
-        return _prog_ckey.match(last_segment)
-    return False  # pragma: no cover
+        rv = _prog_ckey.match(last_segment)
+    return rv
 
 
 def invert_chunk_coords(key):
@@ -373,6 +660,7 @@ def array_metadata_to_zarr(array_metadata):
     array_metadata['fill_value'] = 0  # also if None was requested
     array_metadata['order'] = 'C'
     array_metadata['filters'] = []
+    array_metadata['dimension_separator'] = '.'
 
     compressor_config = array_metadata['compressor']
     compressor_config = compressor_config_to_zarr(compressor_config)

zarr/storage.py

@@ -1065,22 +1065,28 @@ class FSStore(MutableMapping):
         Separator placed between the dimensions of a chunk.
     storage_options : passed to the fsspec implementation
     """
+    _array_meta_key = array_meta_key
+    _group_meta_key = group_meta_key
+    _attrs_key = attrs_key
 
-    _META_KEYS = (attrs_key, group_meta_key, array_meta_key)
-
-    def __init__(self, url, normalize_keys=False, key_separator=None,
+    def __init__(self, url, normalize_keys=True, key_separator=None,
                  mode='w',
                  exceptions=(KeyError, PermissionError, IOError),
                  dimension_separator=None,
                  **storage_options):
         import fsspec
         self.normalize_keys = normalize_keys
+
+        protocol, _ = fsspec.core.split_protocol(url)
+        # set auto_mkdir to True for local file system
+        if protocol in (None, "file") and not storage_options.get("auto_mkdir"):
+            storage_options["auto_mkdir"] = True
+
         self.map = fsspec.get_mapper(url, **storage_options)
         self.fs = self.map.fs  # for direct operations
         self.path = self.fs._strip_protocol(url)
         self.mode = mode
         self.exceptions = exceptions
-
         # For backwards compatibility. Guaranteed to be non-None
         if key_separator is not None:
             dimension_separator = key_separator
@@ -1091,7 +1097,6 @@ def __init__(self, url, normalize_keys=False, key_separator=None,
 
         # Pass attributes to array creation
         self._dimension_separator = dimension_separator
-
         if self.fs.exists(self.path) and not self.fs.isdir(self.path):
             raise FSPathExistNotDir(url)
 
@@ -1100,7 +1105,7 @@ def _normalize_key(self, key):
         if key:
             *bits, end = key.split('/')
 
-            if end not in FSStore._META_KEYS:
+            if end not in (self._array_meta_key, self._group_meta_key, self._attrs_key):
                 end = end.replace('.', self.key_separator)
                 key = '/'.join(bits + [end])
 
@@ -1178,7 +1183,7 @@ def listdir(self, path=None):
             if self.key_separator != "/":
                 return children
             else:
-                if array_meta_key in children:
+                if self._array_meta_key in children:
                     # special handling of directories containing an array to map nested chunk
                     # keys back to standard chunk keys
                     new_children = []