funlib.persistence 0.7.0__tar.gz → 0.8.0__tar.gz

{funlib_persistence-0.7.0 → funlib_persistence-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: funlib.persistence
-Version: 0.7.0
+Version: 0.8.0
 Summary: Interfaces for data (arrays and graphs) and storage formats (databases and file formats)
 Author-email: William Patton <pattonw@hhmi.org>, Jan Funke <funkej@hhmi.org>
 License: MIT

{funlib_persistence-0.7.0 → funlib_persistence-0.8.0}/src/funlib/persistence/__init__.py

@@ -1,4 +1,4 @@
 from .arrays import Array, open_ds, prepare_ds, open_ome_ds, prepare_ome_ds  # noqa
 
-__version__ = "0.7.0"
+__version__ = "0.8.0"
 __version_info__ = tuple(int(i) for i in __version__.split("."))

{funlib_persistence-0.7.0 → funlib_persistence-0.8.0}/src/funlib/persistence/arrays/array.py

@@ -74,6 +74,18 @@ class Array(Freezable):
             offset, voxel_size, axis_names, units, and types. Metadata
             can either be passed in or read from array attributes.
 
+        scheduler (``Optional[str]``):
+
+            The dask scheduler used to compute reads (`__getitem__` and
+            `to_ndarray`); writes do not go through dask. `None` (the default)
+            leaves the choice to dask, which uses its threaded scheduler for
+            arrays. Pass `"single-threaded"` when the caller already
+            parallelises array access (one process or thread per block, as
+            daisy does), where dask's own thread pool only adds contention
+            and can deadlock. See
+            https://docs.dask.org/en/stable/scheduler-overview.html for the
+            available schedulers.
+
     """
 
     data: da.Array
@@ -89,7 +101,9 @@ class Array(Freezable):
         chunks: Optional[Union[int, Sequence[int], str]] = "auto",
         lazy_op: Optional[LazyOp] = None,
         strict_metadata: bool = False,
+        scheduler: Optional[str] = None,
     ):
+        self.scheduler = scheduler
         if not isinstance(data, da.Array):
             self.data = da.from_array(data, chunks=chunks)
         else:
@@ -110,11 +124,13 @@ class Array(Freezable):
         # only used if not reading from zarr and there is no built in `.attrs`
         self._attrs: dict[str, Any] = {}
 
+        # whether the view's chunks still sit on the stored chunk grid (see
+        # `_keeps_chunk_grid`); `chunk_shape` is reported only while they do
+        self._on_chunk_grid = True
+        self.lazy_ops: list[LazyOp] = []
         if lazy_op is not None:
             self.apply_lazy_ops(lazy_op)
-
-        lazy_ops = [] if lazy_op is None else [lazy_op]
-        self.lazy_ops = lazy_ops
+            self.lazy_ops = [lazy_op]
 
         self.freeze()
 
@@ -133,9 +149,67 @@ class Array(Freezable):
             return self._attrs
 
     @property
-    def chunk_shape(self) -> Coordinate:
+    def chunk_shape(self) -> Optional[Coordinate]:
+        """The write-safe chunk shape: the shard shape if sharded, else the chunk
+        shape (an ``open_ds(chunks=...)`` override is reported as given). ``None``
+        for a view that is not writeable or that lazy operations moved off the
+        stored chunk grid."""
+        if not self.is_writeable or not self._on_chunk_grid:
+            return None
+        shards = self.shard_shape
+        return shards if shards is not None else self._dask_chunk_shape()
+
+    @property
+    def inner_chunk_shape(self) -> Optional[Coordinate]:
+        """Chunk shape for efficient reads. Can be smaller than the write-safe
+        :attr:`chunk_shape` in the case of sharding. ``None`` for a view that
+        lazy operations moved off the stored chunk grid."""
+        return self._dask_chunk_shape() if self._on_chunk_grid else None
+
+    @property
+    def shard_shape(self) -> Optional[Coordinate]:
+        """The shard shape in the case of sharded datasets. ``None`` for non
+        sharded or unwritable datasets, or off the stored chunk grid. See
+        https://zarr.readthedocs.io/en/stable/user-guide/performance.html#sharding
+        """
+        shards = getattr(self._source_data, "shards", None)
+        if shards is None or not self.is_writeable or not self._on_chunk_grid:
+            return None
+        return Coordinate(
+            s for s, uncollapsed in zip(shards, self._uncollapsed_dims) if uncollapsed
+        )
+
+    def _dask_chunk_shape(self) -> Coordinate:
         return Coordinate(self.data.chunksize)  # ty: ignore[unresolved-attribute]
 
+    def _keeps_chunk_grid(self, lazy_op) -> bool:
+        """Whether a lazy op leaves this view's chunks on the stored chunk grid:
+        dropping or truncating axes does, a start, a step, a selection list, a
+        mask or an arbitrary callable does not. A ``Roi`` does if it begins on
+        a chunk boundary."""
+        if isinstance(lazy_op, Roi):
+            if (
+                not all(self.uncollapsed_dims(physical=True))
+                or lazy_op.dims != self.spatial_dims
+            ):
+                return True  # apply_lazy_ops refuses the op, with its own message
+            unit = self.chunk_shape
+            if unit is None:
+                return False
+            physical = [c for c, t in zip(unit, self.types) if t in ["space", "time"]]
+            begin = (lazy_op.offset - self.offset) / self.voxel_size
+            return all(b % c == 0 for b, c in zip(begin, physical))
+        if not isinstance(lazy_op, tuple):
+            lazy_op = (lazy_op,)
+        for a in lazy_op:
+            if isinstance(a, int):
+                continue
+            if isinstance(a, slice):
+                if a.start in (None, 0) and a.step in (None, 1):
+                    continue
+            return False
+        return True
+
     def uncollapsed_dims(self, physical: bool = False) -> list[bool]:
         """
         We support lazy slicing of arrays, such as `x = Array(np.ones(5,5,5))` and
@@ -253,17 +327,18 @@ class Array(Freezable):
         )
 
     def apply_lazy_ops(self, lazy_op):
