Repository: dask/dask
Pull request #11650 (2025-01-08 17:07:31)

Add cached version for ``normalize_chunks``
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

cc @dcherian is this suitable for xarray?

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     15 files  ± 0       15 suites  ±0   4h 29m 11s ⏱️ -2s
 17 163 tests + 1   15 968 ✅ + 1   1 195 💤 ±0  0 ❌ ±0 
211 476 runs  +13  194 296 ✅ +10  17 180 💤 +3  0 ❌ ±0 

Results for commit f4c00115. ± Comparison against base commit 7393a778.

<details>
  <summary>This pull request <b>removes</b> 1 and <b>adds</b> 2 tests. <i>Note that renamed tests count towards both.</i></summary>

```
dask.array.tests.test_array_core ‑ test_normalize_chunks
```

```
dask.array.tests.test_array_core ‑ test_normalize_chunks[normalize_chunks]
dask.array.tests.test_array_core ‑ test_normalize_chunks[normalize_chunks_cached]
```
</details>

[test-results]:data:application/gzip;base64,H4sIABe8fmcC/02MywrDIBQFfyW47sJrfPZnivgAaRKL0VXpv1clMVnOHM58kQ+L29FzAvaY0F5CvsiWpHOIW2MODKqqa+67AD6f/NqLMf2kuLzkO3yaBMWG8zos1eEhXEoxHSaVraUJABX84JFWlKhLHmkB8jze053vZRPXNeQKyFODMQDDFDPOrVRyFlQ4X5X02tnZCUkI0+j3B6T98QoaAQAA

dcherian: Yes decent impact.

```
# BEFORE: dask  : 2024.12.1, xarray: 2025.1.0
# Wall time: 1min 12s
#
# AFTER: xarray: 2024.11.1.dev50+gd9365109.d20250108, dask  : 2024.12.1+966.gf4c001150
# Wall time: 49.1 s
%time xr.open_zarr("gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3")
```

```
from dask.array.core import normalize_chunks_cached

normalize_chunks_cached.cache_info()
# CacheInfo(hits=271, misses=2, maxsize=128, currsize=2)
```

Now the big issue is tokenize.
dcherian: This is the issue:
```
from dask.layers import ArraySliceDep

chunks = ((1,)*1_000_000, (721,), (1441,))
%timeit dask.base.tokenize(ArraySliceDep(chunks)) # 55ms
```

for 300 variables, that alone takes 16s. somehow using `id(chunks)` in that tokenization would fix it.

Patch:
diff --git a/dask/array/api.py b/dask/array/api.py
new file mode 100644
index 000000000..b8dab59a0
--- /dev/null
+++ b/dask/array/api.py
@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+from dask.array.core import normalize_chunks, normalize_chunks_cached  # noqa: F401
diff --git a/dask/array/core.py b/dask/array/core.py
index 41c4389cc..3576e660d 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import math
 import operator
 import os
@@ -3007,6 +3008,24 @@ def ensure_int(f):
     return i
 
 
