fix(indexing): delegate source tokenization to Dask

packages/zarr-indexing/changes/4348.bugfix.md

@@ -0,0 +1 @@
+Delegate LazyArray source tokenization to Dask, honoring its registered normalizers, source hooks, and deterministic-token requirements. Remove local content-hashing and UUID fallbacks. Dask remains optional for indexing and reading.

packages/zarr-indexing/docs/guide/integrations.md

@@ -243,3 +243,14 @@ that implementation or reproduce its full worker/GPU lifecycle.
 ·
 **API:** [API reference](../api/index.md)
 </nav>
+
+## Dask tokenization
+
+`LazyArray.__dask_tokenize__()` combines Dask's token for the wrapped source
+with the serialized view transform. Dask owns source hashing, registered
+normalizers, custom source hooks, and deterministic-token requirements.
+Tokenization may read or hash source values. Dask is optional for indexing
+and reading, but required when requesting a Dask token.
+
+The reader and partitioning are omitted because they must preserve values.
+Changing a source after graph construction does not update existing Dask keys.

packages/zarr-indexing/pyproject.toml

@@ -99,9 +99,8 @@ extend = "../../pyproject.toml"
 target-version = "py312"
 
 [tool.ruff.lint.per-file-ignores]
-# Chunk discovery and __dask_tokenize__ deliberately catch Exception: a
-# foreign source's attributes or token hooks may fail, so these paths provide
-# fallback metadata or tokens when ordinary exceptions occur. Configured here (not as
+# Chunk discovery deliberately catches Exception: a foreign source's
+# attributes may fail, so this path provides fallback metadata. Configured here (not as
 # noqa comments) because different ruff versions
 # have differed on whether these rules fire; RUF100 can remove unused noqa comments.
 "src/zarr_indexing/lazy_array.py" = ["BLE001", "S110"]

packages/zarr-indexing/src/zarr_indexing/lazy_array.py