+        # decided before the op moves the offset, recorded only once it applied
+        on_chunk_grid = self._on_chunk_grid and self._keeps_chunk_grid(lazy_op)
         if self._is_slice(lazy_op):
             if not isinstance(lazy_op, tuple):
                 lazy_op = (lazy_op,)
-            for ii, a in enumerate(lazy_op):
+            # Pair each entry with the dimension it indexes BEFORE collapsing any: the
+            # previous in-place walk re-counted uncollapsed dims after each collapse,
+            # so a tuple with two or more ints collapsed the wrong dimensions.
+            dims = [i for i, uc in enumerate(self._uncollapsed_dims) if uc]
+            for a, i in zip(lazy_op, dims):
                 if isinstance(a, int):
-                    for i, uc in enumerate(self._uncollapsed_dims):
-                        if uc:
-                            if ii == 0:
-                                self._uncollapsed_dims[i] = False
-                                break
-                            ii -= 1
+                    self._uncollapsed_dims[i] = False
             self.data = self.data[lazy_op]
         elif isinstance(lazy_op, Roi):
             assert all(self.uncollapsed_dims(physical=True)), (
@@ -282,6 +357,7 @@ class Array(Freezable):
                 f"LazyOp {lazy_op} is not a supported lazy_op. "
                 f"Supported lazy_ops are: {LazyOp}"
             )
+        self._on_chunk_grid = on_chunk_grid
 
     def lazy_op(self, lazy_op: LazyOp):
         """Apply an lazy operation to this array.
@@ -329,7 +405,9 @@ class Array(Freezable):
                     % (roi, self.roi)
                 )
 
-            return self.data[self.__slices(roi, use_lazy_slices=False)].compute()
+            return self.data[self.__slices(roi, use_lazy_slices=False)].compute(
+                scheduler=self.scheduler
+            )
 
         elif isinstance(key, Coordinate):
             coordinate = key
@@ -338,10 +416,10 @@ class Array(Freezable):
                 raise IndexError("Requested coordinate is not contained in this array.")
 
             index = self.__index(coordinate)
-            return self.data[index].compute()
+            return self.data[index].compute(scheduler=self.scheduler)
 
         else:
-            return self.data[key].compute()
+            return self.data[key].compute(scheduler=self.scheduler)
 
     def __setitem__(self, key: Roi | slice | tuple, value: np.ndarray | float | int):
         """Set the data of this array.
@@ -421,7 +499,7 @@ class Array(Freezable):
         if fill_value != 0:
             data[:] = fill_value
 
-        array = Array(data, roi.offset, self.voxel_size)
+        array = Array(data, roi.offset, self.voxel_size, scheduler=self.scheduler)
 
         shared_roi = self.roi.intersect(roi)
 
@@ -438,18 +516,20 @@ class Array(Freezable):
         voxel_roi = (roi - self.offset) / self.voxel_size
 
         if check_chunk_align:
+            chunk_shape = self.chunk_shape
+            assert chunk_shape is not None, "this view has no chunk grid to align to"
             for d in range(roi.dims):
                 end_of_array = roi.get_end()[d] == self.roi.get_end()[d]
 
-                begin_align_with_chunks = voxel_roi.begin[d] % self.chunk_shape[d] == 0
-                shape_align_with_chunks = voxel_roi.shape[d] % self.chunk_shape[d] == 0
+                begin_align_with_chunks = voxel_roi.begin[d] % chunk_shape[d] == 0
+                shape_align_with_chunks = voxel_roi.shape[d] % chunk_shape[d] == 0
 
                 assert begin_align_with_chunks and (
                     shape_align_with_chunks or end_of_array
                 ), (
                     "ROI %s (in voxels: %s) does not align with chunks of "
                     "size %s (mismatch in dimension %d)"
-                    % (roi, voxel_roi, self.chunk_shape, d)
+                    % (roi, voxel_roi, chunk_shape, d)
                 )
 
         roi_slices = tuple(