+@functools.lru_cache
+def normalize_chunks_cached(
+    chunks, shape=None, limit=None, dtype=None, previous_chunks=None
+):
+    """Cached version of normalize_chunks.
+
+    .. note::
+
+        chunks and previous_chunks are expected to be hashable. Dicts and lists aren't
+        allowed for this function.
+
+    See :func:`normalize_chunks` for further documentation.
+    """
+    return normalize_chunks(
+        chunks, shape=shape, limit=limit, dtype=dtype, previous_chunks=previous_chunks
+    )
+
+
 def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks=None):
     """Normalize chunks to tuple of tuples
 
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index 2e0143040..89bf42828 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -521,6 +521,17 @@ Generalized Ufuncs
    gufunc
 
 
+Downstream Libraries API
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. currentmodule:: dask.array.api
+
+.. autosummary::
+   :toctree: generated/
+
+   normalize_chunks_cached
+   normalize_chunks
+
 Internal functions
 ~~~~~~~~~~~~~~~~~~
 
diff --git a/pyproject.toml b/pyproject.toml
index b1f40278b..2f6662ab3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -161,7 +161,7 @@ platform = "linux"
 # platform = win32
 # platform = darwin
 plugins = ["numpy.typing.mypy_plugin"]
-allow_untyped_decorators = false
+disallow_untyped_decorators = true
 ignore_missing_imports = true
 no_implicit_optional = true
 show_error_codes = true
@@ -176,7 +176,7 @@ exclude = "dask/dataframe/dask_expr/"
 module = [
     "dask.order",
 ]
-allow_untyped_defs = false
+disallow_untyped_defs = true
 
 [tool.codespell]
 ignore-words-list = "coo,nd,medias"


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index f3b3ed733..e72ec9d30 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -47,6 +47,7 @@ from dask.array.core import (
     graph_from_arraylike,
     load_store_chunk,
     normalize_chunks,
+    normalize_chunks_cached,
     optimize,
     stack,
     store,
@@ -3056,30 +3057,30 @@ def test_from_array_with_missing_chunks():
     assert d.chunks == da.from_array(x, chunks=(2, 2, 3)).chunks
 
 
-def test_normalize_chunks():
-    assert normalize_chunks(3, (4, 6)) == ((3, 1), (3, 3))
-    assert normalize_chunks(((3, 3), (8,)), (6, 8)) == ((3, 3), (8,))
-    assert normalize_chunks((4, 5), (9,)) == ((4, 5),)
-    assert normalize_chunks((4, 5), (9, 9)) == ((4, 4, 1), (5, 4))
-    assert normalize_chunks(-1, (5, 5)) == ((5,), (5,))
-    assert normalize_chunks((3, -1), (5, 5)) == ((3, 2), (5,))
-    assert normalize_chunks((3, None), (5, 5)) == ((3, 2), (5,))
-    assert normalize_chunks({0: 3}, (5, 5)) == ((3, 2), (5,))
-    assert normalize_chunks([[2, 2], [3, 3]]) == ((2, 2), (3, 3))
-    assert normalize_chunks(10, (30, 5)) == ((10, 10, 10), (5,))
-    assert normalize_chunks((), (0, 0)) == ((0,), (0,))
-    assert normalize_chunks(-1, (0, 3)) == ((0,), (3,))
-    assert normalize_chunks(((float("nan"),),)) == ((np.nan,),)
-
-    assert normalize_chunks("auto", shape=(20,), limit=5, dtype="uint8") == (
-        (5, 5, 5, 5),
-    )
-    assert normalize_chunks(("auto", None), (5, 5), dtype=int) == ((5,), (5,))
+@pytest.mark.parametrize("func", [normalize_chunks, normalize_chunks_cached])
+def test_normalize_chunks(func):
+    assert func(3, (4, 6)) == ((3, 1), (3, 3))
+    assert func(((3, 3), (8,)), (6, 8)) == ((3, 3), (8,))
+    assert func((4, 5), (9,)) == ((4, 5),)
+    assert func((4, 5), (9, 9)) == ((4, 4, 1), (5, 4))
+    assert func(-1, (5, 5)) == ((5,), (5,))
+    assert func((3, -1), (5, 5)) == ((3, 2), (5,))
+    assert func((3, None), (5, 5)) == ((3, 2), (5,))
+    if func is normalize_chunks:
+        assert func({0: 3}, (5, 5)) == ((3, 2), (5,))
+        assert func([[2, 2], [3, 3]]) == ((2, 2), (3, 3))
+    assert func(10, (30, 5)) == ((10, 10, 10), (5,))
+    assert func((), (0, 0)) == ((0,), (0,))
+    assert func(-1, (0, 3)) == ((0,), (3,))
+    assert func(((float("nan"),),)) == ((np.nan,),)
+
+    assert func("auto", shape=(20,), limit=5, dtype="uint8") == ((5, 5, 5, 5),)
+    assert func(("auto", None), (5, 5), dtype=int) == ((5,), (5,))
 
     with pytest.raises(ValueError):
-        normalize_chunks(((10,),), (11,))
+        func(((10,),), (11,))
     with pytest.raises(ValueError):
-        normalize_chunks(((5,), (5,)), (5,))
+        func(((5,), (5,)), (5,))
 
 
 def test_single_element_tuple():