Validate Zarr v3 with zarr-python and make zarr an optional extra

dandi/files/zarr.py

@@ -22,6 +22,7 @@
     from ..upload import ZarrMode
 
 from dandischema.models import BareAsset, DigestType
+from packaging.version import Version
 from pydantic import BaseModel, ConfigDict, ValidationError
 import requests
 from zarr_checksum.tree import ZarrChecksumTree
@@ -127,9 +128,9 @@ def get_zarr_format_version(path: Path) -> Optional[str]:
     return None
 
 
-def _ts_validate_zarr3(path: Path, devel_debug: bool = False) -> list[ValidationResult]:
+def _validate_zarr3(path: Path, devel_debug: bool = False) -> list[ValidationResult]:
     """
-    Validate a Zarr format V3 LocalStore with the tensorstore package
+    Validate a Zarr format V3 LocalStore with zarr-python
 
     Parameters
     ----------
@@ -151,12 +152,9 @@ def _ts_validate_zarr3(path: Path, devel_debug: bool = False) -> list[Validation
 
     Note
     ----
-        Since tensorstore does not support the concept of a Zarr group, this function
-        validates a Zarr format V3 LocalStore by opening all the contained arrays with
-        tensorstore individually.
-
-        This function will no longer be needed once the upgrade to zarr-python 3.x is
-        done and should be removed.
+        zarr-python can partially open a malformed Zarr format V3 store (e.g. valid
+        root metadata, bad child metadata), so this function validates a Zarr format
+        V3 LocalStore by opening all the contained arrays individually.
     """
 
     if not path.is_dir():