@@ -12,8 +12,8 @@
 ```
 
 Selection construction does not read source values: `result()`, `__array__`,
-and eager `__getitem__` perform reads. Tokenization can also read or hash source
-data, depending on the source and tokenization path. `.lazy` operations inspect
+and eager `__getitem__` perform reads. Source tokenization is delegated to
+Dask and may inspect source values. `.lazy` operations inspect
 selection metadata and may copy or process supplied index arrays. Composition
 does not accumulate wrapper layers: a view of a view is still a single transform
 and retains its reader.
@@ -138,7 +138,6 @@
 import json
 import math
 import operator
-import uuid
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Protocol, cast
@@ -173,10 +172,6 @@
 
 __all__ = ["LazyArray", "Partition"]
 
-# Above this declared byte count, the no-Dask token fallback adds a fresh UUID
-# instead of digesting contents. See `_wrapped_token`.
-_TOKEN_DIGEST_LIMIT = 1 << 20
-
 
 def _invoke_reader(
     reader: Reader,
@@ -525,66 +520,6 @@ def _validate_prepared_parts(parts: Sequence[Partition], out_shape: tuple[int, .
         raise ValueError("prepared parts do not tile the view exactly")
 
 
-# --------------------------------------------------------------------------- #
-# Tokenization
-# --------------------------------------------------------------------------- #
-
-
-def _wrapped_token(array: Any) -> Any:
-    """A token for the wrapped array.
-
-    In order of preference: the array's own `__dask_tokenize__`;
-    `dask.base.tokenize` when dask is importable (imported lazily — this package
-    never requires it); otherwise a local fallback that digests the contents of
-    a small array.
-
-    Tokens can differ depending on whether Dask is available and on the source
-    hook. This fallback does not provide a portable content identifier.
-
-    Above `_TOKEN_DIGEST_LIMIT`, or when conversion is unavailable, the local
-    fallback adds a fresh UUID on each call, so repeated calls normally differ.
-    The returned tuple is still equal to itself. Below the limit, conversion
-    can read the source, and the hash is of its NumPy buffer bytes. Object-array
-    buffer bytes contain object references, not a recursive content snapshot.
-    """
-    hook = getattr(array, "__dask_tokenize__", None)
-    if hook is not None:
-        try:
-            return hook()
-        # A failing source hook falls through to the remaining tokenization paths.
-        except Exception:  # pragma: no cover - a hook that refuses to run
-            pass
-    try:
-        # dask is an optional peer, never a dependency of this package, so it is
-        # imported here and its absence is ordinary.
-        from dask.base import tokenize  # pyright: ignore[reportMissingImports]
-    except ImportError:
-        pass
-    else:
-        return tokenize(array)
-
-    shape = tuple(int(s) for s in getattr(array, "shape", ()))
-    dtype = getattr(array, "dtype", None)
-    structural = (type(array).__qualname__, shape, str(dtype))
-    # A fresh identifier per call when contents cannot be identified. It
-    # is the shape and dtype that would otherwise be mistaken for an identity,
-    # so they are kept alongside it for a reader looking at a graph.
-    unidentified = (*structural, "unidentified", uuid.uuid4().hex)
-
-    # Decide whether to digest the contents from the *declared* size. Measuring
-    # it by converting first would read the whole array — a multi-gigabyte store
-    # pulled into memory by a token call, which is the opposite of the point.
-    itemsize = getattr(dtype, "itemsize", None)
-    if not isinstance(itemsize, int) or itemsize * math.prod(shape) > _TOKEN_DIGEST_LIMIT:
-        return unidentified
-    try:
-        contents = np.ascontiguousarray(array)
-    # Failed NumPy conversion leaves this source unidentified.
-    except Exception:
-        return unidentified
-    return (*structural, hashlib.sha256(contents.tobytes()).hexdigest())
-
-
 # --------------------------------------------------------------------------- #
 # The wrapper
 # --------------------------------------------------------------------------- #
@@ -1249,15 +1184,18 @@ def __dask_tokenize__(self) -> Any:
         tokens produce equal tokens; arbitrary semantically equivalent mappings
         are not guaranteed to serialize identically.
 
-        Source tokenization can read or hash data and need not be deterministic
-        on every fallback path; see `_wrapped_token`. The reader and partitioning
-        are omitted under the contract that they preserve values. Cache users
-        must also account for source mutation and the source's token semantics.
+        Dask tokenizes the wrapped source using its normal dispatch and
+        determinism policy. This may read or hash source values. Dask is
+        imported only when this method is called and is otherwise optional.
+        The reader and partitioning are omitted because they must preserve
+        values. Mutating a source does not update keys in existing Dask graphs.
         """
+        from dask.base import tokenize  # pyright: ignore[reportMissingImports]
+
         canonical = json.dumps(self._transform.to_json(), sort_keys=True)
         return (
             type(self).__qualname__,
-            _wrapped_token(self._array),
+            tokenize(self._array),
             hashlib.sha256(canonical.encode()).hexdigest(),
         )
 

packages/zarr-indexing/tests/test_lazy_array.py

@@ -1776,6 +1776,7 @@ def test_nonfirst_partition_transform_directly_addresses_its_array() -> None:
 
 
 def test_partition_token_encodes_its_public_global_transform() -> None:
+    pytest.importorskip("dask.base")
     source = np.arange(8)
     base = LazyArray.from_numpy(source)
     partition_view = list(base.with_parts((4,)).parts())[1].view
@@ -1907,6 +1908,7 @@ def test_with_parts_validates_strictly(parts: Any, match: str) -> None:
 
 def test_dask_token_is_deterministic_and_discriminating() -> None:
     """Same data and same view token alike; a different selection differs."""
+    pytest.importorskip("dask.base")
     data = reference()
     base = LazyArray(data)
     assert base.__dask_tokenize__() == LazyArray(reference()).__dask_tokenize__()
@@ -1923,6 +1925,7 @@ def test_dask_token_is_deterministic_and_discriminating() -> None:
 
 
 def test_reader_and_partitioning_do_not_change_dask_identity() -> None:
+    pytest.importorskip("dask.base")
     base = LazyArray(reference())
     token = base.__dask_tokenize__()
     assert base.with_reader(numpy_reader).__dask_tokenize__() == token
@@ -2000,7 +2003,6 @@ def test_pickle_round_trip() -> None:
     view = LazyArray(reference()).with_parts((2, 2, 2)).lazy[1:6, ::2].lazy.oindex[[3, 0, 0], :, :]
     restored = pickle.loads(pickle.dumps(view))
     assert restored.shape == view.shape
-    assert restored.__dask_tokenize__() == view.__dask_tokenize__()
     np.testing.assert_array_equal(np.asarray(restored.result()), np.asarray(view.result()))
 
 
@@ -2010,7 +2012,7 @@ def test_pickle_round_trip() -> None:
 
 
 def test_dask_from_array_roundtrip() -> None:
-    """A `LazyArray` is a drop-in dask source — no translation ceremony."""
+    """Dask can tokenize and read a wrapper over a Zarr source."""
     da = pytest.importorskip("dask.array")
     source = make_source("zarr")
 
@@ -2435,25 +2437,6 @@ def test_a_masked_source_keeps_its_mask_when_the_view_is_empty(parts: Any) -> No
     assert np.asarray(got).shape == (3, 0), parts
 
 
-def test_a_large_array_without_dask_refuses_to_claim_equality(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """Without Dask, arrays above the digest limit receive distinct fallback tokens."""
-    import sys
-
-    monkeypatch.setitem(sys.modules, "dask.base", None)
-    big = np.zeros(1 << 19, dtype=np.int64)
-    other = big.copy()
-    other[0] = 1
-
-    assert LazyArray(big).__dask_tokenize__() != LazyArray(other).__dask_tokenize__()
-    assert LazyArray(big).__dask_tokenize__() != LazyArray(big).__dask_tokenize__()
-
-    # Below the limit the contents are digested, so equal data still tokens alike.
-    small = np.zeros(8, dtype=np.int64)
-    assert LazyArray(small).__dask_tokenize__() == LazyArray(small.copy()).__dask_tokenize__()
-
-
 # ---------------------------------------------------------------------------
 # Completeness and partition spellings
 # ---------------------------------------------------------------------------
@@ -2576,3 +2559,63 @@ def test_fancy_composition_over_an_empty_axis() -> None:
     scalar = composed.lazy.vindex[..., np.array(1)]
     assert scalar.shape == (2, 0)
     assert np.asarray(scalar.result()).shape == (2, 0)
+
+
+@pytest.mark.parametrize("kind", ["numpy", "object", "masked", "registered", "hook"])
+def test_source_token_uses_dask_policy(kind: str) -> None:
+    dask_base = pytest.importorskip("dask.base")
+
+    class RegisteredArray(ForeignArray):
+        pass
+
+    class VersionedArray(ForeignArray):
+        def __dask_tokenize__(self) -> Any:
+            return ("versioned-source", 1)
+
+    dask_base.normalize_token.register(RegisteredArray, lambda source: ("registered-source", 1))
+    data = np.arange(4)
+    sources = {
+        "numpy": data,
+        "object": data.astype(object),
+        "masked": np.ma.masked_greater(data, 2),
+        "registered": RegisteredArray(data, None),
+        "hook": VersionedArray(data, None),
+    }
+    source = sources[kind]
+    assert LazyArray(source).__dask_tokenize__()[1] == dask_base.tokenize(source)
+
+
+def test_source_token_preserves_dask_determinism_requirement() -> None:
+    dask_base = pytest.importorskip("dask.base")
+    dask_tokenize = pytest.importorskip("dask.tokenize")
+
+    class UnserializableArray(ForeignArray):
+        def __reduce_ex__(self, protocol: int) -> Any:
+            raise TypeError("cannot serialize source")
+
+    source = UnserializableArray(np.arange(4), None)
+    for value in (source, LazyArray(source)):
+        with pytest.raises(dask_tokenize.TokenizationError):
+            dask_base.tokenize(value, ensure_deterministic=True)
+
+
+def test_source_token_preserves_hook_failure() -> None:
+    pytest.importorskip("dask.base")
+
+    class RefusingArray(ForeignArray):
+        def __dask_tokenize__(self) -> Any:
+            raise RuntimeError("source version unavailable")
+
+    with pytest.raises(RuntimeError, match="source version unavailable"):
+        LazyArray(RefusingArray(np.arange(4), None)).__dask_tokenize__()
+
+
+def test_dask_is_only_required_for_tokenization(monkeypatch: pytest.MonkeyPatch) -> None:
+    import sys
+
+    monkeypatch.setitem(sys.modules, "dask.base", None)
+    data = np.arange(4)
+    view = LazyArray(data).lazy[1:]
+    np.testing.assert_array_equal(view.result(), data[1:])
+    with pytest.raises(ModuleNotFoundError):
+        view.__dask_tokenize__()