From 23437492737c3051f413b095ee4ec948d584d77f Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Fri, 13 Sep 2024 17:02:06 -0700
Subject: [PATCH 01/12] switch to using zarr.open_array instead of using the
 Array constructor

---
 dask/array/core.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dask/array/core.py b/dask/array/core.py
index 76ab569f4ee..4c7ebd83bb4 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3635,9 +3635,9 @@ def from_zarr(
             store = zarr.storage.FSStore(url, **storage_options)
         else:
             store = url
-        z = zarr.Array(store, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(store, read_only=True, path=component, **kwargs)
     else:
-        z = zarr.Array(url, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(url, read_only=True, path=component, **kwargs)
     chunks = chunks if chunks is not None else z.chunks
     if name is None:
         name = "from-zarr-" + tokenize(z, component, storage_options, chunks, **kwargs)

From e37704b318388f37808b63d9147b5260e4bdcbae Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Sat, 14 Sep 2024 14:09:12 -0700
Subject: [PATCH 02/12] further changes for zarr3 compat

---
 dask/array/core.py                  | 22 +++++++++++++++++-----
 dask/array/tests/test_array_core.py | 19 ++++++++++++-------
 dask/tests/test_distributed.py      | 16 ++++++++++++----
 3 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/dask/array/core.py b/dask/array/core.py
index 4c7ebd83bb4..d31b1999bb8 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -22,6 +22,7 @@
 
 import numpy as np
 from numpy.typing import ArrayLike
+from packaging.version import Version
 from tlz import accumulate, concat, first, frequencies, groupby, partition
 from tlz.curried import pluck
 
@@ -3632,12 +3633,15 @@ def from_zarr(
         if isinstance(url, os.PathLike):
             url = os.fspath(url)
         if storage_options:
-            store = zarr.storage.FSStore(url, **storage_options)
+            if Version(zarr.__version__) < Version("3.0.0.a0"):
+                store = zarr.storage.FSStore(url, **storage_options)
+            else:
+                zarr.store.RemoteStore(url, **storage_options)
         else:
             store = url
-        z = zarr.open_array(store, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(store=store, read_only=True, path=component, **kwargs)
     else:
-        z = zarr.open_array(url, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(store=url, read_only=True, path=component, **kwargs)
     chunks = chunks if chunks is not None else z.chunks
     if name is None:
         name = "from-zarr-" + tokenize(z, component, storage_options, chunks, **kwargs)
@@ -3706,9 +3710,14 @@ def to_zarr(
             "currently supported by Zarr.%s" % unknown_chunk_message
         )
 
+    if Version(zarr.__version__) < Version("3.0.0.a0"):
+        zarr_mem_store_types = (dict, zarr.storage.MemoryStore, zarr.storage.KVStore)
+    else:
+        zarr_mem_store_types = (zarr.store.MemoryStore,)
+
     if isinstance(url, zarr.Array):
         z = url
-        if isinstance(z.store, (dict, zarr.storage.MemoryStore, zarr.storage.KVStore)):
+        if isinstance(z.store, zarr_mem_store_types):
             try:
                 from distributed import default_client
 
@@ -3751,7 +3760,10 @@ def to_zarr(
     storage_options = storage_options or {}
 
     if storage_options:
-        store = zarr.storage.FSStore(url, **storage_options)
+        if Version(zarr.__version__) < Version("3.0.0.a0"):
+            store = zarr.storage.FSStore(url, **storage_options)
+        else:
+            store = zarr.store.RemoteStore(url, **storage_options)
     else:
         store = url
 
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 14467151861..c890d8b54c9 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -19,6 +19,7 @@
 from operator import add, sub
 from threading import Lock
 
+from packaging.version import Version
 from tlz import concat, merge
 from tlz.curried import identity
 
@@ -4598,15 +4599,17 @@ def test_read_zarr_chunks():
         assert arr.chunks == ((5, 4),)
 
 
-def test_zarr_pass_mapper():
-    pytest.importorskip("zarr")
-    import zarr.storage
+def test_zarr_pass_store():
+    zarr = pytest.importorskip("zarr")
 
     with tmpdir() as d:
-        mapper = zarr.storage.DirectoryStore(d)
+        if Version(zarr.__version__) < Version("3.0.0.a0"):
+            store = zarr.storage.DirectoryStore(d)
+        else:
+            store = zarr.store.LocalStore(d, mode="w")
         a = da.zeros((3, 3), chunks=(1, 1))
-        a.to_zarr(mapper)
-        a2 = da.from_zarr(mapper)
+        a.to_zarr(store)
+        a2 = da.from_zarr(store)
         assert_eq(a, a2)
         assert a2.chunks == a.chunks
 
@@ -4623,7 +4626,9 @@ def test_zarr_group():
         # second time is fine, group exists
         a.to_zarr(d, component="test2", overwrite=False)
         a.to_zarr(d, component="nested/test", overwrite=False)
-        group = zarr.open_group(d, mode="r")
+
+        # currently failing here. zarr may not be creating the root group node
+        group = zarr.open_group(store=d, mode="r")
         assert list(group) == ["nested", "test", "test2"]
         assert "test" in group["nested"]
 
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index 1a9fa1eb9ed..d6754a1facc 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -11,6 +11,8 @@
 from functools import partial
 from operator import add
 
+from packaging.version import Version
+
 from distributed import Client, SchedulerPlugin, WorkerPlugin, futures_of, wait
 from distributed.utils_test import cleanup  # noqa F401
 from distributed.utils_test import client as c  # noqa F401
@@ -361,8 +363,11 @@ def test_zarr_distributed_with_explicit_directory_store(c):
     with tmpdir() as d:
         chunks = (1, 1)
         a = da.zeros((3, 3), chunks=chunks)
-        s = zarr.storage.DirectoryStore(d)
-        z = zarr.creation.open_array(
+        if Version(zarr.__version__) < Version("3.0.0.a0"):
+            s = zarr.storage.DirectoryStore(d)
+        else:
+            s = zarr.store.LocalStore(d, mode="a")
+        z = zarr.open_array(
             shape=a.shape,
             chunks=chunks,
             dtype=a.dtype,
@@ -382,8 +387,11 @@ def test_zarr_distributed_with_explicit_memory_store(c):
 
     chunks = (1, 1)
     a = da.zeros((3, 3), chunks=chunks)
-    s = zarr.storage.MemoryStore()
-    z = zarr.creation.open_array(
+    if Version(zarr.__version__) < Version("3.0.0.a0"):
+        s = zarr.storage.MemoryStore()
+    else:
+        s = zarr.store.MemoryStore()
+    z = zarr.open_array(
         shape=a.shape,
         chunks=chunks,
         dtype=a.dtype,

From 4380993110870a33ff7cef45f78d7512963150db Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Sat, 14 Sep 2024 15:46:48 -0700
Subject: [PATCH 03/12] switch to working branch for upstream test

---
 continuous_integration/scripts/install.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/continuous_integration/scripts/install.sh b/continuous_integration/scripts/install.sh
index 87d50d79bf5..fcc887f0684 100644
--- a/continuous_integration/scripts/install.sh
+++ b/continuous_integration/scripts/install.sh
@@ -27,7 +27,7 @@ if [[ ${UPSTREAM_DEV} ]]; then
         git+https://github.com/dask/distributed \
         git+https://github.com/dask/dask-expr \
         git+https://github.com/dask/fastparquet \
-        git+https://github.com/zarr-developers/zarr-python.git@main
+        git+https://github.com/zarr-developers/zarr-python.git@fix/dask-compat
         # Zarr's default branch (`v3`) is still under development.
         # Explicitly specify `main` until their default branch is ready.
         # https://github.com/zarr-developers/zarr-python/issues/1922

From 916561cf20e9dc14f2ca5127639d7698a667361f Mon Sep 17 00:00:00 2001
From: Joe Hamman <jhamman1@gmail.com>
Date: Mon, 16 Sep 2024 09:18:08 -0700
Subject: [PATCH 04/12] Update dask/array/core.py

Co-authored-by: James Bourbeau <jrbourbeau@users.noreply.github.com>
---
 dask/array/core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dask/array/core.py b/dask/array/core.py
index d31b1999bb8..3c70c8480e4 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3636,7 +3636,7 @@ def from_zarr(
             if Version(zarr.__version__) < Version("3.0.0.a0"):
                 store = zarr.storage.FSStore(url, **storage_options)
             else:
-                zarr.store.RemoteStore(url, **storage_options)
+                store = zarr.store.RemoteStore(url, **storage_options)
         else:
             store = url
         z = zarr.open_array(store=store, read_only=True, path=component, **kwargs)

From 104bc6f3a703664935c74d6e2e6cfaabe5449ec7 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Mon, 16 Sep 2024 17:27:41 -0700
Subject: [PATCH 05/12] update imports

---
 dask/array/core.py                  | 6 ++++--
 dask/array/tests/test_array_core.py | 2 +-
 dask/tests/test_distributed.py      | 8 ++++----
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/dask/array/core.py b/dask/array/core.py
index 3c70c8480e4..75db12cdeef 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3713,7 +3713,7 @@ def to_zarr(
     if Version(zarr.__version__) < Version("3.0.0.a0"):
         zarr_mem_store_types = (dict, zarr.storage.MemoryStore, zarr.storage.KVStore)
     else:
-        zarr_mem_store_types = (zarr.store.MemoryStore,)
+        zarr_mem_store_types = (zarr.storage.MemoryStore,)
 
     if isinstance(url, zarr.Array):
         z = url
@@ -3763,7 +3763,9 @@ def to_zarr(
         if Version(zarr.__version__) < Version("3.0.0.a0"):
             store = zarr.storage.FSStore(url, **storage_options)
         else:
-            store = zarr.store.RemoteStore(url, **storage_options)
+            store = zarr.storage.RemoteStore(
+                url, mode=kwargs.pop("mode", "a"), **storage_options
+            )
     else:
         store = url
 
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index c890d8b54c9..25cc2072992 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -4606,7 +4606,7 @@ def test_zarr_pass_store():
         if Version(zarr.__version__) < Version("3.0.0.a0"):
             store = zarr.storage.DirectoryStore(d)
         else:
-            store = zarr.store.LocalStore(d, mode="w")
+            store = zarr.storage.LocalStore(d, mode="w")
         a = da.zeros((3, 3), chunks=(1, 1))
         a.to_zarr(store)
         a2 = da.from_zarr(store)
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index d6754a1facc..fcc86fa23d0 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -366,13 +366,13 @@ def test_zarr_distributed_with_explicit_directory_store(c):
         if Version(zarr.__version__) < Version("3.0.0.a0"):
             s = zarr.storage.DirectoryStore(d)
         else:
-            s = zarr.store.LocalStore(d, mode="a")
+            s = zarr.storage.LocalStore(d, mode="a")
         z = zarr.open_array(
             shape=a.shape,
             chunks=chunks,
             dtype=a.dtype,
             store=s,
-            mode="a",
+            # mode="a",
         )
         a.to_zarr(z)
         a2 = da.from_zarr(d)
@@ -390,13 +390,13 @@ def test_zarr_distributed_with_explicit_memory_store(c):
     if Version(zarr.__version__) < Version("3.0.0.a0"):
         s = zarr.storage.MemoryStore()
     else:
-        s = zarr.store.MemoryStore()
+        s = zarr.storage.MemoryStore(mode="a")
     z = zarr.open_array(
         shape=a.shape,
         chunks=chunks,
         dtype=a.dtype,
         store=s,
-        mode="a",
+        # mode="a",
     )
     with pytest.raises(RuntimeError, match="distributed scheduler"):
         a.to_zarr(z)

From e3aab8320513d37da8f5aa618c1ccf52ebdb3be0 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Sun, 29 Sep 2024 14:51:55 -0700
Subject: [PATCH 06/12] fixup test

---
 dask/array/tests/test_array_core.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 25cc2072992..8d56ead0179 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -4627,9 +4627,8 @@ def test_zarr_group():
         a.to_zarr(d, component="test2", overwrite=False)
         a.to_zarr(d, component="nested/test", overwrite=False)
 
-        # currently failing here. zarr may not be creating the root group node
         group = zarr.open_group(store=d, mode="r")
-        assert list(group) == ["nested", "test", "test2"]
+        assert set(group) == {"nested", "test", "test2"}
         assert "test" in group["nested"]
 
         a2 = da.from_zarr(d, component="test")

From 228deb97d41845ae9d26bdc6a8bab0ad26d18ce9 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Sun, 29 Sep 2024 15:06:01 -0700
Subject: [PATCH 07/12] fixup distributed test

---
 dask/tests/test_distributed.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index fcc86fa23d0..1ab5fb67636 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -372,7 +372,7 @@ def test_zarr_distributed_with_explicit_directory_store(c):
             chunks=chunks,
             dtype=a.dtype,
             store=s,
-            # mode="a",
+            mode="a",
         )
         a.to_zarr(z)
         a2 = da.from_zarr(d)
@@ -396,7 +396,7 @@ def test_zarr_distributed_with_explicit_memory_store(c):
         chunks=chunks,
         dtype=a.dtype,
         store=s,
-        # mode="a",
+        mode="a",
     )
     with pytest.raises(RuntimeError, match="distributed scheduler"):
         a.to_zarr(z)

From 1134808e8d4a11c4989be54421d96e63b6a01437 Mon Sep 17 00:00:00 2001
From: Joe Hamman <jhamman1@gmail.com>
Date: Tue, 1 Oct 2024 09:27:43 -0700
Subject: [PATCH 08/12] Update continuous_integration/scripts/install.sh

---
 continuous_integration/scripts/install.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/continuous_integration/scripts/install.sh b/continuous_integration/scripts/install.sh
index fcc887f0684..36da38bb589 100644
--- a/continuous_integration/scripts/install.sh
+++ b/continuous_integration/scripts/install.sh
@@ -27,7 +27,7 @@ if [[ ${UPSTREAM_DEV} ]]; then
         git+https://github.com/dask/distributed \
         git+https://github.com/dask/dask-expr \
         git+https://github.com/dask/fastparquet \
-        git+https://github.com/zarr-developers/zarr-python.git@fix/dask-compat
+        git+https://github.com/zarr-developers/zarr-python.git@v3
         # Zarr's default branch (`v3`) is still under development.
         # Explicitly specify `main` until their default branch is ready.
         # https://github.com/zarr-developers/zarr-python/issues/1922

From a04473094ff817a043e2752a523930e0794e5c90 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Sun, 6 Oct 2024 22:34:56 -0700
Subject: [PATCH 09/12] fixup order / xarray tests

---
 dask/tests/test_order.py  | 7 ++++++-
 dask/tests/test_sizeof.py | 8 ++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 4dda35041ee..69474981848 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -4,6 +4,7 @@
 from collections import defaultdict
 
 import pytest
+from packaging.version import Version
 
 import dask
 from dask import delayed
@@ -891,7 +892,11 @@ def test_array_store_final_order(tmpdir):
     arrays = [da.ones((110, 4), chunks=(100, 2)) for i in range(4)]
     x = da.concatenate(arrays, axis=0).rechunk((100, 2))
 
-    store = zarr.DirectoryStore(tmpdir)
+    if Version(zarr.__version__) < Version("3.0.0.a0"):
+        store = zarr.storage.DirectoryStore(tmpdir)
+    else:
+        store = zarr.storage.LocalStore(str(tmpdir), mode="w")
+
     root = zarr.group(store, overwrite=True)
     dest = root.empty_like(name="dest", data=x, chunks=x.chunksize, overwrite=True)
     d = x.store(dest, lock=False, compute=False)
diff --git a/dask/tests/test_sizeof.py b/dask/tests/test_sizeof.py
index dd99d27d5db..5093f67a93f 100644
--- a/dask/tests/test_sizeof.py
+++ b/dask/tests/test_sizeof.py
@@ -5,6 +5,7 @@
 from array import array
 
 import pytest
+from packaging.version import Version
 
 from dask.multiprocessing import get_context
 from dask.sizeof import sizeof
@@ -276,12 +277,19 @@ def test_xarray():
 
 def test_xarray_not_in_memory():
     xr = pytest.importorskip("xarray")
+    zarr = pytest.importorskip("zarr")
     np = pytest.importorskip("numpy")
     pytest.importorskip("zarr")
 
     ind = np.arange(-66, 67, 1).astype(float)
     arr = np.random.random((len(ind),))
 
+    # TODO: remove this conditional after consolidated metadata lands in v3
+    if Version(zarr.__version__) > Version("3.0.0.a0") and Version(
+        zarr.__version__
+    ) < Version("3.0.0"):
+        pytest.xfail("consolidated metadata and xarray support is not complete")
+
     with tmpdir() as path:
         xr.DataArray(
             arr,

From 9721f0253f732f4536a75a85f1f5d40e6c8440b8 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Thu, 10 Oct 2024 11:20:03 -0700
Subject: [PATCH 10/12] cleanup resources in fixture

---
 dask/tests/test_distributed.py | 36 ++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index 1ab5fb67636..f28e21f13f1 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -342,10 +342,33 @@ def test_futures_in_graph(c):
     assert xxyy3.compute(scheduler="dask.distributed") == ((1 + 1) + (2 + 2)) + 10
 
 
-def test_zarr_distributed_roundtrip(c):
+@pytest.fixture(scope="function")
+def zarr(c):
+    zarr_lib = pytest.importorskip("zarr")
+    # Zarr-Python 3 lazily allocates a dedicated thread/IO loop
+    # for to execute async tasks. To avoid having this thread
+    # be picked up as a "leaked thread", we manually trigger it's
+    # creation before using zarr
+    try:
+        _ = zarr_lib.core.sync._get_loop()
+        _ = zarr_lib.core.sync._get_executor()
+        yield zarr_lib
+    except AttributeError:
+        yield zarr_lib
+    finally:
+        # Zarr-Python 3 lazily allocates a IO thread, a thread pool executor, and
+        # an IO loop. Here we clean up these resources to avoid leaking threads
+        # In normal operations, this is done as by an atexit handler when Zarr
+        # is shutting down.
+        try:
+            zarr_lib.core.sync.cleanup_resources()
+        except AttributeError:
+            pass
+
+
+def test_zarr_distributed_roundtrip(c, zarr):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
-    pytest.importorskip("zarr")
 
     with tmpdir() as d:
         a = da.zeros((3, 3), chunks=(1, 1))
@@ -355,10 +378,9 @@ def test_zarr_distributed_roundtrip(c):
         assert a2.chunks == a.chunks
 
 
-def test_zarr_distributed_with_explicit_directory_store(c):
+def test_zarr_distributed_with_explicit_directory_store(c, zarr):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
-    zarr = pytest.importorskip("zarr")
 
     with tmpdir() as d:
         chunks = (1, 1)
@@ -380,10 +402,9 @@ def test_zarr_distributed_with_explicit_directory_store(c):
         assert a2.chunks == a.chunks
 
 
-def test_zarr_distributed_with_explicit_memory_store(c):
+def test_zarr_distributed_with_explicit_memory_store(c, zarr):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
-    zarr = pytest.importorskip("zarr")
 
     chunks = (1, 1)
     a = da.zeros((3, 3), chunks=chunks)
@@ -402,10 +423,9 @@ def test_zarr_distributed_with_explicit_memory_store(c):
         a.to_zarr(z)
 
 
-def test_zarr_in_memory_distributed_err(c):
+def test_zarr_in_memory_distributed_err(c, zarr):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
-    zarr = pytest.importorskip("zarr")
 
     chunks = (1, 1)
     a = da.ones((3, 3), chunks=chunks)

From eadc5f57311e1ec1bfe976a2f09e2dd9de843ba4 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Thu, 10 Oct 2024 11:27:35 -0700
Subject: [PATCH 11/12] point at working branch

---
 continuous_integration/scripts/install.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/continuous_integration/scripts/install.sh b/continuous_integration/scripts/install.sh
index 36da38bb589..8f7fdea8df5 100644
--- a/continuous_integration/scripts/install.sh
+++ b/continuous_integration/scripts/install.sh
@@ -27,7 +27,7 @@ if [[ ${UPSTREAM_DEV} ]]; then
         git+https://github.com/dask/distributed \
         git+https://github.com/dask/dask-expr \
         git+https://github.com/dask/fastparquet \
-        git+https://github.com/zarr-developers/zarr-python.git@v3
+        git+https://github.com/jhamman/zarr-python.git@fix/thread-cleanup
         # Zarr's default branch (`v3`) is still under development.
         # Explicitly specify `main` until their default branch is ready.
         # https://github.com/zarr-developers/zarr-python/issues/1922

From 0094664884cc5e2bcb3a16cad635468add43dcb2 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Thu, 10 Oct 2024 13:08:34 -0700
Subject: [PATCH 12/12] tidy

---
 continuous_integration/scripts/install.sh |  5 +---
 dask/array/core.py                        | 30 +++++++++++++++--------
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/continuous_integration/scripts/install.sh b/continuous_integration/scripts/install.sh
index 8f7fdea8df5..fd72ee15251 100644
--- a/continuous_integration/scripts/install.sh
+++ b/continuous_integration/scripts/install.sh
@@ -27,10 +27,7 @@ if [[ ${UPSTREAM_DEV} ]]; then
         git+https://github.com/dask/distributed \
         git+https://github.com/dask/dask-expr \
         git+https://github.com/dask/fastparquet \
-        git+https://github.com/jhamman/zarr-python.git@fix/thread-cleanup
-        # Zarr's default branch (`v3`) is still under development.
-        # Explicitly specify `main` until their default branch is ready.
-        # https://github.com/zarr-developers/zarr-python/issues/1922
+        git+https://github.com/zarr-developers/zarr-python
     mamba uninstall --force numpy pandas scipy numexpr numba sparse scikit-image h5py
     python -m pip install --no-deps --pre --retries 10 \
         -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple \
diff --git a/dask/array/core.py b/dask/array/core.py
index 75db12cdeef..7d4f023fa61 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -13,7 +13,7 @@
 from bisect import bisect
 from collections import defaultdict
 from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence
-from functools import partial, reduce, wraps
+from functools import lru_cache, partial, reduce, wraps
 from itertools import product, zip_longest
 from numbers import Integral, Number
 from operator import add, mul
@@ -3583,6 +3583,16 @@ def from_array(
     return Array(dsk, name, chunks, meta=meta, dtype=getattr(x, "dtype", None))
 
 
+@lru_cache
+def _zarr_v3() -> bool:
+    try:
+        import zarr
+    except ImportError:
+        return False
+    else:
+        return Version(zarr.__version__).major >= 3
+
+
 def from_zarr(
     url,
     component=None,
@@ -3633,10 +3643,10 @@ def from_zarr(
         if isinstance(url, os.PathLike):
             url = os.fspath(url)
         if storage_options:
-            if Version(zarr.__version__) < Version("3.0.0.a0"):
-                store = zarr.storage.FSStore(url, **storage_options)
-            else:
+            if _zarr_v3():
                 store = zarr.store.RemoteStore(url, **storage_options)
+            else:
+                store = zarr.storage.FSStore(url, **storage_options)
         else:
             store = url
         z = zarr.open_array(store=store, read_only=True, path=component, **kwargs)
@@ -3710,10 +3720,10 @@ def to_zarr(
             "currently supported by Zarr.%s" % unknown_chunk_message
         )
 
-    if Version(zarr.__version__) < Version("3.0.0.a0"):
-        zarr_mem_store_types = (dict, zarr.storage.MemoryStore, zarr.storage.KVStore)
-    else:
+    if _zarr_v3():
         zarr_mem_store_types = (zarr.storage.MemoryStore,)
+    else:
+        zarr_mem_store_types = (dict, zarr.storage.MemoryStore, zarr.storage.KVStore)
 
     if isinstance(url, zarr.Array):
         z = url
@@ -3760,12 +3770,12 @@ def to_zarr(
     storage_options = storage_options or {}
 
     if storage_options:
-        if Version(zarr.__version__) < Version("3.0.0.a0"):
-            store = zarr.storage.FSStore(url, **storage_options)
-        else:
+        if _zarr_v3():
             store = zarr.storage.RemoteStore(
                 url, mode=kwargs.pop("mode", "a"), **storage_options
             )
+        else:
+            store = zarr.storage.FSStore(url, **storage_options)
     else:
         store = url