@@ -216,17 +214,17 @@ def _ts_validate_zarr3(path: Path, devel_debug: bool = False) -> list[Validation
             else:
                 # Check if the directory is a Zarr array
                 if meta.node_type == "array":
-                    results.extend(_ts_validate_zarr3_array(Path(root), devel_debug))
+                    results.extend(_validate_zarr3_array(Path(root), devel_debug))
                     dirs.clear()  # Skip subdirectories
 
     return results
 
 
-def _ts_validate_zarr3_array(
+def _validate_zarr3_array(
     path: Path, devel_debug: bool = False
 ) -> list[ValidationResult]:
     """
-    Validate a Zarr format V3 array in a LocalStore with the tensorstore package
+    Validate a Zarr format V3 array in a LocalStore with zarr-python
 
     Parameters
     ----------
@@ -239,41 +237,31 @@ def _ts_validate_zarr3_array(
     -------
     list[ValidationResult]
         A list of validation results representing validation errors encountered
-
-    Note
-    ----
-        This function will no longer be needed once the upgrade to zarr-python 3.x is
-        done and should be removed.
     """
     # Avoid heavy import by importing within function
-    from importlib.metadata import version
-
-    import tensorstore as ts  # type: ignore[import]
+    import zarr
 
     results: list[ValidationResult] = []
 
-    # TensorStore spec describing where and how to read the Zarr array
-    spec = {"driver": "zarr3", "kvstore": {"driver": "file", "path": str(path)}}
-
     try:
-        ts.open(spec, read=True, write=False).result()
+        zarr.open_array(path, mode="r", zarr_format=3)
     except Exception as e:
         if devel_debug:
             raise
         results.append(
             ValidationResult(
-                id="zarr.tensorstore_cannot_open",
+                id="zarr.cannot_open",
                 origin=Origin(
                     type=OriginType.INTERNAL,
-                    validator=Validator.tensorstore,
-                    validator_version=version("tensorstore"),
+                    validator=Validator.zarr,
+                    validator_version=zarr.__version__,
                     standard=Standard.ZARR,
                     standard_version="3",
                 ),
                 scope=Scope.FILE,
                 origin_result=e,
                 severity=Severity.ERROR,
-                message="Error opening Zarr array with tensorstore",
+                message="Error opening Zarr array",
                 path=path,
             )
         )
@@ -474,8 +462,39 @@ def get_validation_errors(
         devel_debug: bool = False,
         missing_file_content: MissingFileContent | None = None,
     ) -> list[ValidationResult]:
-        # Avoid heavy import by importing within function:
-        import zarr
+        # zarr-python is an optional dependency (the `zarr` extra), and only
+        # zarr-python 3.x can read Zarr format V3.
+        zarr_version: str | None
+        try:
+            # Avoid heavy import by importing within function:
+            import zarr
+        except ImportError:
+            zarr_version = None
+        else:
+            zarr_version = zarr.__version__
+        if zarr_version is None or Version(zarr_version).major < 3:
+            found = (
+                "it is not installed"
+                if zarr_version is None
+                else f"found {zarr_version}"
+            )
+            return [
+                ValidationResult(
+                    origin=ORIGIN_VALIDATION_DANDI_ZARR,
+                    severity=Severity.ERROR,
+                    id="dandi_zarr.zarr_python_required",
+                    scope=Scope.FILE,
+                    path=self.filepath,
+                    message=(
+                        f"Validating Zarr assets requires zarr-python 3.x ({found}). "
+                        "Install it with `pip install 'dandi[zarr]'`."
+                    ),
+                )
+            ] + super().get_validation_errors(
+                schema_version=schema_version,
+                devel_debug=devel_debug,
+                missing_file_content=missing_file_content,
+            )
 
         errors: list[ValidationResult] = []
         origin_internal_zarr: Origin = Origin(
@@ -485,20 +504,22 @@ def get_validation_errors(
             standard=Standard.ZARR,
         )
 
-        # Zarr V3 stores are validated by `_ts_validate_zarr3` regardless of
-        # zarr-python version. With zarr-python 2.x, `zarr.open()` cannot
-        # handle V3 at all (raises `PathNotFoundError`); with zarr-python 3.x
-        # it can sometimes partially open malformed V3 stores (e.g. valid
-        # root metadata, bad child metadata), which would cause us to miss
-        # the `_ts_validate_zarr3` path. Detecting the format up front keeps
-        # the validation behaviour consistent across versions.
+        # Zarr V3 stores are validated by `_validate_zarr3`, which opens every
+        # array individually: `zarr.open()` can sometimes partially open
+        # malformed V3 stores (e.g. valid root metadata, bad child metadata),
+        # which would cause us to miss errors.
         format_version = get_zarr_format_version(self.filepath)
         if format_version == "3":
-            errors.extend(_ts_validate_zarr3(self.filepath, devel_debug))
+            errors.extend(_validate_zarr3(self.filepath, devel_debug))
             data = None
         else:
             try:
-                data = zarr.open(str(self.filepath), mode="r")
+                # Consolidated metadata is skipped because zarr-python 3.x
+                # cannot parse some V2 array metadata written by zarr-python
+                # 2.x, such as object arrays with a fill value of 0 or the
+                # `json2` codec, both found in NWB files written by hdmf-zarr.
+                # `_is_empty_group` already tolerates such members.
+                data = zarr.open(str(self.filepath), mode="r", use_consolidated=False)
             except Exception as e:
                 if devel_debug:
                     raise

dandi/tests/test_files.py

@@ -694,7 +694,7 @@ def test_validate_valid_zarr3(path: str) -> None:
         ),
         # A store with a corrupt zarr.json for an array (missing fields other than
         # `node_type`)
-        ("array_v3_corrupt_zarr_json.zarr", {"zarr.tensorstore_cannot_open"}),
+        ("array_v3_corrupt_zarr_json.zarr", {"zarr.cannot_open"}),
     ],
 )
 def test_validate_invalid_zarr3(path: str, expected_result_ids: set[str]) -> None:
@@ -712,3 +712,55 @@ def test_validate_invalid_zarr3(path: str, expected_result_ids: set[str]) -> Non
 
     result_ids = {r.id for r in zf.get_validation_errors()}
     assert result_ids == expected_result_ids
