Support for Zarr v3

Author: Altay Sansal

src/mdio/api/accessor.py

@@ -186,13 +186,13 @@ def _validate_store(self, storage_options):
         if storage_options is None:
             storage_options = {}
 
-        self.store = process_url(
+        self.url = process_url(
             url=self.url,
-            mode=self.mode,
-            storage_options=storage_options,
-            memory_cache_size=self._memory_cache_size,
             disk_cache=self._disk_cache,
         )
+        self.store = zarr.open(
+            self.url, mode=self.mode, storage_options=storage_options
+        ).store
 
     def _connect(self):
         """Open the zarr root."""
@@ -375,12 +375,12 @@ def stats(self, value: dict) -> None:
     @property
     def _metadata_group(self) -> zarr.Group:
         """Get metadata zarr.group handle."""
-        return self.root.metadata
+        return self.root["metadata"]
 
     @property
     def _data_group(self) -> zarr.Group:
         """Get data zarr.Group handle."""
-        return self.root.data
+        return self.root["data"]
 
     def __getitem__(self, item: int | tuple) -> npt.ArrayLike | da.Array | tuple:
         """Data getter."""

src/mdio/api/convenience.py

@@ -5,8 +5,8 @@
 from typing import TYPE_CHECKING
 
 import zarr
+from numcodecs import Blosc
 from tqdm.auto import tqdm
-from zarr import Blosc
 
 from mdio.api.io_utils import process_url
 from mdio.core.indexing import ChunkIterator
@@ -21,6 +21,7 @@
     from mdio import MDIOReader
 
 
+# TODO(Altay): This is not implemented in Zarr v3 yet.
 def copy_mdio(  # noqa: PLR0913
     source: MDIOReader,
     dest_path_or_buffer: str,
@@ -242,6 +243,7 @@ def rechunk_batch(
     write_rechunked_values(source, suffix_list, *plan)
 
 
+# TODO(Altay): This needs to be validated
 def rechunk(
     source: MDIOAccessor,
     chunks: tuple[int, ...],

src/mdio/api/io_utils.py

@@ -2,40 +2,28 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 import dask.array as da
 import zarr
-from zarr.storage import FSStore
+from zarr.storage._utils import normalize_path
 
 
 def process_url(
     url: str,
-    mode: str,
-    storage_options: dict[str, Any],
-    memory_cache_size: int,
     disk_cache: bool,
-) -> FSStore:
+) -> str:
     """Check read/write access to FSStore target and return FSStore with double caching.
 
-    It can use an in-memory Least Recently Used (LRU) cache implementation from
-    Zarr, and optionally, a file cache (`simplecache` protocol from FSSpec) that
-    is useful for remote stores.
-
-    File cache is only valid for remote stores. The LRU caching works
-    on both remote and local.
+    It can optionally use a file cache (`simplecache` protocol from fsspec) that
+    is useful for remote stores. File cache is only useful for remote stores.
 
     The `storage_options` argument represents a set of parameters to be passed
-    to the FSSpec backend. Note that the format of `storage_options` is
+    to the fsspec backend. Note that the format of `storage_options` is
     different if `disk_cache` is enabled or disabled, since `disk_cache`
     interanlly uses the simplecache protocol.
 
     Args:
-        url: FSSpec compliant url
-        mode: Toggle for overwriting existing store
-        storage_options: Storage options for the storage backend.
-        memory_cache_size: Maximum in memory LRU cache size in bytes.
-        disk_cache: This enables FSSpec's `simplecache` if True.
+        url: fsspec compliant url
+        disk_cache: This enables fsspec's `simplecache` if True.
 
     Returns:
         Store with augmentations like cache, write verification etc.
@@ -52,7 +40,6 @@ def process_url(
         ...     url="s3://bucket/key",
         ...     mode="r",
         ...     storage_options={"key": "my_key", "secret": "my_secret"},
-        ...     memory_cache_size=0,
         ...     disk_cache=False,
         ... )
 
@@ -64,7 +51,6 @@ def process_url(
         ...     url="s3://bucket/key",
         ...     mode="r",
         ...     storage_options={"s3": {"key": "my_key", "secret": "my_secret"}},
-        ...     memory_cache_size=0,
         ...     disk_cache=True,
         ... )
 
@@ -77,35 +63,13 @@ def process_url(
         ...         "s3": {"key": "my_key", "secret": "my_secret"},
         ...         "simplecache": {"cache_storage": "custom/local/cache/path"},
         ...     },
-        ...     memory_cache_size=0,
         ...     disk_cache=True,
         ... )
     """
     if disk_cache is True:
         url = "::".join(["simplecache", url])
 
-    # Strip whitespaces and slashes from end of string
-    url = url.rstrip("/ ")
-
-    # Flag for checking write access
-    check = True if mode == "w" else False
-
-    # TODO: Turning off write checking now because zarr has a bug.
-    #  Get rid of this once bug is fixed.
-    check = False
-
-    store = FSStore(
-        url=url,
-        check=check,
-        create=check,
-        mode=mode,
-        **storage_options,
-    )
-
-    if memory_cache_size != 0:
-        store = zarr.storage.LRUStoreCache(store=store, max_size=memory_cache_size)
-
-    return store
+    return normalize_path(url)
 
 
 def open_zarr_array(group_handle: zarr.Group, name: str) -> zarr.Array:

src/mdio/converters/segy.py

@@ -343,6 +343,8 @@ def segy_to_mdio(  # noqa: C901
         ...     grid_overrides={"HasDuplicates": True},
         ... )
     """
+    zarr.config.set({"default_zarr_format": 2, "write_empty_chunks": False})
+
     if index_names is None:
         index_names = [f"dim_{i}" for i in range(len(index_bytes))]
 
@@ -364,13 +366,8 @@ def segy_to_mdio(  # noqa: C901
     if storage_options_output is None:
         storage_options_output = {}
 
-    store = process_url(
-        url=mdio_path_or_buffer,
-        mode="w",
-        storage_options=storage_options_output,
-        memory_cache_size=0,  # Making sure disk caching is disabled,
-        disk_cache=False,  # Making sure disk caching is disabled
-    )
+    url = process_url(url=mdio_path_or_buffer, disk_cache=False)
+    root_group = zarr.open_group(url, mode="w", storage_options=storage_options_output)
 
     # Open SEG-Y with MDIO's SegySpec. Endianness will be inferred.
     mdio_spec = mdio_segy_spec()
@@ -406,42 +403,43 @@ def segy_to_mdio(  # noqa: C901
         logger.warning(f"Ingestion grid shape: {grid.shape}.")
         raise GridTraceCountError(np.sum(grid.live_mask), num_traces)
 
-    zarr_root = create_zarr_hierarchy(
-        store=store,
+    root_group = create_zarr_hierarchy(
+        root_group=root_group,
         overwrite=overwrite,
     )
 
     # Get UTC time, then add local timezone information offset.
     iso_datetime = datetime.now(timezone.utc).isoformat()
 
-    write_attribute(name="created", zarr_group=zarr_root, attribute=iso_datetime)
-    write_attribute(name="api_version", zarr_group=zarr_root, attribute=API_VERSION)
+    write_attribute(name="created", zarr_group=root_group, attribute=iso_datetime)
+    write_attribute(name="api_version", zarr_group=root_group, attribute=API_VERSION)
 
     dimensions_dict = [dim.to_dict() for dim in dimensions]
-    write_attribute(name="dimension", zarr_group=zarr_root, attribute=dimensions_dict)
+    write_attribute(name="dimension", zarr_group=root_group, attribute=dimensions_dict)
 
     # Write trace count
     trace_count = np.count_nonzero(grid.live_mask)
-    write_attribute(name="trace_count", zarr_group=zarr_root, attribute=trace_count)
+    write_attribute(name="trace_count", zarr_group=root_group, attribute=trace_count)
 
     # Note, live mask is not chunked since it's bool and small.
-    zarr_root["metadata"].create_dataset(
-        data=grid.live_mask,
+    live_mask_arr = root_group["metadata"].create_array(
         name="live_mask",
         shape=grid.shape[:-1],
-        chunks=-1,
-        dimension_separator="/",
+        chunks=grid.shape[:-1],
+        dtype="bool",
+        chunk_key_encoding={"name": "v2", "separator": "/"},
     )
+    live_mask_arr[...] = grid.live_mask[...]
 
     write_attribute(
         name="text_header",
-        zarr_group=zarr_root["metadata"],
+        zarr_group=root_group["metadata"],
         attribute=text_header.split("\n"),
     )
 
     write_attribute(
         name="binary_header",
-        zarr_group=zarr_root["metadata"],
+        zarr_group=root_group["metadata"],
         attribute=binary_header.to_dict(),
     )
 
@@ -470,8 +468,8 @@ def segy_to_mdio(  # noqa: C901
     stats = blocked_io.to_zarr(
         segy_file=segy,
         grid=grid,
-        data_root=zarr_root["data"],
-        metadata_root=zarr_root["metadata"],
+        data_root=root_group["data"],
+        metadata_root=root_group["metadata"],
         name="_".join(["chunked", suffix]),
         dtype="float32",
         chunks=chunksize,
@@ -480,17 +478,7 @@ def segy_to_mdio(  # noqa: C901
     )
 
     for key, value in stats.items():
-        write_attribute(name=key, zarr_group=zarr_root, attribute=value)
-
-    # Non-cached store for consolidating metadata.
-    # If caching is enabled the metadata may fall out of cache hence
-    # creating an incomplete `.zmetadata` file.
-    store_nocache = process_url(
-        url=mdio_path_or_buffer,
-        mode="r+",
-        storage_options=storage_options_output,
-        memory_cache_size=0,  # Making sure disk caching is disabled,
-        disk_cache=False,  # Making sure disk caching is disabled
-    )
+        write_attribute(name=key, zarr_group=root_group, attribute=value)
 
-    zarr.consolidate_metadata(store_nocache)
+    # Finalize Zarr for fast open
+    zarr.consolidate_metadata(root_group.store)

src/mdio/segy/blocked_io.py

@@ -11,9 +11,9 @@
 import numpy as np
 from dask.array import Array
 from dask.array import map_blocks
+from numcodecs import Blosc
 from psutil import cpu_count
 from tqdm.auto import tqdm
-from zarr import Blosc
 from zarr import Group
 
 from mdio.core import Grid
@@ -90,26 +90,24 @@ def to_zarr(
             "install 'zfpy' or install mdio with `--extras lossy`"
         )
 
-    trace_array = data_root.create_dataset(
+    trace_array = data_root.create_array(
         name=name,
         shape=grid.shape,
         compressor=trace_compressor,
         chunks=chunks,
-        dimension_separator="/",
-        write_empty_chunks=False,
+        chunk_key_encoding={"name": "v2", "separator": "/"},
         **kwargs,
     )
 
     # Get header dtype in native order (little-endian 99.9% of the time)
     header_dtype = segy_file.spec.trace.header.dtype.newbyteorder("=")
-    header_array = metadata_root.create_dataset(
+    header_array = metadata_root.create_array(
         name="_".join([name, "trace_headers"]),
         shape=grid.shape[:-1],  # Same spatial shape as data
         chunks=chunks[:-1],  # Same spatial chunks as data
+        chunk_key_encoding={"name": "v2", "separator": "/"},
         compressor=header_compressor,
         dtype=header_dtype,
-        dimension_separator="/",
-        write_empty_chunks=False,
     )
 
     # Initialize chunk iterator (returns next chunk slice indices each iteration)

src/mdio/segy/helpers_segy.py

@@ -1,18 +1,21 @@
 """Helper functions for tinkering with SEG-Y related Zarr."""
 
-from zarr import Group
-from zarr import open_group
+from typing import TYPE_CHECKING
+
 from zarr.errors import ContainsGroupError
-from zarr.storage import FSStore
 
 from mdio.core.exceptions import MDIOAlreadyExistsError
 
 
-def create_zarr_hierarchy(store: FSStore, overwrite: bool) -> Group:
+if TYPE_CHECKING:
+    from zarr.core import Group
+
+
+def create_zarr_hierarchy(root_group: "Group", overwrite: bool) -> "Group":
     """Create `zarr` hierarchy for SEG-Y files.
 
     Args:
-        store: Output path where the converted output is written.
+        root_group: Output root group where data will be written.
         overwrite: Toggle for overwriting existing store.
 
     Returns:
@@ -21,14 +24,12 @@ def create_zarr_hierarchy(store: FSStore, overwrite: bool) -> Group:
     Raises:
         MDIOAlreadyExistsError: If a file with data already exists.
     """
-    root_group = open_group(store=store)
-
     try:
         root_group.create_group(name="data", overwrite=overwrite)
         root_group.create_group(name="metadata", overwrite=overwrite)
     except ContainsGroupError as e:
         msg = (
-            f"An MDIO file with data already exists at {store.path}. "
+            f"An MDIO file with data already exists at {root_group.store.path}. "
             "If this is intentional, please specify 'overwrite=True'."
         )
         raise MDIOAlreadyExistsError(msg) from e