Skip to content

Commit a26876d

Browse files
committed
fix(indexing): delegate source tokenization to Dask
Assisted-by: Codex:GPT-6
1 parent 3e727e7 commit a26876d

8 files changed

Lines changed: 80 additions & 289 deletions

File tree

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
Define an explicit source token contract: hash supported plain NumPy arrays field by field without structured padding, delegate source hooks and propagate their errors, and reject unsupported sources without converting or serializing them. Object fields, dtype metadata, and non-string field titles require a source hook. Remove optional-Dask and per-call UUID fallbacks. Document hashing costs, mutation limits, and `name=False` for Dask wrapping opaque sources.
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/index.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -225,8 +225,8 @@ example that materializes the selected data.
225225
(`numpy.add(view, 1)` does so; `numpy.shape(view)` and `numpy.ndim(view)`
226226
can use metadata without reading values)
227227

228-
Dask tokenization hashes plain NumPy data or delegates to a source hook.
229-
Other sources require [an explicit hook or Dask naming opt-out](integrations.md#dask-tokenization-and-source-mutation).
228+
Dask tokenization may also inspect values, depending on the wrapped source
229+
and tokenization path.
230230

231231
Python arithmetic such as `view + 1` raises `TypeError` instead: this
232232
wrapper defers indexing, not a general compute graph.

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

Lines changed: 10 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -244,39 +244,13 @@ that implementation or reproduce its full worker/GPU lifecycle.
244244
**API:** [API reference](../api/index.md)
245245
</nav>
246246

247-
## Dask tokenization and source mutation
248-
249-
`LazyArray.__dask_tokenize__()` combines the serialized transform with a source
250-
value token. Supported plain NumPy arrays are hashed in full on every call,
251-
including large arrays. Structured fields are hashed separately, excluding
252-
padding; overlapping fields are each visited. Time scales with the bytes
253-
visited, and temporary memory with the largest field buffer. Equal dtype
254-
representations, shapes, field bytes in C order, and serialized transforms have
255-
equal tokens. This is byte equality within fields, so signed zeros and distinct
256-
NaN bit patterns can produce different tokens despite numerical equivalence.
257-
258-
Other sources must define an explicit `__dask_tokenize__` hook. This includes
259-
object arrays (also structured object fields), array subclasses, memory-mapped
260-
arrays, and remote arrays. Dtype metadata and non-string field titles also
261-
require a source hook, including in nested field and subarray dtypes: they can
262-
contain arbitrary Python objects with no guaranteed value-based representation.
263-
Known `numpy.memmap` and `mmap.mmap` backing is
264-
rejected through ndarray base and memoryview object chains. Arbitrary buffer
265-
provenance cannot be inferred from a plain ndarray. The wrapper does not
266-
convert or serialize unsupported sources to discover their values. Unsupported sources raise `TypeError`, and
267-
exceptions from explicit hooks propagate. Installing Dask does not change this
268-
policy. A hook must describe the source's values or immutable version, and owns
269-
its determinism and any I/O it performs.
270-
271-
For an opaque source such as a Zarr array, use
272-
`dask.array.from_array(view, chunks=..., name=False)` to request a fresh graph
273-
name without content tokenization. This opts out of content-based task sharing.
274-
See Dask's [from_array documentation](https://docs.dask.org/en/stable/generated/dask.array.from_array.html)
275-
and [tokenization contract](https://docs.dask.org/en/stable/custom-collections.html#implementing-deterministic-hashing).
276-
277-
Tokens describe values at tokenization time, not a snapshot. Mutating a source
278-
after building a graph does not update existing Dask keys or invalidate cached
279-
results. Concurrent mutation during hashing is unsupported. Applications must
280-
manage source lifetimes and versions; these tokens are not persistent cache
281-
identities. Readers and partitioning are omitted because they must preserve
282-
values.
247+
## Dask tokenization
248+
249+
`LazyArray.__dask_tokenize__()` combines Dask's token for the wrapped source
250+
with the serialized view transform. Dask owns source hashing, registered
251+
normalizers, custom source hooks, and deterministic-token requirements.
252+
Tokenization may read or hash source values. Dask is optional for indexing
253+
and reading, but required when requesting a Dask token.
254+
255+
The reader and partitioning are omitted because they must preserve values.
256+
Changing a source after graph construction does not update existing Dask keys.

‎packages/zarr-indexing/docs/index.md‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,8 @@ you want its values:
2828
```
2929

3030
Composing these selections does not read source values; the example reads them
31-
at `result()`. Construction inspects source metadata. Dask tokenization hashes
32-
plain NumPy data or delegates to an explicit source hook; other sources need
33-
[a naming opt-out or source hook](guide/integrations.md#dask-tokenization-and-source-mutation). [Lazy views compose](guide/index.md#lazy-views-compose) shows how
31+
at `result()`. Construction inspects source metadata, and Dask tokenization can
32+
inspect source values. [Lazy views compose](guide/index.md#lazy-views-compose) shows how
3433
the chain stays one description, and where the materialization boundary is.
3534

3635
## Learn more

‎packages/zarr-indexing/examples/lazy_indexing_dask/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,11 @@ ways of deferring an indexing operation.
77
The example shows how to:
88

99
- Pass a `LazyArray` — over a Zarr array or over a view of one — to
10-
`dask.array.from_array` with `name=False` to bypass source tokenization
10+
`dask.array.from_array`
1111
- Build one Dask task per partition from `parts()`, compute them in parallel, and
1212
place each result with the partition's `out_selection`
1313
- Read `is_complete` to inspect coverage of a partition cell
14-
- Inspect `__dask_tokenize__` for equal NumPy source/selection pairs;
14+
- Inspect `__dask_tokenize__` for the example's equal source/selection pairs;
1515
token equality can support task deduplication but does not promise persistent caching
1616
- Measure what a task graph costs for indexing-only work, against composing the
1717
same selections into one transform

‎packages/zarr-indexing/examples/lazy_indexing_dask/lazy_indexing_dask.py‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -41,14 +41,13 @@ def test_from_array(source: zarr.Array) -> None:
4141

4242
# `from_array` needs `shape`, `dtype`, and `__getitem__`, which the wrapper
4343
# provides. Each Dask block reads its own region through the wrapper.
44-
# Zarr has no source token hook: request a fresh Dask graph name.
45-
array = da.from_array(lazy, chunks=(10, 10), name=False)
44+
array = da.from_array(lazy, chunks=(10, 10))
4645
print(array)
4746
assert np.array_equal(array.compute(scheduler="threads"), source[:])
4847

4948
# A view works the same way, and its shape is the shape of the selection.
5049
view = LazyArray(source).lazy[5:35, 3:27]
51-
array = da.from_array(view, chunks=(10, 10), name=False)
50+
array = da.from_array(view, chunks=(10, 10))
5251
assert array.shape == (30, 24)
5352
assert np.array_equal(array.compute(scheduler="threads"), source[5:35, 3:27])
5453

@@ -83,9 +82,8 @@ def read(part: object) -> np.ndarray:
8382
print(f"{len(complete)} of {len(parts)} parts cover their chunk completely")
8483

8584

86-
def test_tokenize() -> None:
85+
def test_tokenize(source: zarr.Array) -> None:
8786
"""Check token equality for these unchanged source and selection pairs."""
88-
source = np.arange(40 * 30).reshape(40, 30)
8987
lazy = LazyArray(source)
9088

9189
# These wrappers have equal tokens despite being different Python objects.

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

Lines changed: 10 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212
```
1313
1414
Selection construction does not read source values: `result()`, `__array__`,
15-
and eager `__getitem__` perform reads. Tokenization hashes plain NumPy source
16-
data or delegates to an explicit source hook. `.lazy` operations inspect
15+
and eager `__getitem__` perform reads. Source tokenization is delegated to
16+
Dask and may inspect source values. `.lazy` operations inspect
1717
selection metadata and may copy or process supplied index arrays. Composition
1818
does not accumulate wrapper layers: a view of a view is still a single transform
1919
and retains its reader.
@@ -137,7 +137,6 @@
137137
import hashlib
138138
import json
139139
import math
140-
import mmap
141140
import operator
142141
from collections.abc import Sequence
143142
from dataclasses import dataclass, field
@@ -521,66 +520,6 @@ def _validate_prepared_parts(parts: Sequence[Partition], out_shape: tuple[int, .
521520
raise ValueError("prepared parts do not tile the view exactly")
522521

523522

524-
# --------------------------------------------------------------------------- #
525-
# Tokenization
526-
# --------------------------------------------------------------------------- #
527-
528-
529-
def _validate_token_dtype(dtype: np.dtype[Any]) -> None:
530-
"""Require dtype attributes with a deterministic, value-based representation."""
531-
if dtype.metadata is not None:
532-
raise TypeError("Dtype metadata requires an explicit source __dask_tokenize__ hook.")
533-
if dtype.subdtype is not None:
534-
_validate_token_dtype(dtype.subdtype[0])
535-
if dtype.fields is not None:
536-
for field in dtype.fields.values():
537-
if len(field) == 3 and type(field[2]) is not str:
538-
raise TypeError(
539-
"Non-string field titles require an explicit source __dask_tokenize__ hook."
540-
)
541-
_validate_token_dtype(field[0])
542-
543-
544-
def _numpy_values_token(array: np.ndarray[Any, Any]) -> Any:
545-
"""Hash field values in C order, excluding structured padding bytes."""
546-
if array.dtype.fields is not None:
547-
return tuple(_numpy_values_token(array[name]) for name in array.dtype.names or ())
548-
return hashlib.sha256(array.tobytes(order="C")).hexdigest()
549-
550-
551-
def _wrapped_token(array: Any) -> Any:
552-
"""Identify supported source contents without converting foreign arrays.
553-
554-
Explicit hooks own their versioning, determinism, and I/O contract.
555-
Supported plain NumPy arrays are hashed field by field on every call. Other
556-
sources must provide a hook; neither serialization nor identity is a safe
557-
substitute for a source's value/version contract.
558-
"""
559-
hook = getattr(array, "__dask_tokenize__", None)
560-
if hook is not None:
561-
return hook()
562-
if type(array) is not np.ndarray or array.dtype.hasobject:
563-
msg = (
564-
"Tokenization requires a source __dask_tokenize__ hook or a plain "
565-
"NumPy ndarray without object fields. For Dask from_array, use "
566-
"name=False to bypass content tokenization."
567-
)
568-
raise TypeError(msg)
569-
base: Any = array.base
570-
while isinstance(base, (np.ndarray, memoryview, mmap.mmap)):
571-
if isinstance(base, (np.memmap, mmap.mmap)):
572-
msg = "Memory-mapped sources require an explicit __dask_tokenize__ hook."
573-
raise TypeError(msg)
574-
base = base.obj if isinstance(base, memoryview) else base.base
575-
_validate_token_dtype(array.dtype)
576-
return (
577-
"numpy.ndarray",
578-
array.shape,
579-
repr(array.dtype),
580-
_numpy_values_token(array),
581-
)
582-
583-
584523
# --------------------------------------------------------------------------- #
585524
# The wrapper
586525
# --------------------------------------------------------------------------- #
@@ -1245,27 +1184,18 @@ def __dask_tokenize__(self) -> Any:
12451184
tokens produce equal tokens; arbitrary semantically equivalent mappings
12461185
are not guaranteed to serialize identically.
12471186
1248-
Supported plain NumPy arrays are hashed in full on each call. Structured
1249-
fields are hashed separately, excluding padding; overlapping fields
1250-
are each visited. Time scales with the bytes visited, and temporary
1251-
memory with the largest field buffer. Dtype metadata and non-string
1252-
field titles require an explicit source hook, including in nested dtypes.
1253-
Known NumPy/mmap backing is rejected through ndarray base and
1254-
memoryview object chains; arbitrary buffer provenance is not inferred.
1255-
Array subclasses, object arrays, and foreign sources require a source `__dask_tokenize__` hook; unsupported sources
1256-
raise `TypeError`. Hooks own determinism, versioning, and any I/O, and
1257-
hook exceptions propagate. Installing Dask does not change this policy.
1258-
1259-
The reader and partitioning are omitted under the contract that they
1260-
preserve values. A token describes the source at tokenization time;
1261-
mutation after graph construction does not invalidate existing Dask
1262-
keys. Concurrent mutation during hashing is unsupported. This method
1263-
does not provide a persistent cache identity or a source snapshot.
1187+
Dask tokenizes the wrapped source using its normal dispatch and
1188+
determinism policy. This may read or hash source values. Dask is
1189+
imported only when this method is called and is otherwise optional.
1190+
The reader and partitioning are omitted because they must preserve
1191+
values. Mutating a source does not update keys in existing Dask graphs.
12641192
"""
1193+
from dask.base import tokenize # pyright: ignore[reportMissingImports]
1194+
12651195
canonical = json.dumps(self._transform.to_json(), sort_keys=True)
12661196
return (
12671197
type(self).__qualname__,
1268-
_wrapped_token(self._array),
1198+
tokenize(self._array),
12691199
hashlib.sha256(canonical.encode()).hexdigest(),
12701200
)
12711201

0 commit comments

Comments
 (0)