+
+
+@pytest.mark.ai_generated
+def test_validate_zarr3_string_array(tmp_path: Path) -> None:
+    # Variable-length strings are an extension data type in Zarr format V3,
+    # and NWB files store most of their text in them
+    zarr_path = tmp_path / "strings.zarr"
+    group = zarr.open_group(zarr_path, mode="w", zarr_format=3)
+    group.create_array("identifier", shape=(1,), dtype=str)[:] = ["session"]
+    assert dandi_file(zarr_path).get_validation_errors() == []
+
+
+@pytest.mark.ai_generated
+def test_validate_zarr2_unparseable_consolidated_metadata(tmp_path: Path) -> None:
+    # zarr-python 2.x wrote a fill value of 0 for object arrays, which
+    # zarr-python 3.x cannot parse.  hdmf-zarr wrote such arrays in NWB files
+    # (e.g. `session_start_time`) and consolidated their metadata at the root.
+    zarr_path = tmp_path / "example.zarr"
+    array_meta = {
+        "chunks": [1],
+        "compressor": None,
+        "dtype": "|O",
+        "fill_value": 0,
+        "filters": [{"id": "vlen-bytes"}],
+        "order": "C",
+        "shape": [1],
+        "zarr_format": 2,
+    }
+    metadata = {".zgroup": {"zarr_format": 2}, "session_start_time/.zarray": array_meta}
+    (zarr_path / "session_start_time").mkdir(parents=True)
+    for name, meta in metadata.items():
+        (zarr_path / name).write_text(json.dumps(meta))
+    (zarr_path / ".zmetadata").write_text(
+        json.dumps({"metadata": metadata, "zarr_consolidated_format": 1})
+    )
+    assert dandi_file(zarr_path).get_validation_errors() == []
+
+
+@pytest.mark.ai_generated
+@pytest.mark.parametrize("zarr_version", [None, "2.18.7"])
+def test_validate_zarr_requires_zarr_python_3(
+    zarr_version: str | None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    # `None` stands for zarr-python not being installed
+    if zarr_version is None:
+        monkeypatch.setitem(sys.modules, "zarr", None)
+    else:
+        monkeypatch.setattr(zarr, "__version__", zarr_version)
+    zf = dandi_file(Path(__file__).parent / VALID_STORES_PATH / "single_array.zarr")
+    assert [r.id for r in zf.get_validation_errors()] == [
+        "dandi_zarr.zarr_python_required"
+    ]

pyproject.toml

@@ -81,13 +81,9 @@ dependencies = [
     "ruamel.yaml >= 0.16.0, < 1",
     "semantic-version >= 2.10.0",
     "tenacity >= 8.2.0",
-    # 0.1.62 first release built for numpy 2.x
-    "tensorstore >= 0.1.62",
     # possibly silently incomplete downloads: https://github.com/dandi/dandi-cli/issues/1500
     "urllib3 >= 2.0.2",
     "yarl >= 1.9.1, < 2",
-    # Floor raised: numcodecs 0.11 / hdmf-zarr chain requires >= 2.18.
-    "zarr >= 2.18.0, < 4",
     "zarr_checksum ~= 0.4.0",
 ]
 dynamic = ["version"]
@@ -121,15 +117,24 @@ test = [
     "responses >= 0.23.0, != 0.24.0, != 0.25.5",
     # 5.0 switched to urllib3.connection.VerifiedHTTPSConnection (urllib3 2.x compat)
     "vcrpy >= 5.0.0",
+    "dandi[zarr]",
 ]
 tools = [
     "boto3",
 ]
+# zarr-python is needed only to validate Zarr assets
+zarr = [
+    # Reading strings in Zarr format V3 warns before 3.1.1, and the expected
+    # digests in the tests assume 3.1.4, which stopped writing
+    # `consolidated_metadata: null`.
+    "zarr >= 3.1.4, < 4",
+]
 all = [
     "dandi[extensions]",
     "dandi[extras]",
     "dandi[test]",
     "dandi[tools]",
+    "dandi[zarr]",
 ]
 
 # Development-only dependency sets, not published in built distributions.