Merge remote-tracking branch 'origin/zarr-v3' into zarrv3_with_uv

Author: BrianMichell

.github/workflows/tests.yml

@@ -14,19 +14,17 @@ jobs:
         include:
           - { python: "3.13", os: "ubuntu-latest", session: "pre-commit" }
           - { python: "3.13", os: "ubuntu-latest", session: "safety" }
+          #          - { python: "3.13", os: "ubuntu-latest", session: "mypy" }
           #          - { python: "3.12", os: "ubuntu-latest", session: "mypy" }
           #          - { python: "3.11", os: "ubuntu-latest", session: "mypy" }
-          #          - { python: "3.10", os: "ubuntu-latest", session: "mypy" }
           - { python: "3.13", os: "ubuntu-latest", session: "tests" }
           - { python: "3.12", os: "ubuntu-latest", session: "tests" }
           - { python: "3.11", os: "ubuntu-latest", session: "tests" }
-          - { python: "3.10", os: "ubuntu-latest", session: "tests" }
           - { python: "3.13", os: "windows-latest", session: "tests" }
           - { python: "3.13", os: "macos-latest", session: "tests" }
+          #          - { python: "3.13", os: "ubuntu-latest", session: "typeguard" }
           #          - { python: "3.12", os: "ubuntu-latest", session: "typeguard" }
           #          - { python: "3.11", os: "ubuntu-latest", session: "typeguard" }
-          #          - { python: "3.10", os: "ubuntu-latest", session: "typeguard" }
-          #          - { python: "3.10", os: "ubuntu-latest", session: "xdoctest" }
           - { python: "3.13", os: "ubuntu-latest", session: "docs-build" }
 
     env:

noxfile.py

@@ -14,7 +14,7 @@
 
 
 package = "mdio"
-python_versions = ["3.13", "3.12", "3.11", "3.10"]
+python_versions = ["3.13", "3.12", "3.11"]
 nox.needs_version = ">= 2024.10.9"
 nox.options.sessions = (
     "pre-commit",

pyproject.toml

@@ -3,7 +3,7 @@ name = "multidimio"
 version = "0.8.5"
 description = "Cloud-native, scalable, and user-friendly multi dimensional energy data!"
 authors = [{ name = "Altay Sansal", email = "[email protected]" }]
-requires-python = ">=3.10,<3.14"
+requires-python = ">=3.11,<3.14"
 readme = "README.md"
 license = { "text" = "Apache-2.0" }
 keywords = [
@@ -25,7 +25,7 @@ classifiers = [
 dependencies = [
     "click (>=8.1.7,<9.0.0)",
     "click-params (>=0.5.0,<0.6.0)",
-    "zarr (>=2.18.2,<3.0.0)",
+    "zarr (>=3.0.2,<4.0.0)",
     "dask (>=2024.12.0)",
     "tqdm (>=4.67.0,<5.0.0)",
     "psutil (>=6.1.0,<7.0.0)",

src/mdio/api/accessor.py

@@ -186,32 +186,33 @@ 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,
         )
 
-    def _connect(self):
-        """Open the zarr root."""
         try:
-            if self.mode in {"r", "r+"}:
-                self.root = zarr.open_consolidated(store=self.store, mode=self.mode)
-            elif self.mode == "w":
-                self.root = zarr.open(store=self.store, mode="r+")
-            else:
-                msg = f"Invalid mode: {self.mode}"
-                raise ValueError(msg)
-        except KeyError as e:
+            self.store = zarr.open(
+                self.url, mode=self.mode, storage_options=storage_options
+            ).store
+        except FileNotFoundError as e:
             msg = (
-                f"MDIO file not found or corrupt at {self.store.path}. "
+                f"MDIO file not found or corrupt at {self.url}. "
                 "Please check the URL or ensure it is not a deprecated "
                 "version of MDIO file."
             )
             raise MDIONotFoundError(msg) from e
 
+    def _connect(self):
+        """Open the zarr root."""
+        if self.mode in {"r", "r+"}:
+            self.root = zarr.open_consolidated(store=self.store, mode=self.mode)
+        elif self.mode == "w":
+            self.root = zarr.open(store=self.store, mode="r+")
+        else:
+            msg = f"Invalid mode: {self.mode}"
+            raise ValueError(msg)
+
     def _deserialize_grid(self):
         """Deserialize grid from Zarr metadata."""
         self.grid = Grid.from_zarr(self.root)
@@ -375,12 +376,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
@@ -51,20 +51,22 @@ def copy_mdio(  # noqa: PLR0913
             Default is None (will assume anonymous).
         overwrite: Overwrite destination or not.
 
+    Raises:
+        NotImplementedError: because Zarr v3 doesn't support copy.
     """
     if storage_options is None:
         storage_options = {}
 
     dest_store = process_url(
         url=dest_path_or_buffer,
-        mode="w",
-        storage_options=storage_options,
-        memory_cache_size=0,
         disk_cache=False,
     )
 
     if_exists = "replace" if overwrite is True else "raise"
 
+    # TODO(Altay): Update this function when Zarr v3 supports copy.
+    raise NotImplementedError("Zarr version 3.0.0+ does not support copy yet.")
+
     zarr.copy_store(
         source=source.store,
         dest=dest_store,
@@ -242,6 +244,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,27 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 import dask.array as da
 import zarr
-from zarr.storage import FSStore
 
 
 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 +39,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 +50,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 +62,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 url
 
 
 def open_zarr_array(group_handle: zarr.Group, name: str) -> zarr.Array:

src/mdio/commands/info.py

@@ -95,7 +95,7 @@ def parse_grid(grid: Grid) -> dict[str, dict[str, int | str]]:
 def parse_access_patterns(reader: MDIOReader) -> dict[str, Any]:
     """Extract access patterns and their info."""
     access_pattern_dict = {}
-    for name, array in reader._data_group.items():
+    for name, array in reader._data_group.arrays():
         pattern = name.replace("chunked_", "")
         chunks = str(array.chunks)
         format_ = str(array.dtype)

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)