Merge pull request #895 from m-albert/build_pyramid_zarr_creation_options

Added optional parameter open_array_kwargs to build_pyramid

Author: jluethi

fractal_tasks_core/pyramids.py

@@ -15,12 +15,14 @@
 import logging
 import pathlib
 from typing import Callable
+from typing import Mapping
 from typing import Optional
 from typing import Sequence
 from typing import Union
 
 import dask.array as da
 import numpy as np
+import zarr.errors
 
 logger = logging.getLogger(__name__)
 
@@ -33,6 +35,7 @@ def build_pyramid(
     coarsening_xy: int = 2,
     chunksize: Optional[Sequence[int]] = None,
     aggregation_function: Optional[Callable] = None,
+    open_array_kwargs: Optional[Mapping] = None,
 ) -> None:
 
     """
@@ -48,6 +51,7 @@ def build_pyramid(
         coarsening_xy: Linear coarsening factor between subsequent levels.
         chunksize: Shape of a single chunk.
         aggregation_function: Function to be used when downsampling.
+        open_array_kwargs: Additional arguments for zarr.open.
     """
 
     # Clean up zarrurl
@@ -100,10 +104,33 @@ def build_pyramid(
             f"{str(newlevel_rechunked)}"
         )
 
+        if open_array_kwargs is None:
+            open_array_kwargs = {}
+
+        # If overwrite is false, check that the array doesn't exist yet
+        if not overwrite:
+            try:
+                zarr.open(f"{zarrurl}/{ind_level}", mode="r")
+                raise ValueError(
+                    f"While building the pyramids, pyramid level {ind_level} "
+                    "already existed, but `build_pyramid` was called with "
+                    f"{overwrite=}."
+                )
+            except zarr.errors.PathNotFoundError:
+                pass
+
+        zarrarr = zarr.open(
+            f"{zarrurl}/{ind_level}",
+            shape=newlevel_rechunked.shape,
+            chunks=newlevel_rechunked.chunksize,
+            dtype=newlevel_rechunked.dtype,
+            mode="w",
+            **open_array_kwargs,
+        )
+
         # Write zarr and store output (useful to construct next level)
         previous_level = newlevel_rechunked.to_zarr(
-            zarrurl,
-            component=f"{ind_level}",
+            zarrarr,
             overwrite=overwrite,
             compute=True,
             return_stored=True,

tests/test_unit_pyramid_creation.py

@@ -64,3 +64,36 @@ def test_build_pyramid(tmp_path):
     assert level_3.chunksize == (9, 9)
     assert level_4.shape == (3, 3)
     assert level_5.shape == (1, 1)
+
+    # Succeed
+    zarrurl = str(tmp_path / "F.zarr")
+    da.zeros(shape=(8, 8)).to_zarr(f"{zarrurl}/0")
+    build_pyramid(
+        zarrurl=zarrurl,
+        coarsening_xy=2,
+        num_levels=3,
+        open_array_kwargs={"write_empty_chunks": False, "fill_value": 0},
+    )
+    level_1 = da.from_zarr(f"{zarrurl}/1")
+    level_2 = da.from_zarr(f"{zarrurl}/2")
+    assert level_1.shape == (4, 4)
+    assert level_2.shape == (2, 2)
+    # check that the empty chunks are not written to disk
+    assert not (tmp_path / "F.zarr/1/0.0").exists()
+    assert not (tmp_path / "F.zarr/2/0.0").exists()
+
+
+def test_build_pyramid_overwrite(tmp_path):
+    # Succeed
+    zarrurl = str(tmp_path / "D.zarr")
+    da.ones(shape=(8, 8)).to_zarr(f"{zarrurl}/0")
+    build_pyramid(zarrurl=zarrurl, coarsening_xy=2, num_levels=3)
+    # Should fail because overwrite is not set
+    with pytest.raises(ValueError):
+        build_pyramid(
+            zarrurl=zarrurl, coarsening_xy=2, num_levels=3, overwrite=False
+        )
+    # Should work
+    build_pyramid(
+        zarrurl=zarrurl, coarsening_xy=2, num_levels=3, overwrite=True
+    )